Simplify vars_of_global usage

4 files changed, 21 insertions, 18 deletions

diff --git a/kernel/constr.ml b/kernel/constr.ml
index 7ffde3107b..b490aa5092 100644
--- a/kernel/constr.ml
+++ b/kernel/constr.ml
@@ -407,6 +407,12 @@ let destCoFix c = match kind c with
   | CoFix cofix -> cofix
   | _ -> raise DestKO
 
+let destRef c = let open GlobRef in match kind c with
+  | Var x -> VarRef x, Univ.Instance.empty
+  | Const (c,u) -> ConstRef c, u
+  | Ind (ind,u) -> IndRef ind, u
+  | Construct (c,u) -> ConstructRef c, u
+  | _ -> raise DestKO
 
 (******************************************************************)
 (* Flattening and unflattening of embedded applications and casts *)
diff --git a/kernel/constr.mli b/kernel/constr.mli
index 12819ac39d..c012f04260 100644
--- a/kernel/constr.mli
+++ b/kernel/constr.mli
@@ -343,6 +343,8 @@ val destFix : constr -> fixpoint
 
 val destCoFix : constr -> cofixpoint
 
+val destRef : constr -> GlobRef.t Univ.puniverses
+
 (** {6 Equality} *)
 
 (** [equal a b] is true if [a] equals [b] modulo alpha, casts,
diff --git a/kernel/environ.ml b/kernel/environ.ml
index 757c8773b7..3b7e3ae544 100644
--- a/kernel/environ.ml
+++ b/kernel/environ.ml
@@ -560,26 +560,22 @@ let universes_of_global env r =
 
 (* Returns the list of global variables in a term *)
 
-let vars_of_global env constr =
-  match kind constr with
-      Var id -> Id.Set.singleton id
-    | Const (kn, _) -> lookup_constant_variables kn env
-    | Ind (ind, _) -> lookup_inductive_variables ind env
-    | Construct (cstr, _) -> lookup_constructor_variables cstr env
-    (** FIXME: is Proj missing? *)
-    | _ -> raise Not_found
+let vars_of_global env gr =
+  let open GlobRef in
+  match gr with
+  | VarRef id -> Id.Set.singleton id
+  | ConstRef kn -> lookup_constant_variables kn env
+  | IndRef ind -> lookup_inductive_variables ind env
+  | ConstructRef cstr -> lookup_constructor_variables cstr env
 
 let global_vars_set env constr =
   let rec filtrec acc c =
-    let acc =
-      match kind c with
-      | Var _ | Const _ | Ind _ | Construct _ ->
-	  Id.Set.union (vars_of_global env c) acc
-      | _ ->
-	  acc in
-    Constr.fold filtrec acc c
+    match destRef c with
+    | gr, _ ->
+      Id.Set.union (vars_of_global env gr) acc
+    | exception DestKO -> Constr.fold filtrec acc c
   in
-    filtrec Id.Set.empty constr
+  filtrec Id.Set.empty constr
 
 
 (* [keep_hyps env ids] keeps the part of the section context of [env] which
diff --git a/kernel/environ.mli b/kernel/environ.mli
index 6e313faab0..3d653bcfe0 100644
--- a/kernel/environ.mli
+++ b/kernel/environ.mli
@@ -276,8 +276,7 @@ val universes_of_global : env -> GlobRef.t -> AUContext.t
 
 val global_vars_set : env -> constr -> Id.Set.t
 
-(** the constr must be a global reference *)
-val vars_of_global : env -> constr -> Id.Set.t
+val vars_of_global : env -> GlobRef.t -> Id.Set.t
 
 (** closure of the input id set w.r.t. dependency *)
 val really_needed : env -> Id.Set.t -> Id.Set.t