Merge PR #6506: Fast rel lookup

7 files changed, 82 insertions, 52 deletions

diff --git a/kernel/cClosure.ml b/kernel/cClosure.ml
index 11616da7b3..b1181157e1 100644
--- a/kernel/cClosure.ml
+++ b/kernel/cClosure.ml
@@ -258,7 +258,7 @@ type 'a infos_cache = {
   i_repr : 'a infos -> constr -> 'a;
   i_env : env;
   i_sigma : existential -> constr option;
-  i_rels : constr option array;
+  i_rels : (Context.Rel.Declaration.t * Pre_env.lazy_val) Range.t;
   i_tab : 'a KeyTable.t }
 
 and 'a infos = {
@@ -282,13 +282,16 @@ let ref_value_cache ({i_cache = cache} as infos)  ref =
     let body =
       match ref with
 	| RelKey n ->
-            let len = Array.length cache.i_rels in
-            let i = n - 1 in
-            let () = if i < 0 || len <= i then raise Not_found in
-            begin match Array.unsafe_get cache.i_rels i with
-            | None -> raise Not_found
-            | Some t -> lift n t
-            end
+          let open Context.Rel.Declaration in
+          let i = n - 1 in
+          let (d, _) =
+            try Range.get cache.i_rels i
+            with Invalid_argument _ -> raise Not_found
+          in
+          begin match d with
+          | LocalAssum _ -> raise Not_found
+          | LocalDef (_, t, _) -> lift n t
+          end
 	| VarKey id -> assoc_defined id cache.i_env
 	| ConstKey cst -> constant_value_in cache.i_env cst
     in
@@ -303,26 +306,13 @@ let ref_value_cache ({i_cache = cache} as infos)  ref =
 let evar_value cache ev =
   cache.i_sigma ev
 
-let defined_rels flags env =
-(*  if red_local_const (snd flags) then*)
-  let ctx = rel_context env in
-  let len = List.length ctx in
-  let ans = Array.make len None in
-  let open Context.Rel.Declaration in
-  let iter i = function
-    | LocalAssum _ -> ()
-    | LocalDef (_,b,_) -> Array.unsafe_set ans i (Some b)
-  in
-  let () = List.iteri iter ctx in
-  ans
-(*  else (0,[])*)
-
 let create mk_cl flgs env evars =
+  let open Pre_env in
   let cache =
     { i_repr = mk_cl;
       i_env = env;
       i_sigma = evars;
-      i_rels = defined_rels flgs env;
+      i_rels = (Environ.pre_env env).env_rel_context.env_rel_map;
       i_tab = KeyTable.create 17 }
   in { i_flags = flgs; i_cache = cache }
 
diff --git a/kernel/csymtable.ml b/kernel/csymtable.ml
index 2ffe36fcf7..712c66f11d 100644
--- a/kernel/csymtable.ml
+++ b/kernel/csymtable.ml
@@ -198,7 +198,7 @@ and slot_for_fv env fv =
       let rv = Pre_env.lookup_rel_val i env in
       begin match force_lazy_val rv with
       | None ->
-	 env.env_rel_context |> Context.Rel.lookup i |> RelDecl.get_value |> fill_fv_cache rv i val_of_rel env_of_rel
+        env |> Pre_env.lookup_rel i |> RelDecl.get_value |> fill_fv_cache rv i val_of_rel env_of_rel
       | Some (v, _) -> v
       end
   | FVuniv_var idu ->
diff --git a/kernel/environ.ml b/kernel/environ.ml
index 3c86129fed..93dc2f9a72 100644
--- a/kernel/environ.ml
+++ b/kernel/environ.ml
@@ -60,18 +60,17 @@ let deactivated_guard env = not (typing_flags env).check_guarded
 let universes env = env.env_stratification.env_universes
 let named_context env = env.env_named_context.env_named_ctx
 let named_context_val env = env.env_named_context
-let rel_context env = env.env_rel_context
+let rel_context env = env.env_rel_context.env_rel_ctx
 let opaque_tables env = env.indirect_pterms
 let set_opaque_tables env indirect_pterms = { env with indirect_pterms }
 
 let empty_context env =
-  match env.env_rel_context, env.env_named_context.env_named_ctx with
+  match env.env_rel_context.env_rel_ctx, env.env_named_context.env_named_ctx with
   | [], [] -> true
   | _ -> false
 
 (* Rel context *)
-let lookup_rel n env =
-  Context.Rel.lookup n env.env_rel_context
+let lookup_rel = lookup_rel
 
 let evaluable_rel n env =
   is_local_def (lookup_rel n env)
@@ -88,13 +87,12 @@ let push_rec_types (lna,typarray,_) env =
 
 let fold_rel_context f env ~init =
   let rec fold_right env =
-    match env.env_rel_context with
-    | [] -> init
-    | rd::rc ->
+    match match_rel_context_val env.env_rel_context with
+    | None -> init
+    | Some (rd, _, rc) ->
 	let env =
 	  { env with
 	    env_rel_context = rc;
-	    env_rel_val = List.tl env.env_rel_val;
 	    env_nb_rel = env.env_nb_rel - 1 } in
 	f env rd (fold_right env)
   in fold_right env
@@ -144,16 +142,21 @@ let evaluable_named id env =
 let reset_with_named_context ctxt env =
   { env with
     env_named_context = ctxt;
-    env_rel_context = Context.Rel.empty;
-    env_rel_val = [];
+    env_rel_context = empty_rel_context_val;
     env_nb_rel = 0 }
 
 let reset_context = reset_with_named_context empty_named_context_val
 
 let pop_rel_context n env =
+  let rec skip n ctx =
+    if Int.equal n 0 then ctx
+    else match match_rel_context_val ctx with
+    | None -> invalid_arg "List.skipn"
+    | Some (_, _, ctx) -> skip (pred n) ctx
+  in
   let ctxt = env.env_rel_context in
   { env with
-    env_rel_context = List.skipn n ctxt;
+    env_rel_context = skip n ctxt;
     env_nb_rel = env.env_nb_rel - n }
 
 let fold_named_context f env ~init =
diff --git a/kernel/nativecode.ml b/kernel/nativecode.ml
index c558e9ed03..ffe19510a6 100644
--- a/kernel/nativecode.ml
+++ b/kernel/nativecode.ml
@@ -1830,7 +1830,7 @@ and apply_fv env sigma univ (fv_named,fv_rel) auxdefs ml =
   in
   let auxdefs = List.fold_right get_rel_val fv_rel auxdefs in
   let auxdefs = List.fold_right get_named_val fv_named auxdefs in
-  let lvl = Context.Rel.length env.env_rel_context in
+  let lvl = Context.Rel.length env.env_rel_context.env_rel_ctx in
   let fv_rel = List.map (fun (n,_) -> MLglobal (Grel (lvl-n))) fv_rel in
   let fv_named = List.map (fun (id,_) -> MLglobal (Gnamed id)) fv_named in
   let aux_name = fresh_lname Anonymous in
@@ -1838,8 +1838,8 @@ and apply_fv env sigma univ (fv_named,fv_rel) auxdefs ml =
 
 and compile_rel env sigma univ auxdefs n =
   let open Context.Rel.Declaration in
-  let decl = Context.Rel.lookup n env.env_rel_context in
-  let n = Context.Rel.length env.env_rel_context - n in
+  let decl = Pre_env.lookup_rel n env in
+  let n = List.length env.env_rel_context.env_rel_ctx - n in
   match decl with
   | LocalDef (_,t,_) ->
       let code = lambda_of_constr env sigma t in
diff --git a/kernel/nativelambda.ml b/kernel/nativelambda.ml
index 160a90dc2f..b333b0fb9c 100644
--- a/kernel/nativelambda.ml
+++ b/kernel/nativelambda.ml
@@ -639,7 +639,7 @@ let optimize lam =
 let lambda_of_constr env sigma c =
   set_global_env env;
   let env = Renv.make () in
-  let ids = List.rev_map RelDecl.get_name !global_env.env_rel_context in
+  let ids = List.rev_map RelDecl.get_name !global_env.env_rel_context.env_rel_ctx in
   Renv.push_rels env (Array.of_list ids);
   let lam = lambda_of_constr env sigma c in
 (*  if Flags.vm_draw_opt () then begin
diff --git a/kernel/pre_env.ml b/kernel/pre_env.ml
index 4ef89f8c0e..3ef15105ae 100644
--- a/kernel/pre_env.ml
+++ b/kernel/pre_env.ml
@@ -67,11 +67,15 @@ type named_context_val = {
   env_named_map : (Context.Named.Declaration.t * lazy_val) Id.Map.t;
 }
 
+type rel_context_val = {
+  env_rel_ctx : Context.Rel.t;
+  env_rel_map : (Context.Rel.Declaration.t * lazy_val) Range.t;
+}
+
 type env = {
   env_globals       : globals;           (* globals = constants + inductive types + modules + module-types *)
   env_named_context : named_context_val; (* section variables *)
-  env_rel_context   : Context.Rel.t;
-  env_rel_val       : lazy_val list;
+  env_rel_context   : rel_context_val;
   env_nb_rel        : int;
   env_stratification : stratification;
   env_typing_flags  : typing_flags;
@@ -84,6 +88,11 @@ let empty_named_context_val = {
   env_named_map = Id.Map.empty;
 }
 
+let empty_rel_context_val = {
+  env_rel_ctx = [];
+  env_rel_map = Range.empty;
+}
+
 let empty_env = {
   env_globals = {
     env_constants = Cmap_env.empty;
@@ -91,8 +100,7 @@ let empty_env = {
     env_modules = MPmap.empty;
     env_modtypes = MPmap.empty};
   env_named_context = empty_named_context_val;
-  env_rel_context = Context.Rel.empty;
-  env_rel_val = [];
+  env_rel_context = empty_rel_context_val;
   env_nb_rel = 0;
   env_stratification = {
     env_universes = UGraph.initial_universes;
@@ -106,21 +114,39 @@ let empty_env = {
 
 let nb_rel env = env.env_nb_rel
 
+let push_rel_context_val d ctx = {
+  env_rel_ctx = Context.Rel.add d ctx.env_rel_ctx;
+  env_rel_map = Range.cons (d, ref VKnone) ctx.env_rel_map;
+}
+
+let match_rel_context_val ctx = match ctx.env_rel_ctx with
+| [] -> None
+| decl :: rem ->
+  let (_, lval) = Range.hd ctx.env_rel_map in
+  let ctx = { env_rel_ctx = rem; env_rel_map = Range.tl ctx.env_rel_map } in
+  Some (decl, lval, ctx)
+
 let push_rel d env =
-  let rval = ref VKnone in
     { env with
-      env_rel_context = Context.Rel.add d env.env_rel_context;
-      env_rel_val = rval :: env.env_rel_val;
+      env_rel_context = push_rel_context_val d env.env_rel_context;
       env_nb_rel = env.env_nb_rel + 1 }
 
+let lookup_rel n env =
+  try fst (Range.get env.env_rel_context.env_rel_map (n - 1))
+  with Invalid_argument _ -> raise Not_found
+
 let lookup_rel_val n env =
-  try List.nth env.env_rel_val (n - 1)
-  with Failure _ -> raise Not_found
+  try snd (Range.get env.env_rel_context.env_rel_map (n - 1))
+  with Invalid_argument _ -> raise Not_found
+
+let rel_skipn n ctx = {
+  env_rel_ctx = Util.List.skipn n ctx.env_rel_ctx;
+  env_rel_map = Range.skipn n ctx.env_rel_map;
+}
 
 let env_of_rel n env =
   { env with
-    env_rel_context = Util.List.skipn n env.env_rel_context;
-    env_rel_val = Util.List.skipn n env.env_rel_val;
+    env_rel_context = rel_skipn n env.env_rel_context;
     env_nb_rel = env.env_nb_rel - n
   }
 
diff --git a/kernel/pre_env.mli b/kernel/pre_env.mli
index fef530c872..335ca1057f 100644
--- a/kernel/pre_env.mli
+++ b/kernel/pre_env.mli
@@ -45,11 +45,15 @@ type named_context_val = private {
   env_named_map : (Context.Named.Declaration.t * lazy_val) Id.Map.t;
 }
 
+type rel_context_val = private {
+  env_rel_ctx : Context.Rel.t;
+  env_rel_map : (Context.Rel.Declaration.t * lazy_val) Range.t;
+}
+
 type env = {
     env_globals       : globals;
     env_named_context : named_context_val;
-    env_rel_context   : Context.Rel.t;
-    env_rel_val       : lazy_val list;
+    env_rel_context   : rel_context_val;
     env_nb_rel        : int;
     env_stratification : stratification;
     env_typing_flags  : typing_flags;
@@ -63,8 +67,15 @@ val empty_env : env
 
 (** Rel context *)
 
+val empty_rel_context_val : rel_context_val
+val push_rel_context_val :
+  Context.Rel.Declaration.t -> rel_context_val -> rel_context_val
+val match_rel_context_val  :
+  rel_context_val -> (Context.Rel.Declaration.t * lazy_val * rel_context_val) option
+
 val nb_rel         : env -> int
 val push_rel       : Context.Rel.Declaration.t -> env -> env
+val lookup_rel     : int -> env -> Context.Rel.Declaration.t
 val lookup_rel_val : int -> env -> lazy_val
 val env_of_rel     : int -> env -> env