Datalog: add [not_equal] and [filter] predicates (#4018)

Co-authored-by: Basile Clément <[email protected]>

Author: Ekdohibs

middle_end/flambda2/datalog/cursor.ml

@@ -15,6 +15,10 @@
 
 open Datalog_imports
 
+type _ value_repr = Int_repr : int value_repr
+
+let int_repr = Int_repr
+
 (* Note: we don't use [with_name] here to avoid the extra indirection during
    execution. *)
 type vm_action =
@@ -25,6 +29,12 @@ type vm_action =
       * string
       * string list
       -> vm_action
+  | Unless_eq :
+      'k option ref * 'k option ref * string * string * 'k value_repr
+      -> vm_action
+  | Filter :
+      ('k Constant.hlist -> bool) * 'k Option_ref.hlist * string list
+      -> vm_action
 
 type action =
   | Bind_iterator :
@@ -39,6 +49,11 @@ let unless id cell args =
     (Unless
        (Table.Id.is_trie id, cell, args.values, Table.Id.name id, args.names))
 
+let unless_eq repr cell1 cell2 =
+  VM_action (Unless_eq (cell1.value, cell2.value, cell1.name, cell2.name, repr))
+
+let filter f args = VM_action (Filter (f, args.values, args.names))
+
 type binder = Bind_table : ('t, 'k, 'v) Table.Id.t * 't ref -> binder
 
 type actions = { mutable rev_actions : action list }
@@ -55,6 +70,14 @@ let pp_cursor_action ff = function
          ~pp_sep:(fun ff () -> Format.fprintf ff ", ")
          Format.pp_print_string)
       l_names
+  | Unless_eq (_x1, _x2, x1_name, x2_name, _repr) ->
+    Format.fprintf ff "if %s == %s:@ continue" x1_name x2_name
+  | Filter (_f, _args, args_names) ->
+    Format.fprintf ff "<filter>(%a)"
+      (Format.pp_print_list
+         ~pp_sep:(fun ff () -> Format.fprintf ff ", ")
+         Format.pp_print_string)
+      args_names
 
 module Order : sig
   type t
@@ -323,6 +346,14 @@ let evaluate = function
          (Trie.find_opt is_trie (Option_ref.get args) cell.contents)
     then Virtual_machine.Skip
     else Virtual_machine.Accept
+  | Unless_eq (cell1, cell2, _cell1_name, _cell2_name, Int_repr) ->
+    if Int.equal (Option.get !cell1) (Option.get !cell2)
+    then Virtual_machine.Skip
+    else Virtual_machine.Accept
+  | Filter (f, args, _args_names) ->
+    if f (Option_ref.get args)
+    then Virtual_machine.Accept
+    else Virtual_machine.Skip
 
 let naive_iter cursor db f =
   with_bound_cursor ~callback:f cursor db @@ fun () ->

middle_end/flambda2/datalog/cursor.mli

@@ -15,6 +15,10 @@
 
 open Datalog_imports
 
+type 'a value_repr
+
+val int_repr : int value_repr
+
 type action
 
 val bind_iterator :
@@ -23,6 +27,12 @@ val bind_iterator :
 val unless :
   ('t, 'k, 'v) Table.Id.t -> 't ref -> 'k Option_ref.hlist with_names -> action
 
+val unless_eq :
+  'k value_repr -> 'k option ref with_name -> 'k option ref with_name -> action
+
+val filter :
+  ('k Constant.hlist -> bool) -> 'k Option_ref.hlist with_names -> action
+
 type actions
 
 val add_action : actions -> action -> unit

middle_end/flambda2/datalog/datalog.ml

@@ -156,20 +156,19 @@ let rec find_last_binding0 : type a. order:_ -> _ -> a Term.hlist -> _ =
 let find_last_binding post_level args =
   find_last_binding0 ~order:Cursor.Order.parameters post_level args
 
+let compile_term : 'a Term.t -> 'a option ref with_name = function
+  | Constant cte -> { value = ref (Some cte); name = "<constant>" }
+  | Parameter param -> { value = param.cell; name = param.name }
+  | Variable var -> Cursor.Level.use_output var
+
 let rec compile_terms : type a. a Term.hlist -> a Option_ref.hlist with_names =
  fun vars ->
   match vars with
   | [] -> { values = []; names = [] }
-  | term :: terms -> (
+  | term :: terms ->
+    let { value; name } = compile_term term in
     let { values; names } = compile_terms terms in
-    match term with
-    | Constant cte ->
-      { values = ref (Some cte) :: values; names = "<constant>" :: names }
-    | Parameter param ->
-      { values = param.cell :: values; names = param.name :: names }
-    | Variable var ->
-      let { value; name } = Cursor.Level.use_output var in
-      { values = value :: values; names = name :: names })
+    { values = value :: values; names = name :: names }
 
 let unless_atom id args k info =
   let refs = compile_terms args in
@@ -180,6 +179,23 @@ let unless_atom id args k info =
   Cursor.add_action post_level (Cursor.unless id r refs);
   k info
 
+let unless_eq repr arg1 arg2 k info =
+  let ref1 = compile_term arg1 in
+  let ref2 = compile_term arg2 in
+  let post_level =
+    find_last_binding (Cursor.initial_actions info.context) [arg1; arg2]
+  in
+  Cursor.add_action post_level (Cursor.unless_eq repr ref1 ref2);
+  k info
+
+let filter f args k info =
+  let refs = compile_terms args in
+  let post_level =
+    find_last_binding (Cursor.initial_actions info.context) args
+  in
+  Cursor.add_action post_level (Cursor.filter f refs);
+  k info
+
 type callback =
   | Callback :
       { func : 'a Constant.hlist -> unit;

middle_end/flambda2/datalog/datalog.mli

@@ -63,6 +63,19 @@ val unless_atom :
   ('p, 'a) program ->
   ('p, 'a) program
 
+val unless_eq :
+  'k Cursor.value_repr ->
+  'k Term.t ->
+  'k Term.t ->
+  ('p, 'a) program ->
+  ('p, 'a) program
+
+val filter :
+  ('k Constant.hlist -> bool) ->
+  'k Term.hlist ->
+  ('p, 'a) program ->
+  ('p, 'a) program
+
 type callback
 
 val create_callback :

middle_end/flambda2/datalog/flambda2_datalog.ml

@@ -60,6 +60,9 @@ module Datalog = struct
       | { repr = Patricia_tree_repr; _ } :: (_ :: _ as columns) ->
         Trie.patricia_tree_of_trie (is_trie columns)
 
+    let key_repr : type t k v. (t, k, v) id -> k Cursor.value_repr = function
+      | { repr = Patricia_tree_repr; _ } -> Cursor.int_repr
+
     module type S = sig
       type t
 
@@ -245,20 +248,33 @@ module Datalog = struct
 
   let deduce = Schedule.deduce
 
+  type equality =
+    | Equality : 'k Cursor.value_repr * 'k Term.t * 'k Term.t -> equality
+
+  type filter = Filter : ('k Constant.hlist -> bool) * 'k Term.hlist -> filter
+
   type hypothesis =
     [ `Atom of atom
-    | `Not_atom of atom ]
+    | `Not_atom of atom
+    | `Distinct of equality
+    | `Filter of filter ]
 
   let atom id args = `Atom (Atom (id, args))
 
   let not (`Atom atom) = `Not_atom atom
 
+  let distinct c x y = `Distinct (Equality (Column.key_repr c, x, y))
+
+  let filter f args = `Filter (Filter (f, args))
+
   let where predicates f =
     List.fold_left
       (fun f predicate ->
         match predicate with
         | `Atom (Atom (id, args)) -> where_atom id args f
-        | `Not_atom (Atom (id, args)) -> unless_atom id args f)
+        | `Not_atom (Atom (id, args)) -> unless_atom id args f
+        | `Distinct (Equality (repr, t1, t2)) -> unless_eq repr t1 t2 f
+        | `Filter (Filter (p, args)) -> Datalog.filter p args f)
       f predicates
 
   module Cursor = struct

middle_end/flambda2/datalog/flambda2_datalog.mli

@@ -132,9 +132,15 @@ module Datalog : sig
 
   type atom
 
+  type equality
+
+  type filter
+
   type hypothesis =
     [ `Atom of atom
-    | `Not_atom of atom ]
+    | `Not_atom of atom
+    | `Distinct of equality
+    | `Filter of filter ]
 
   (** [atom rel args] represents the application of relation [rel] to the
       arguments [args].
@@ -153,6 +159,12 @@ module Datalog : sig
 
   val not : [< `Atom of atom] -> [> `Not_atom of atom]
 
+  val distinct :
+    (_, 'k, _) Column.id -> 'k Term.t -> 'k Term.t -> [> `Distinct of equality]
+
+  val filter :
+    ('k Constant.hlist -> bool) -> 'k Term.hlist -> [> `Filter of filter]
+
   type database
 
   val print : Format.formatter -> database -> unit