hconstr v4: share terms with identical binders

kernel/typeops.ml

@@ -732,25 +732,6 @@ let self x = x.self
 
 let refcount x = x.refcount
 
-type local_env = {
-  (* unique identifiers for each binder crossed *)
-  rels : int Range.t;
-  (* global counter *)
-  cnt : int ref;
-}
-
-let empty_env () = {
-  rels = Range.empty;
-  cnt = ref 0;
-}
-
-let push_rel env =
-  incr env.cnt;
-  {
-    rels = Range.cons !(env.cnt) env.rels;
-    cnt = env.cnt;
-  }
-
 module Tbl = struct
   type key = t
 
@@ -798,6 +779,58 @@ module Tbl = struct
 
 end
 
+type local_env = {
+  (* unique identifiers for each binder crossed *)
+  rels : int Range.t;
+  (* global counter *)
+  cnt : int ref;
+  assum_uids : int Tbl.t;
+  (* the surrounding table is for the body, the inner table for the type *)
+  letin_uids : int Tbl.t Tbl.t;
+}
+
+let empty_env () = {
+  rels = Range.empty;
+  cnt = ref 0;
+  assum_uids = Tbl.create 47;
+  letin_uids = Tbl.create 47;
+}
+
+let push_unknown_rel env =
+  incr env.cnt;
+  { env with rels = Range.cons !(env.cnt) env.rels }
+
+let push_assum t env =
+  let uid = match Tbl.find_opt env.assum_uids t with
+  | Some uid -> uid
+  | None ->
+    incr env.cnt;
+    let uid = !(env.cnt) in
+    Tbl.add env.assum_uids t uid;
+    uid
+  in
+  { env with rels = Range.cons uid env.rels }
+
+let push_letin ~body ~typ env =
+  let uid = match Tbl.find_opt env.letin_uids body with
+    | Some tbl -> begin match Tbl.find_opt tbl typ with
+        | Some uid -> uid
+        | None ->
+          incr env.cnt;
+          let uid = !(env.cnt) in
+          Tbl.add tbl typ uid;
+          uid
+      end
+    | None ->
+      incr env.cnt;
+      let uid = !(env.cnt) in
+      let tbl = Tbl.create 3 in
+      Tbl.add tbl typ uid;
+      Tbl.add env.letin_uids body tbl;
+      uid
+  in
+  { env with rels = Range.cons uid env.rels }
+
 let hash_annot = hash_annot Name.hash
 
 let hash_array hashf a = Array.fold_left (fun hash x -> Hashset.Combine.combine hash (hashf x)) 0 a
@@ -912,18 +945,18 @@ and of_constr_aux tbl local_env c =
   | Prod (na,t,c) ->
     let na = hcons_annot na in
     let t = of_constr tbl local_env t in
-    let c = of_constr tbl (push_rel local_env) c in
+    let c = of_constr tbl (push_assum t local_env) c in
     Prod (na, t, c)
   | Lambda (na, t, c) ->
     let na = hcons_annot na in
     let t = of_constr tbl local_env t in
-    let c = of_constr tbl (push_rel local_env) c in
+    let c = of_constr tbl (push_assum t local_env) c in
     Lambda (na,t,c)
   | LetIn (na,b,t,c) ->
     let na = hcons_annot na in
     let b = of_constr tbl local_env b in
     let t = of_constr tbl local_env t in
-    let c = of_constr tbl (push_rel local_env) c in
+    let c = of_constr tbl (push_letin ~body:b ~typ:t local_env) c in
     LetIn (na,b,t,c)
   | App (h,args) ->
     let h = of_constr tbl local_env h in
@@ -946,7 +979,7 @@ and of_constr_aux tbl local_env c =
   | Case (ci,u,pms,(p,r),iv,c,bl) ->
     let of_ctx (bnd, c) =
       let () = hcons_inplace hcons_annot bnd in
-      let local_env = Array.fold_left (fun local_env _ -> push_rel local_env) local_env bnd in
+      let local_env = Array.fold_left (fun local_env _ -> push_unknown_rel local_env) local_env bnd in
       let c = of_constr tbl local_env c in
       bnd, c
     in
@@ -964,13 +997,13 @@ and of_constr_aux tbl local_env c =
   | Fix (ln,(lna,tl,bl)) ->
     let () = hcons_inplace hcons_annot lna in
     let tl = Array.map (of_constr tbl local_env) tl in
-    let body_env = Array.fold_left (fun env _ -> push_rel env) local_env tl in
+    let body_env = Array.fold_left (fun env t -> push_assum t env) local_env tl in
     let bl = Array.map (of_constr tbl body_env) bl in
     Fix (ln,(lna,tl,bl))
   | CoFix (ln,(lna,tl,bl)) ->
     let () = hcons_inplace hcons_annot lna in
     let tl = Array.map (of_constr tbl local_env) tl in
-    let body_env = Array.fold_left (fun env _ -> push_rel env) local_env tl in
+    let body_env = Array.fold_left (fun env t -> push_assum t env) local_env tl in
     let bl = Array.map (of_constr tbl body_env) bl in
     CoFix (ln,(lna,tl,bl))
   | Proj (p,r,c) ->
@@ -991,7 +1024,7 @@ let dbg = CDebug.create ~name:"hconstr" ()
 
 let of_constr env c =
   let local_env = empty_env () in
-  let local_env = iterate push_rel (Environ.nb_rel env) local_env in
+  let local_env = iterate push_unknown_rel (Environ.nb_rel env) local_env in
   let tbl = Tbl.create 57 in
   let c = of_constr tbl local_env c in
   dbg Pp.(fun () ->