Emit transactional update functions

src/emit.ml

@@ -133,7 +133,7 @@ let line buffer text = Buffer.add_string buffer (text ^ "\n")
 
 let indexed prefix index = Printf.sprintf "%s%d" prefix index
 
-let binder_names of_expr =
+let binders_of of_expr =
   let binders = ref [] in
   let rec walk expr =
     match expr.Anf.a with
@@ -155,6 +155,9 @@ let binder_names of_expr =
         ()
   in
   walk of_expr;
+  !binders
+
+let names_of_idents idents =
   let used = Hashtbl.create 32 in
   List.map
     (fun ident ->
@@ -164,7 +167,13 @@ let binder_names of_expr =
       in
       Hashtbl.replace used candidate ();
       (Ident.stamp ident, candidate))
-    !binders
+    idents
+
+let binder_names of_expr = names_of_idents (binders_of of_expr)
+
+let names_with_suffix suffix names = List.map (fun (stamp, name) -> (stamp, name ^ suffix)) names
+
+let atom_expr span ty atom = { Anf.a = Anf.AAtom atom; aty = ty; aspan = span }
 
 let var_name names ident =
   match Util.assoc_opt (Ident.stamp ident) names with
@@ -531,28 +540,306 @@ let emit_output_functions plan buffer =
     line buffer
       (Printf.sprintf
          "    \"[ \" ^ String.concat \"; \" (List.rev (List.rev_map (fun (key, item) -> Printf.sprintf \"(%%d, %%s)\" key (show_%s item)) (Delta_runtime.Pure_map.bindings value))) ^ \" ]\""
-         (type_name (output_value_type plan))))
+         (type_name (output_value_type plan)));
+    line buffer "";
+    line buffer "  let output_change_to_string (change : output_change) : string =";
+    line buffer
+      (Printf.sprintf
+         "    String.concat \"; \" (List.rev (List.rev_map (fun (key, item) -> match item with Ins value -> Printf.sprintf \"insert %%d %%s\" key (show_%s value) | Rem value -> Printf.sprintf \"remove %%d %%s\" key (show_%s value) | Rep (old_value, new_value) -> Printf.sprintf \"replace %%d %%s %%s\" key (show_%s old_value) (show_%s new_value)) change))"
+         (type_name (output_value_type plan)) (type_name (output_value_type plan))
+         (type_name (output_value_type plan)) (type_name (output_value_type plan))))
   else (
     line buffer "  let apply_output_change (value : output) (change : output_change) : output = value + change";
     line buffer "";
-    line buffer "  let output_to_string (value : output) : string = string_of_int value");
+    line buffer "  let output_to_string (value : output) : string = string_of_int value";
+    line buffer "";
+    line buffer "  let output_change_to_string (change : output_change) : string = Printf.sprintf \"%+d\" change");
   line buffer ""
 
 let emit_signature plan buffer =
   line buffer "  type state";
   line buffer "  type output";
   line buffer "  type output_change";
+  line buffer
+    (Printf.sprintf "  type batch_op = Insert of int * %s | Remove of int | Replace of int * %s"
+       (ocaml_ty (input_type plan)) (ocaml_ty (input_type plan)));
   line buffer (Printf.sprintf "  val init : (int * %s) list -> state" (ocaml_ty (input_type plan)));
+  line buffer
+    "  val apply_batch : state -> batch_op list -> (state * output_change) Delta_runtime.outcome";
   line buffer "  val result : state -> output";
   line buffer "  val apply_output_change : output -> output_change -> output";
   line buffer "  val output_to_string : output -> string";
+  line buffer "  val output_change_to_string : output_change -> string";
   line buffer "  val counters : unit -> Delta_runtime.counters";
   line buffer "  val reset_counters : unit -> unit"
 
+let emit_batch_type plan buffer =
+  line buffer
+    (Printf.sprintf "  type batch_op = Insert of int * %s | Remove of int | Replace of int * %s"
+       (ocaml_ty (input_type plan)) (ocaml_ty (input_type plan)));
+  line buffer ""
+
+let rec emit_delta_functions plan buffer =
+  List.iter
+    (fun node ->
+      match node.Graph.n_kind with
+      | Graph.Map (parameter, body) ->
+          let names = names_of_idents (parameter :: binders_of body) in
+          let old_names = names_with_suffix "_old" names in
+          let new_names = names_with_suffix "_new" names in
+          let parameter_name = sanitize (Ident.display parameter) in
+          let element = node_input_element plan node in
+          line buffer
+            (Printf.sprintf "  let delta_%d (%s_old : %s) (%s_new : %s) : %s =" node.Graph.n_id
+               parameter_name (ocaml_ty element) parameter_name (ocaml_ty element)
+               (change_ty body.Anf.aty));
+          let rec statements expr =
+            match expr.Anf.a with
+            | Anf.ALet (ident, bound, rest) ->
+                let name = var_name names ident in
+                line buffer (Printf.sprintf "    let %s_old = %s in" name (render_expr old_names bound));
+                line buffer (Printf.sprintf "    let %s_new = %s in" name (render_expr new_names bound));
+                statements rest
+            | _ -> ()
+          in
+          line buffer "    Delta_runtime.count_scalar_delta query_counters;";
+          statements body;
+          line buffer (Printf.sprintf "    %s" (render_delta plan old_names new_names body));
+          line buffer ""
+      | _ -> ())
+    plan.Graph.pl_nodes
+
+and render_delta plan old_names new_names expr =
+  match expr.Anf.a with
+  | Anf.AAtom (Anf.AInt _ | Anf.ABool _ | Anf.AString _ | Anf.AUnit) -> "empty_" ^ type_name expr.Anf.aty
+  | Anf.AAtom (Anf.AVar _) -> (
+      match Types.repr expr.Anf.aty with
+      | Types.TInt -> Printf.sprintf "(%s - %s)" (render_expr new_names expr) (render_expr old_names expr)
+      | _ ->
+          Printf.sprintf "(chg_%s %s %s)" (type_name expr.Anf.aty) (render_expr old_names expr)
+            (render_expr new_names expr))
+  | Anf.ABinop (Syntax.Add, left, right) ->
+      Printf.sprintf "(%s + %s)"
+        (render_delta_atom plan old_names new_names expr.Anf.aspan left)
+        (render_delta_atom plan old_names new_names expr.Anf.aspan right)
+  | Anf.ABinop (Syntax.Sub, left, right) ->
+      Printf.sprintf "(%s - %s)"
+        (render_delta_atom plan old_names new_names expr.Anf.aspan left)
+        (render_delta_atom plan old_names new_names expr.Anf.aspan right)
+  | Anf.ABinop (Syntax.Mul, left, right) ->
+      let left_old = render_expr old_names (atom_expr expr.Anf.aspan Types.TInt left) in
+      let right_old = render_expr old_names (atom_expr expr.Anf.aspan Types.TInt right) in
+      let left_delta = render_delta_atom plan old_names new_names expr.Anf.aspan left in
+      let right_delta = render_delta_atom plan old_names new_names expr.Anf.aspan right in
+      Printf.sprintf "((%s * %s) + (%s * %s) + (%s * %s))" left_old right_delta right_old left_delta
+        left_delta right_delta
+  | _ ->
+      Printf.sprintf "(chg_%s %s %s)" (type_name expr.Anf.aty) (render_expr old_names expr)
+        (render_expr new_names expr)
+
+and render_delta_atom plan old_names new_names span atom =
+  render_delta plan old_names new_names (atom_expr span Types.TInt atom)
+
+let emit_update_functions plan buffer =
+  let nodes = plan.Graph.pl_nodes in
+  List.iteri
+    (fun index node ->
+      let keyword = if index = 0 then "let rec" else "and" in
+      let input_element = node_input_element plan node in
+      let consumers = plan.Graph.pl_consumers.(node.Graph.n_id) in
+      let forward old_change new_change =
+        Printf.sprintf "    (match (%s, %s) with\n     | None, None -> state\n     | _ ->\n%s)"
+          old_change new_change
+          (Util.join "\n"
+             (List.map
+                (fun consumer ->
+                  Printf.sprintf "        update_%d state key %s %s" consumer old_change new_change)
+                consumers))
+      in
+      line buffer
+        (Printf.sprintf
+           "  %s update_%d (state : state) (key : int) (up_old : %s option) (up_new : %s option) : state ="
+           keyword node.Graph.n_id (ocaml_ty input_element) (ocaml_ty input_element));
+      match node.Graph.n_kind with
+      | Graph.Source ->
+          line buffer
+            "    let input = (match up_new with None -> Delta_runtime.Pure_map.remove key state.input | Some value -> Delta_runtime.Pure_map.add key value state.input) in";
+          if consumers = [] then line buffer "    { state with input }"
+          else (
+            line buffer "    let state = { state with input } in";
+            line buffer (forward "up_old" "up_new"))
+      | Graph.Filter (_, _) ->
+          let cache = cache_field node in
+          line buffer
+            (Printf.sprintf "    let old_member = Delta_runtime.Pure_map.find_opt key state.%s in" cache);
+          line buffer
+            (Printf.sprintf
+               "    let new_member = (match up_new with None -> None | Some value -> if predicate_%d value then Some value else None) in"
+               node.Graph.n_id);
+          line buffer
+            (Printf.sprintf
+               "    let cached = (match (old_member, new_member) with None, None -> state.%s | None, Some value -> Delta_runtime.Pure_map.add key value state.%s | Some _, None -> Delta_runtime.Pure_map.remove key state.%s | Some _, Some value -> Delta_runtime.Pure_map.add key value state.%s) in"
+               cache cache cache cache);
+          line buffer (Printf.sprintf "    let state = { state with %s = cached } in" cache);
+          if consumers = [] then line buffer "    state"
+          else line buffer (forward "old_member" "new_member")
+      | Graph.Map (_, _) ->
+          let cached_values = node.Graph.n_cache = Graph.Cached_values in
+          let element = node.Graph.n_element in
+          (if cached_values then
+             line buffer
+               (Printf.sprintf "    let old_value = Delta_runtime.Pure_map.find_opt key state.%s in"
+                  (cache_field node))
+           else line buffer "    let old_value = None in");
+          line buffer
+            (Printf.sprintf
+               "    let new_value = (match up_new with None -> None | Some value -> %s) in"
+               (if cached_values then
+                  Printf.sprintf
+                    "(match up_old with None -> Some (mapping_%d value) | Some previous -> (match old_value with Some cached_value -> Some (apply_%s cached_value (delta_%d previous value)) | None -> Some (mapping_%d value)))"
+                    node.Graph.n_id (type_name element) node.Graph.n_id node.Graph.n_id
+                else Printf.sprintf "Some (mapping_%d value)" node.Graph.n_id));
+          (if cached_values then
+             let cache = cache_field node in
+             line buffer
+               (Printf.sprintf
+                  "    let cached = (match (old_value, new_value) with None, None -> state.%s | None, Some value -> Delta_runtime.Pure_map.add key value state.%s | Some _, None -> Delta_runtime.Pure_map.remove key state.%s | Some _, Some value -> Delta_runtime.Pure_map.add key value state.%s) in"
+                  cache cache cache cache);
+             line buffer (Printf.sprintf "    let state = { state with %s = cached } in" cache)
+           else
+             line buffer
+               (Printf.sprintf
+                  "    let old_value = (match up_old with None -> None | Some previous -> Some (mapping_%d previous)) in"
+                  node.Graph.n_id));
+          if consumers = [] then line buffer "    state" else line buffer (forward "old_value" "new_value")
+      | Graph.Sum ->
+          line buffer
+            "    let delta = (match up_new with None -> 0 | Some value -> value) - (match up_old with None -> 0 | Some value -> value) in";
+          if consumers = [] then
+            line buffer
+              (Printf.sprintf "    { state with %s = state.%s + delta }" (accumulator_field node)
+                 (accumulator_field node))
+          else (
+            line buffer
+              (Printf.sprintf "    let state = { state with %s = state.%s + delta } in"
+                 (accumulator_field node) (accumulator_field node));
+            line buffer (forward "None" "None"))
+      | Graph.Count ->
+          line buffer
+            "    let delta = (match (up_old, up_new) with None, None -> 0 | None, Some _ -> 1 | Some _, None -> -1 | Some _, Some _ -> 0) in";
+          if consumers = [] then
+            line buffer
+              (Printf.sprintf "    { state with %s = state.%s + delta }" (accumulator_field node)
+                 (accumulator_field node))
+          else (
+            line buffer
+              (Printf.sprintf "    let state = { state with %s = state.%s + delta } in"
+                 (accumulator_field node) (accumulator_field node));
+            line buffer (forward "None" "None"));
+      line buffer "")
+    nodes
+
+let root_observation plan =
+  match plan.Graph.pl_result with
+  | Graph.Result_collection id -> (
+      let node = Graph.node_of_id plan id in
+      match node.Graph.n_kind with
+      | Graph.Source -> "Delta_runtime.Pure_map.find_opt key state.input"
+      | Graph.Filter _ | Graph.Map _ ->
+          Printf.sprintf "Delta_runtime.Pure_map.find_opt key state.%s" (cache_field node)
+      | Graph.Sum | Graph.Count -> Printf.sprintf "state.%s" (accumulator_field node))
+  | Graph.Result_scalar _ -> "0"
+
+let emit_apply_key plan buffer =
+  let input_element = input_type plan in
+  let source = match source_node plan with Some node -> node.Graph.n_id | None -> -1 in
+  let scalar_env =
+    List.map
+      (fun (stamp, node_id) ->
+        (stamp, Printf.sprintf "state.%s" (accumulator_field (Graph.node_of_id plan node_id))))
+      plan.Graph.pl_scalar_bindings
+  in
+  line buffer
+    (Printf.sprintf
+       "  let apply_key (state : state) (key : int) (up_old : %s option) (up_new : %s option) : state * output_change ="
+       (ocaml_ty input_element) (ocaml_ty input_element));
+  line buffer (Printf.sprintf "    let before = %s in" (root_observation plan));
+  line buffer (Printf.sprintf "    let state = update_%d state key up_old up_new in" source);
+  (match plan.Graph.pl_result with
+  | Graph.Result_collection id -> (
+      let node = Graph.node_of_id plan id in
+      match node.Graph.n_kind with
+      | Graph.Sum | Graph.Count ->
+          line buffer
+            (Printf.sprintf "    (state, state.%s - before)" (accumulator_field node))
+      | Graph.Source | Graph.Filter _ | Graph.Map _ ->
+          line buffer (Printf.sprintf "    let after = %s in" (root_observation plan));
+          line buffer
+            "    let change = (match (before, after) with None, None -> [] | None, Some value -> [ (key, Ins value) ] | Some value, None -> [ (key, Rem value) ] | Some old_value, Some new_value -> if old_value = new_value then [] else [ (key, Rep (old_value, new_value)) ]) in";
+          line buffer "    (state, change)")
+  | Graph.Result_scalar expr ->
+      let names = names_of_idents (binders_of expr) in
+      let env = names @ scalar_env in
+      line buffer (Printf.sprintf "    let after = %s in" (render_expr env expr));
+      line buffer "    (state, after - before)");
+  line buffer ""
+
+let emit_apply_batch plan buffer =
+  let input_element = input_type plan in
+  let empty_output = if is_collection_output plan then "[]" else "0" in
+  let combine = if is_collection_output plan then "(!output @ change)" else "(!output + change)" in
+  line buffer
+    (Printf.sprintf
+       "  let apply_batch (state : state) (ops : batch_op list) : (state * output_change) Delta_runtime.outcome ="
+       );
+  line buffer "    let rec validate validated touched pending =";
+  line buffer "      match pending with";
+  line buffer "      | [] -> Delta_runtime.Success (validated, touched)";
+  line buffer
+    "      | Insert (key, value) :: rest -> if Delta_runtime.Pure_map.mem key validated then Delta_runtime.Failure (Printf.sprintf \"cannot insert key %d: it is already present\" key) else validate (Delta_runtime.Pure_map.add key value validated) (key :: touched) rest";
+  line buffer
+    "      | Remove key :: rest -> if Delta_runtime.Pure_map.mem key validated then validate (Delta_runtime.Pure_map.remove key validated) (key :: touched) rest else Delta_runtime.Failure (Printf.sprintf \"cannot remove key %d: it is not present\" key)";
+  line buffer
+    "      | Replace (key, value) :: rest -> if Delta_runtime.Pure_map.mem key validated then validate (Delta_runtime.Pure_map.add key value validated) (key :: touched) rest else Delta_runtime.Failure (Printf.sprintf \"cannot replace key %d: it is not present\" key)";
+  line buffer "    in";
+  line buffer "    let original = state.input in";
+  line buffer "    let rec collect validated keys pending =";
+  line buffer "      match pending with";
+  line buffer "      | [] -> List.rev keys";
+  line buffer "      | key :: rest -> (";
+  line buffer
+    "          match (Delta_runtime.Pure_map.find_opt key original, Delta_runtime.Pure_map.find_opt key validated) with";
+  line buffer "          | None, None -> collect validated keys rest";
+  line buffer "          | None, Some value -> collect validated ((key, None, Some value) :: keys) rest";
+  line buffer "          | Some value, None -> collect validated ((key, Some value, None) :: keys) rest";
+  line buffer
+    "          | Some old_value, Some new_value -> if old_value = new_value then collect validated keys rest else collect validated ((key, Some old_value, Some new_value) :: keys) rest)";
+  line buffer "    in";
+  line buffer "    match validate original [] ops with";
+  line buffer "    | Delta_runtime.Failure message -> Delta_runtime.Failure message";
+  line buffer "    | Delta_runtime.Success (validated, touched) ->";
+  line buffer "        let items = collect validated [] (List.sort_uniq compare touched) in";
+  line buffer "        (try";
+  line buffer "           let current = ref state in";
+  line buffer (Printf.sprintf "           let output = ref %s in" empty_output);
+  line buffer "           List.iter (fun (key, up_old, up_new) ->";
+  line buffer "             let (next_state, change) = apply_key !current key up_old up_new in";
+  line buffer "             current := next_state;";
+  line buffer (Printf.sprintf "             output := %s) items;" combine);
+  line buffer "           Delta_runtime.Success (!current, !output)";
+  line buffer "         with";
+  line buffer "         | Division_by_zero -> Delta_runtime.Failure \"division by zero while applying the batch\")";
+  line buffer ""
+
 let emit_structure plan buffer =
   emit_output_types plan buffer;
   emit_state plan buffer;
+  emit_batch_type plan buffer;
+  emit_delta_functions plan buffer;
   emit_node_functions plan buffer;
+  emit_update_functions plan buffer;
+  emit_apply_key plan buffer;
+  emit_apply_batch plan buffer;
   emit_init plan buffer;
   emit_result plan buffer;
   emit_output_functions plan buffer;
@@ -572,3 +859,5 @@ let program_to_string plan =
   emit_structure plan buffer;
   line buffer "end";
   Buffer.contents buffer
+
+