Build collection dependencies and cache layouts

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 ()