[subst] A refactoring to do less allocations

Signed-off-by: Dmitrii.Kosarev a.k.a. Kakadu <kakadu@pm.me>
PLTools · Jan 12, 2025 · 449e680 · 449e680
1 parent aafc2e3
commit 449e680
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 45 deletions.
diff --git a/src/core/Disequality.ml b/src/core/Disequality.ml
@@ -185,12 +185,8 @@ module Disjunct :
           | None ->
             (* not unifiable --- always distinct *)
             raise Disequality_fulfilled
-          | Some ([], _) -> raise Disequality_violated
-          | Some (bnds, _subst) ->
-              (* TODO(Kakadu): reconstruction of map from binding list could hurt performance  *)
-              let rez = Subst.varmap_of_bindings bnds in
-              (* log "Disjunct.recheck returns %a" pp rez; *)
-              rez)
+          | Some (delta, _) when Term.VarMap.is_empty delta -> raise Disequality_violated
+          | Some (bnds, _subst) -> bnds)
       | Violated       ->
         let unchecked = Term.VarMap.remove var t in
         (* log "       unchecked: %a" pp unchecked; *)

diff --git a/src/core/Subst.ml b/src/core/Subst.ml
@@ -160,54 +160,59 @@ let log fmt =
   then Format.kasprintf (Format.printf "%s\n%!") fmt
   else Format.ifprintf Format.std_formatter fmt
 
-let unify ?(subsume=false) ?(scope=Term.Var.non_local_scope) env subst x y =
+let ext ~scope ~env add_delta var term (prefix, subst) =
+  let subst = extend ~scope env subst var term in
+  (add_delta Binding.{var; term} prefix, subst)
+
+let rec unify_helper subsume ext env x y acc =
+  (* log "unify '%s' and '%s'" (Term.show x) (Term.show y); *)
+  let open Term in
+  fold2 x y ~init:acc
+    ~fvar:(fun ((_, subst) as acc) x y ->
+      match walk env subst x, walk env subst y with
+      | Var x, Var y      ->
+        if Var.equal x y then acc else ext x (Term.repr y) acc
+      | Var x, Value y    -> ext x y acc
+      | Value x, Var y    -> ext y x acc
+      | Value x, Value y  -> unify_helper subsume ext env x y acc
+    )
+    ~fval:(fun acc x y ->
+        if x = y then acc else raise Unification_failed
+    )
+    ~fk:(fun ((_, subst) as acc) l v y ->
+        if subsume && (l = Term.R)
+        then raise Unification_failed
+        else match walk env subst v with
+        | Var v    -> ext v y acc
+        | Value x  -> unify_helper subsume ext env x y acc
+    )
+
+let unify_gen ?(subsume=false) ?(scope=Term.Var.non_local_scope) add_delta empty_delta env subst x y =
   (* The idea is to do the unification and collect the unification prefix during the process *)
-  let extend var term (prefix, subst) =
-    let subst = extend ~scope env subst var term in
-    (Binding.({var; term})::prefix, subst)
-  in
-  let rec helper x y acc =
-    (* log "unify '%s' and '%s'" (Term.show x) (Term.show y); *)
-    let open Term in
-    fold2 x y ~init:acc
-      ~fvar:(fun ((_, subst) as acc) x y ->
-        match walk env subst x, walk env subst y with
-        | Var x, Var y      ->
-          if Var.equal x y then acc else extend x (Term.repr y) acc
-        | Var x, Value y    -> extend x y acc
-        | Value x, Var y    -> extend y x acc
-        | Value x, Value y  -> helper x y acc
-      )
-      ~fval:(fun acc x y ->
-          if x = y then acc else raise Unification_failed
-      )
-      ~fk:(fun ((_, subst) as acc) l v y ->
-          if subsume && (l = Term.R)
-          then raise Unification_failed
-          else match walk env subst v with
-          | Var v    -> extend v y acc
-          | Value x  -> helper x y acc
-      )
-  in
   try
     let x, y = Term.(repr x, repr y) in
-    Some (helper x y ([], subst))
+    Some (unify_helper subsume (ext ~scope ~env add_delta) env x y (empty_delta, subst))
+  with Term.Different_shape _ | Unification_failed | Occurs_check -> None
+
+let unify ?(subsume=false) ?(scope=Term.Var.non_local_scope) =
+  unify_gen ~subsume ~scope List.cons []
+
+let unify_map env subst map : (Obj.t Term.VarMap.t * t) option =
+  let add_delta {Binding.var; term} m = Term.VarMap.add var term m in
+  let subsume = false in
+  let scope = Term.Var.non_local_scope in
+  try
+    Stdlib.Option.some @@
+    Term.VarMap.fold (fun var term acc ->
+      unify_helper subsume (ext  ~scope ~env add_delta) env (Obj.magic var) term acc
+    ) map (Term.VarMap.empty, subst)
   with Term.Different_shape _ | Unification_failed | Occurs_check -> None
 
 let apply env subst x = Obj.magic @@
   map env subst (Term.repr x)
     ~fvar:(fun v -> Term.repr v)
     ~fval:(fun x -> Term.repr x)
 
-let unify_map env subst map =
-  let vars, terms =
-    Term.VarMap.fold (fun v term acc -> (v :: fst acc, term :: snd acc)) map ([],[])
-  in
-  (* log "var   = %s" (Term.show (Obj.magic (apply env subst vars))); *)
-  (* log "terms = %s" (Term.show (Obj.magic (apply env subst terms))); *)
-  unify env subst (Obj.magic vars) (Obj.magic terms)
-
-
 let freevars env subst x =
   Env.freevars env @@ apply env subst x
 

diff --git a/src/core/Subst.mli b/src/core/Subst.mli
@@ -67,7 +67,7 @@ val freevars : Env.t -> t -> 'a -> Term.VarSet.t
  *)
 val unify : ?subsume:bool -> ?scope:Term.Var.scope -> Env.t -> t -> 'a -> 'a -> (Binding.t list * t) option
 
-val unify_map: Env.t -> t -> Term.t Term.VarMap.t -> (Binding.t list * t) option
+val unify_map: Env.t -> t -> Term.t Term.VarMap.t -> (Obj.t Term.VarMap.t * t) option
 
 val merge_disjoint : Env.t -> t -> t -> t