Merge branch 'compile-clause-early' of https://github.com/LPCIC/elpi into compile-clause-early

Author: FissoreD

src/compiler.ml

@@ -718,7 +718,7 @@ end = struct (* {{{ *)
            | Some "Map" -> Some Map
            | Some "Hash" -> Some HashMap
            | Some "DTree" -> Some DiscriminationTree
-           | Some s -> error ~loc ("unknown indexing directive " ^ s) in
+           | Some s -> error ~loc ("unknown indexing directive " ^ s ^ ". Valid ones are: Map, Hash, DTree.") in
          begin match r with
            | None -> aux_tatt (Some (Structured.Index(i,it))) rest
            | Some (Structured.Index _) -> duplicate_err "index"

src/discrimination_tree.ml

@@ -165,6 +165,14 @@ module Trie = struct
 
   let is_empty x = x == empty
 
+  let rec replace p x = function
+    | Node { data; other; listTailVariable; map } ->
+        Node {
+          data = data |> List.map (fun y -> if p y then x else y);
+          other = other |> Option.map (replace p x);
+          listTailVariable = listTailVariable |> Option.map (replace p x);
+          map = map |> Ptmap.map (replace p x);
+        }
 
   let add (a : Path.t) v t =
     let max = ref 0 in
@@ -383,6 +391,7 @@ let retrieve cmp_data path { t } =
   let r = call (retrieve ~pos:0 path t) in
   Bl.of_list @@ List.sort cmp_data r
 
+let replace p x i = { i with t = Trie.replace p x i.t }
 
 module Internal = struct
   let kConstant = kConstant

src/discrimination_tree.mli

@@ -51,6 +51,8 @@ val empty_dt : 'b list -> 'a t
 *)
 val retrieve : ('a -> 'a -> int) -> Path.t -> 'a t -> 'a Bl.scan
 
+val replace : ('a -> bool) -> 'a -> 'a t -> 'a t
+
 (***********************************************************)
 (* Printers                                                *)
 (***********************************************************)

src/runtime.ml

@@ -2710,7 +2710,25 @@ let add1clause ~depth { idx; time; times } ~time_dir ~insert ~empty ~cons ~copy
           StrMap.add id clause.timestamp times in
   let idx =
     try
-      add1clause2 ~depth idx ~insert ~empty ~copy graft grafting_reference predicate clause (Ptmap.find predicate idx);
+      (* TODO: do this only at compile time *)
+      match graft with
+      | Some (Elpi_parser.Ast.Structured.Replace _) ->
+          Ptmap.map (function
+          | TwoLevelIndex {
+            argno; mode;
+            all_clauses;
+            flex_arg_clauses;
+            arg_idx;
+            } -> TwoLevelIndex {
+              argno; mode;
+              all_clauses = insert graft grafting_reference clause all_clauses;
+              flex_arg_clauses = insert graft grafting_reference clause flex_arg_clauses;
+              arg_idx = Ptmap.map (fun l -> insert graft grafting_reference clause l) arg_idx;
+              }
+          | BitHash { mode; args; args_idx } -> BitHash { mode; args; args_idx = Ptmap.map (fun l -> insert graft grafting_reference clause l) args_idx }
+          | IndexWithDiscriminationTree {mode; arg_depths; args_idx; } -> IndexWithDiscriminationTree {mode; arg_depths; args_idx = Discrimination_tree.replace (fun x -> x.timestamp = grafting_reference) clause args_idx; }
+          ) idx
+      | _ -> add1clause2 ~depth idx ~insert ~empty ~copy graft grafting_reference predicate clause (Ptmap.find predicate idx);
     with
     | Not_found ->
       match classify_clause_argno ~depth 0 [] clause.args with

tests/sources/graft_replace_ok.elpi

@@ -1,3 +1,4 @@
+:index (1) % "Hash" "DTree"
 pred p o:int.
 :name "replace_me"
 p 1.
@@ -6,4 +7,5 @@ p 1.
 p 2.
 
 main :-
-  std.findall (p _) [p 2].
+  std.findall (p _) [p 2],
+  std.findall (p 1) [].