Fast environment lookup for rels.

We take advantage of the range structure to get a O(log n) retrieval of values
bound to a rel in an environment.

6 files changed, 69 insertions, 29 deletions

diff --git a/kernel/csymtable.ml b/kernel/csymtable.ml
index 2ffe36fcf..712c66f11 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 1afab453a..9d76768d9 100644
--- a/kernel/environ.ml
+++ b/kernel/environ.ml
@@ -58,18 +58,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)
@@ -86,13 +85,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
@@ -142,16 +140,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 c558e9ed0..ffe19510a 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 160a90dc2..b333b0fb9 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 c5254b453..d685e269f 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;
@@ -85,6 +89,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;
@@ -92,8 +101,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;
@@ -108,21 +116,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 054ae1743..563acb0a2 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;
@@ -64,8 +68,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