Build collection dependencies and cache layouts

example/count_large.delta

@@ -0,0 +1,9 @@
+type order = {
+  customer : string;
+  total : int;
+}
+
+input orders : collection order
+
+query count_large =
+  orders |> filter (fun o -> o.total > 500) |> count

example/revenue.delta

@@ -0,0 +1,11 @@
+type order = {
+  customer : string;
+  total : int;
+}
+
+input orders : collection order
+
+let tax n = n * 20 / 100
+
+query revenue =
+  orders |> map (fun o -> tax o.total) |> sum

src/graph.ml

@@ -0,0 +1,202 @@
+type node_kind =
+  | Source
+  | Filter of Ident.t * Anf.expr
+  | Map of Ident.t * Anf.expr
+  | Sum
+  | Count
+
+type cache =
+  | No_cache
+  | Cached_values
+  | Accumulator
+
+type node = {
+  n_id : int;
+  n_kind : node_kind;
+  n_input : int option;
+  n_element : Types.t;
+  n_cache : cache;
+  n_span : Location.span;
+}
+
+type result =
+  | Result_collection of int
+  | Result_scalar of Anf.expr
+
+type plan = {
+  pl_nodes : node list;
+  pl_result : result;
+  pl_output : Types.t;
+  pl_output_element : Types.t option;
+  pl_consumers : int list array;
+}
+
+let node_of_id plan id =
+  let rec search = function
+    | [] -> Diagnostic.error Location.none "internal error: unknown plan node %d" id
+    | node :: rest -> if node.n_id = id then node else search rest
+  in
+  search plan.pl_nodes
+
+let build program =
+  let nodes = ref [] in
+  let table = Hashtbl.create 16 in
+  let consumers = Hashtbl.create 16 in
+  let input_node = ref None in
+  let add kind input element cache span =
+    let id = List.length !nodes in
+    nodes :=
+      !nodes
+      @ [
+          {
+            n_id = id;
+            n_kind = kind;
+            n_input = input;
+            n_element = element;
+            n_cache = cache;
+            n_span = span;
+          };
+        ];
+    (match input with
+    | Some up ->
+        let existing = try Hashtbl.find consumers up with Not_found -> [] in
+        Hashtbl.replace consumers up (id :: existing)
+    | None -> ());
+    id
+  in
+  let element_of span ty =
+    match Types.repr ty with
+    | Types.TCollection element -> element
+    | _ -> Diagnostic.error span "internal error: expected a collection type"
+  in
+  let find node_id =
+    let rec search = function
+      | [] -> Diagnostic.error Location.none "internal error: unknown plan node %d" node_id
+      | node :: rest -> if node.n_id = node_id then node else search rest
+    in
+    search !nodes
+  in
+  let source span =
+    match !input_node with
+    | Some id -> id
+    | None ->
+        let id = add Source None program.Anf.ap_input_element Cached_values span in
+        input_node := Some id;
+        id
+  in
+  let node_of_atom atom span =
+    match atom with
+    | Anf.AVar ident -> (
+        match Util.hashtbl_find_opt table (Ident.stamp ident) with
+        | Some id -> id
+        | None ->
+            if Ident.stamp ident = Ident.stamp program.Anf.ap_input then source span
+            else
+              Diagnostic.error span "internal error: `%s` is not a collection in this plan"
+                (Ident.display ident))
+    | _ -> Diagnostic.error span "internal error: expected a collection variable"
+  in
+  let rec bind_bound bound =
+    let span = bound.Anf.aspan in
+    match bound.Anf.a with
+    | Anf.AFilter (source_atom, parameter, body) ->
+        let input = node_of_atom source_atom span in
+        Some (add (Filter (parameter, body)) (Some input) (find input).n_element Cached_values span)
+    | Anf.AMap (source_atom, parameter, body) ->
+        let input = node_of_atom source_atom span in
+        Some (add (Map (parameter, body)) (Some input) (element_of span bound.Anf.aty) Cached_values span)
+    | Anf.ASum source_atom ->
+        let input = node_of_atom source_atom span in
+        Some (add Sum (Some input) Types.TInt Accumulator span)
+    | Anf.ACount source_atom ->
+        let input = node_of_atom source_atom span in
+        Some (add Count (Some input) Types.TInt Accumulator span)
+    | Anf.AAtom (Anf.AVar ident) -> (
+        match Util.hashtbl_find_opt table (Ident.stamp ident) with
+        | Some id -> Some id
+        | None ->
+            if Ident.stamp ident = Ident.stamp program.Anf.ap_input then Some (source span) else None)
+    | _ -> None
+  in
+  let rec walk expr =
+    match expr.Anf.a with
+    | Anf.ALet (ident, bound, body) ->
+        (match bind_bound bound with
+        | Some id -> Hashtbl.replace table (Ident.stamp ident) id
+        | None -> ());
+        walk body
+    | Anf.AAtom (Anf.AVar ident) -> (
+        match Util.hashtbl_find_opt table (Ident.stamp ident) with
+        | Some id -> Result_collection id
+        | None ->
+            if Ident.stamp ident = Ident.stamp program.Anf.ap_input then
+              Result_collection (source expr.Anf.aspan)
+            else Result_scalar expr)
+    | _ -> Result_scalar expr
+  in
+  let result = walk program.Anf.ap_query_body in
+  let output = program.Anf.ap_query_body.Anf.aty in
+  let output_element = match Types.repr output with Types.TCollection element -> Some element | _ -> None in
+  (match (result, Types.repr output) with
+  | Result_collection id, Types.TCollection _ ->
+      let node = find id in
+      (match node.n_kind with
+      | Sum | Count ->
+          Diagnostic.error node.n_span
+            "internal error: a collection query cannot end in an aggregate"
+      | Source | Filter _ | Map _ -> ())
+  | Result_collection id, Types.TInt -> (
+      let node = find id in
+      match node.n_kind with
+      | Sum | Count -> ()
+      | _ ->
+          Diagnostic.error node.n_span "internal error: an integer query must end in an aggregate")
+  | Result_scalar _, Types.TInt -> ()
+  | Result_collection _, other | Result_scalar _, other ->
+      Diagnostic.error program.Anf.ap_query_body.Anf.aspan
+        "internal error: unexpected plan result for output type %s" (Types.pp other));
+  let consumers =
+    Array.init (List.length !nodes) (fun id ->
+        match Util.hashtbl_find_opt consumers id with Some ids -> List.rev ids | None -> [])
+  in
+  {
+    pl_nodes = !nodes;
+    pl_result = result;
+    pl_output = output;
+    pl_output_element = output_element;
+    pl_consumers = consumers;
+  }
+
+let kind_to_string node =
+  match node.n_kind with
+  | Source -> "source"
+  | Filter (parameter, _) -> "filter (fun " ^ Ident.display parameter ^ " -> ...)"
+  | Map (parameter, _) -> "map (fun " ^ Ident.display parameter ^ " -> ...)"
+  | Sum -> "sum"
+  | Count -> "count"
+
+let cache_to_string = function
+  | No_cache -> "none"
+  | Cached_values -> "values"
+  | Accumulator -> "accumulator"
+
+let result_to_string plan =
+  match plan.pl_result with
+  | Result_collection id -> Printf.sprintf "node %d" id
+  | Result_scalar expr -> "scalar " ^ Anf.to_string expr
+
+let dump plan =
+  let lines =
+    List.map
+      (fun node ->
+        let input = match node.n_input with Some id -> Printf.sprintf " over node %d" id | None -> "" in
+        Printf.sprintf "  node %d: %s%s : %s [cache: %s, consumers: %s]" node.n_id
+          (kind_to_string node) input (Types.pp node.n_element) (cache_to_string node.n_cache)
+          (Util.join "," (List.map string_of_int plan.pl_consumers.(node.n_id))))
+      plan.pl_nodes
+  in
+  Util.join "\n"
+    ([ Printf.sprintf "output : %s" (Types.pp plan.pl_output);
+       Printf.sprintf "result : %s" (result_to_string plan) ]
+    @ lines)
+  ^ "\n"

src/graph.mli

@@ -0,0 +1,38 @@
+type node_kind =
+  | Source
+  | Filter of Ident.t * Anf.expr
+  | Map of Ident.t * Anf.expr
+  | Sum
+  | Count
+
+type cache =
+  | No_cache
+  | Cached_values
+  | Accumulator
+
+type node = {
+  n_id : int;
+  n_kind : node_kind;
+  n_input : int option;
+  n_element : Types.t;
+  n_cache : cache;
+  n_span : Location.span;
+}
+
+type result =
+  | Result_collection of int
+  | Result_scalar of Anf.expr
+
+type plan = {
+  pl_nodes : node list;
+  pl_result : result;
+  pl_output : Types.t;
+  pl_output_element : Types.t option;
+  pl_consumers : int list array;
+}
+
+val build : Anf.program -> plan
+val node_of_id : plan -> int -> node
+val dump : plan -> string
+val kind_to_string : node -> string
+val cache_to_string : cache -> string

src/main.ml

@@ -95,15 +95,18 @@ let specialize_source path = Specialize.program (infer_source path)
 
 let anf_source path = Anf.program (specialize_source path)
 
+let plan_source path = Graph.build (anf_source path)
+
 let frontend_unavailable () =
   Diagnostic.error Location.none "the delta pipeline beyond parsing is not implemented in this revision"
 
 let check path =
-  ignore (anf_source path)
+  ignore (plan_source path)
 
 let dump stage path =
   match stage with
   | "anf" -> print_string (Anf.program_to_string (anf_source path))
+  | "delta" -> print_string (Graph.dump (plan_source path))
   | _ ->
   let program = infer_source path in
   match stage with
@@ -111,12 +114,12 @@ let dump stage path =
   | _ -> frontend_unavailable ()
 
 let emit path output =
-  ignore (anf_source path);
+  ignore (plan_source path);
   ignore output;
   frontend_unavailable ()
 
 let build path output =
-  ignore (anf_source path);
+  ignore (plan_source path);
   ignore output;
   frontend_unavailable ()
 

test/test_incremental.ml

@@ -120,3 +120,80 @@ let cases =
         in
         check_equal_string "wrapped" (string_of_int min_int) (Value.to_string result) );
   ]
+
+let plan text = Graph.build (Anf.program (Specialize.program (infer text)))
+
+let plan_of_fixture name = Graph.build (Anf.program (Specialize.program (infer (read_fixture name))))
+
+let graph_cases =
+  [
+    ( "the example plan is a source, a filter and a map",
+      fun () ->
+        let plan = plan_of_fixture "expensive_order.delta" in
+        check_equal_int "three nodes" 3 (List.length plan.Graph.pl_nodes);
+        let nodes = plan.Graph.pl_nodes in
+        check "source first" (match (List.nth nodes 0).Graph.n_kind with Graph.Source -> true | _ -> false);
+        check "filter second"
+          (match (List.nth nodes 1).Graph.n_kind with Graph.Filter _ -> true | _ -> false);
+        check "map last" (match (List.nth nodes 2).Graph.n_kind with Graph.Map _ -> true | _ -> false);
+        check_equal_string "output" "collection (string, int)" (Types.pp plan.Graph.pl_output) );
+    ( "the revenue plan ends in a sum accumulator",
+      fun () ->
+        let plan = plan_of_fixture "revenue.delta" in
+        let root = Graph.node_of_id plan (match plan.Graph.pl_result with Graph.Result_collection id -> id | _ -> -1) in
+        check "sum" (match root.Graph.n_kind with Graph.Sum -> true | _ -> false);
+        check "accumulator cache" (root.Graph.n_cache = Graph.Accumulator);
+        check_equal_string "linear output" "int" (Types.pp plan.Graph.pl_output) );
+    ( "the count query counts the retained rows",
+      fun () ->
+        let plan = plan_of_fixture "count_large.delta" in
+        check_equal_int "three nodes" 3 (List.length plan.Graph.pl_nodes);
+        let root =
+          match plan.Graph.pl_result with
+          | Graph.Result_collection id -> Graph.node_of_id plan id
+          | Graph.Result_scalar _ ->
+              fail "plan" "expected a node result";
+              { Graph.n_id = -1; n_kind = Graph.Source; n_input = None; n_element = Types.TInt; n_cache = Graph.No_cache; n_span = Location.none }
+        in
+        check "count" (match root.Graph.n_kind with Graph.Count -> true | _ -> false) );
+    ( "collection nodes record their consumers",
+      fun () ->
+        let plan = plan_of_fixture "expensive_order.delta" in
+        check_equal_string "source consumers" "1" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(0)));
+        check_equal_string "filter consumers" "2" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(1)));
+        check_equal_string "map has no consumers" "" (Util.join "," (List.map string_of_int plan.Graph.pl_consumers.(2))) );
+    ( "an identity query has a single source node",
+      fun () ->
+        let plan = plan "input rows : collection int\nquery q = rows\n" in
+        check_equal_int "one node" 1 (List.length plan.Graph.pl_nodes);
+        check "the result is the source"
+          (match plan.Graph.pl_result with Graph.Result_collection 0 -> true | _ -> false) );
+    ( "a plan dump is deterministic",
+      fun () ->
+        Ident.reset ();
+        Types.reset ();
+        Anf.reset ();
+        let first = Graph.dump (plan_of_fixture "expensive_order.delta") in
+        Ident.reset ();
+        Types.reset ();
+        Anf.reset ();
+        let second = Graph.dump (plan_of_fixture "expensive_order.delta") in
+        check_equal_string "identical" first second );
+    ( "an integer query that only uses the input is a plain plan",
+      fun () ->
+        let plan = plan "input rows : collection int\nquery q = rows |> sum\n" in
+        check_equal_int "two nodes" 2 (List.length plan.Graph.pl_nodes);
+        check "integer output" (Types.repr plan.Graph.pl_output = Types.TInt) );
+    ( "a constant integer query produces no collection nodes",
+      fun () ->
+        let plan = plan "input rows : collection int\nquery q = 40 + 2\n" in
+        check_equal_int "no nodes" 0 (List.length plan.Graph.pl_nodes);
+        check "scalar result" (match plan.Graph.pl_result with Graph.Result_scalar _ -> true | _ -> false) );
+    ( "shared collection temporaries are reuse",
+      fun () ->
+        let plan =
+          plan
+            "input rows : collection int\nquery q = rows |> filter (fun r -> r > 0) |> map (fun r -> r + 1) |> sum\n"
+        in
+        check_equal_int "four nodes" 4 (List.length plan.Graph.pl_nodes) );
+  ]

test/test_main.ml

@@ -379,5 +379,6 @@ let () =
   Test_harness.run_suite "changes" Test_change.law_cases;
   Test_harness.run_suite "batches" Test_change.batch_cases;
   Test_harness.run_suite "interpret" Test_incremental.cases;
+  Test_harness.run_suite "graph" Test_incremental.graph_cases;
   Printf.printf "%d cases, %d failures\n" (Test_harness.case_count ()) (Test_harness.failure_count ());
   exit (if Test_harness.failure_count () = 0 then 0 else 1)