Index keys instead of simply global references.

1 files changed, 102 insertions, 22 deletions

diff --git a/library/keys.ml b/library/keys.ml
index 5c295a385..c661e67fb 100644
--- a/library/keys.ml
+++ b/library/keys.ml
@@ -2,45 +2,125 @@ open Globnames
 open Term
 open Libobject
 
+
+type key = 
+  | KGlob of global_reference
+  | KLam 
+  | KLet
+  | KProd
+  | KSort
+  | KEvar
+  | KCase 
+  | KFix 
+  | KCoFix
+  | KRel 
+  | KMeta
+
+module KeyOrdered = struct
+  type t = key
+
+  let hash gr =
+    match gr with
+    | KGlob gr -> 10 + RefOrdered.hash gr
+    | KLam -> 0
+    | KLet -> 1
+    | KProd -> 2
+    | KSort -> 3
+    | KEvar -> 4
+    | KCase -> 5
+    | KFix -> 6
+    | KCoFix -> 7
+    | KRel -> 8
+    | KMeta -> 9
+
+  let compare gr1 gr2 =
+    match gr1, gr2 with
+    | KGlob gr1, KGlob gr2 -> RefOrdered.compare gr1 gr2
+    | _, KGlob _ -> -1
+    | KGlob _, _ -> 1
+    | k, k' -> Int.compare (hash k) (hash k')
+    
+  let equal k1 k2 =
+    match k1, k2 with
+    | KGlob gr1, KGlob gr2 -> RefOrdered.equal gr1 gr2
+    | _, KGlob _ -> false
+    | KGlob _, _ -> false
+    | k, k' -> k == k'
+end
+
+module Keymap = HMap.Make(KeyOrdered)
+module Keyset = Keymap.Set
+
 (* Union-find structure for references to be considered equivalent *)
 
-module RefUF = Unionfind.Make(Refset)(Refmap)
+module KeyUF = Unionfind.Make(Keyset)(Keymap)
 
-let keys = (* Summary.ref ( *)RefUF.create ()(* ) ~name:"Keys_decl" *)
+let keys = (* Summary.ref ( *)KeyUF.create ()(* ) ~name:"Keys_decl" *)
 
-let add_key k v = RefUF.union k v keys
+let add_keys k v = KeyUF.union k v keys
 
-let equiv_keys k k' = RefUF.find k keys == RefUF.find k' keys
+let equiv_keys k k' =
+  k == k' || KeyUF.find k keys == KeyUF.find k' keys
 
 (** Registration of keys as an object *)
 
-let load_key _ (_,(ref,ref')) =
-  add_key ref ref'
+let load_keys _ (_,(ref,ref')) =
+  add_keys ref ref'
+
+let cache_keys o =
+  load_keys 1 o
 
-let cache_key o =
-  load_key 1 o
+let subst_key subst k = 
+  match k with
+  | KGlob gr -> KGlob (subst_global_reference subst gr)
+  | _ -> k
 
-let subst_key (subst,(ref,ref')) =
-  (subst_global_reference subst ref,
-   subst_global_reference subst ref')
+let subst_keys (subst,(k,k')) =
+  (subst_key subst k, subst_key subst k')
 
-let discharge_key (_,refs) =
-  match refs with 
-  | VarRef _, _ | _, VarRef _ -> None
-  | ref, ref' -> Some (pop_global_reference ref, pop_global_reference ref')
+let discharge_key = function
+  | KGlob (VarRef _) -> None
+  | KGlob g -> Some (KGlob (pop_global_reference g))
+  | x -> Some x
 
-let rebuild_key (ref,ref') = (ref, ref')
+let discharge_keys (_,(k,k')) =
+  match discharge_key k, discharge_key k' with 
+  | Some x, Some y -> Some (x, y)
+  | _ -> None
 
-type key_obj = global_reference * global_reference
+let rebuild_keys (ref,ref') = (ref, ref')
+
+type key_obj = key * key
 
 let inKeys : key_obj -> obj =
   declare_object {(default_object "KEYS") with
-    cache_function = cache_key;
-    load_function = load_key;
-    subst_function = subst_key;
+    cache_function = cache_keys;
+    load_function = load_keys;
+    subst_function = subst_keys;
     classify_function = (fun x -> Substitute x);
-    discharge_function = discharge_key;
-    rebuild_function = rebuild_key }
+    discharge_function = discharge_keys;
+    rebuild_function = rebuild_keys }
 
 let declare_keys ref ref' =
   Lib.add_anonymous_leaf (inKeys (ref,ref'))
+
+let rec constr_key c =
+  let open Globnames in 
+  match kind_of_term c with
+  | Const (c, _) -> KGlob (ConstRef c)
+  | Ind (i, u) -> KGlob (IndRef i)
+  | Construct (c,u) -> KGlob (ConstructRef c)
+  | Var id -> KGlob (VarRef id)
+  | App (f, _) -> constr_key f
+  | Proj (p, _) -> KGlob (ConstRef p)
+  | Cast (p, _, _) -> constr_key p
+  | Lambda _ -> KLam 
+  | Prod _ -> KProd
+  | Case _ -> KCase
+  | Fix _ -> KFix
+  | CoFix _ -> KCoFix
+  | Rel _ -> KRel
+  | Meta _ -> KMeta
+  | Evar _ -> KEvar
+  | Sort _ -> KSort 
+  | LetIn _ -> KLet