Merge PR #6506: Fast rel lookup

10 files changed, 211 insertions, 52 deletions

diff --git a/clib/clib.mllib b/clib/clib.mllib
index bd5ddb152..0b5d9826f 100644
--- a/clib/clib.mllib
+++ b/clib/clib.mllib
@@ -12,6 +12,7 @@ CString
 CStack
 
 Int
+Range
 HMap
 Bigint
 
diff --git a/clib/range.ml b/clib/range.ml
new file mode 100644
index 000000000..86a078633
--- /dev/null
+++ b/clib/range.ml
@@ -0,0 +1,91 @@
+(************************************************************************)
+(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
+(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2016     *)
+(*   \VV/  **************************************************************)
+(*    //   *      This file is distributed under the terms of the       *)
+(*         *       GNU Lesser General Public License Version 2.1        *)
+(************************************************************************)
+
+type 'a tree =
+| Leaf of 'a
+| Node of 'a * 'a tree * 'a tree
+
+type 'a t = Nil | Cons of int * 'a tree * 'a t
+
+let oob () = invalid_arg "index out of bounds"
+
+let empty = Nil
+
+let cons x l = match l with
+| Cons (h1, t1, Cons (h2, t2, rem)) ->
+  if Int.equal h1 h2 then Cons (1 + h1 + h2, Node (x, t1, t2), rem)
+  else Cons (1, Leaf x, l)
+| _ -> Cons (1, Leaf x, l)
+
+let is_empty = function
+| Nil -> true
+| _ -> false
+
+let rec tree_get h t i = match t with
+| Leaf x ->
+  if i = 0 then x else oob ()
+| Node (x, t1, t2) ->
+  if i = 0 then x
+  else
+    let h = h / 2 in
+    if i <= h then tree_get h t1 (i - 1) else tree_get h t2 (i - h - 1)
+
+let rec get l i = match l with
+| Nil -> oob ()
+| Cons (h, t, rem) ->
+  if i < h then tree_get h t i else get rem (i - h)
+
+let length l =
+  let rec length accu = function
+  | Nil -> accu
+  | Cons (h, _, l) -> length (h + accu) l
+  in
+  length 0 l
+
+let rec tree_map f = function
+| Leaf x -> Leaf (f x)
+| Node (x, t1, t2) -> Node (f x, tree_map f t1, tree_map f t2)
+
+let rec map f = function
+| Nil -> Nil
+| Cons (h, t, l) -> Cons (h, tree_map f t, map f l)
+
+let rec tree_fold_left f accu = function
+| Leaf x -> f accu x
+| Node (x, t1, t2) ->
+  tree_fold_left f (tree_fold_left f (f accu x) t1) t2
+
+let rec fold_left f accu = function
+| Nil -> accu
+| Cons (_, t, l) -> fold_left f (tree_fold_left f accu t) l
+
+let rec tree_fold_right f t accu = match t with
+| Leaf x -> f x accu
+| Node (x, t1, t2) ->
+  f x (tree_fold_right f t1 (tree_fold_right f t2 accu))
+
+let rec fold_right f l accu = match l with
+| Nil -> accu
+| Cons (_, t, l) -> tree_fold_right f t (fold_right f l accu)
+
+let hd = function
+| Nil -> failwith "hd"
+| Cons (_, Leaf x, _) -> x
+| Cons (_, Node (x, _, _), _) -> x
+
+let tl = function
+| Nil -> failwith "tl"
+| Cons (_, Leaf _, l) -> l
+| Cons (h, Node (_, t1, t2), l) ->
+  let h = h / 2 in
+  Cons (h, t1, Cons (h, t2, l))
+
+let rec skipn n l =
+  if n = 0 then l
+  else if is_empty l then failwith "List.skipn"
+  else skipn (pred n) (tl l)
diff --git a/clib/range.mli b/clib/range.mli
new file mode 100644
index 000000000..ae7684ffa
--- /dev/null
+++ b/clib/range.mli
@@ -0,0 +1,37 @@
+(************************************************************************)
+(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
+(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2016     *)
+(*   \VV/  **************************************************************)
+(*    //   *      This file is distributed under the terms of the       *)
+(*         *       GNU Lesser General Public License Version 2.1        *)
+(************************************************************************)
+
+(** Skewed lists
+
+    This is a purely functional datastructure isomorphic to usual lists, except
+    that it features a O(log n) lookup while preserving the O(1) cons operation.
+
+*)
+
+(** {5 Constructors} *)
+
+type +'a t
+
+val empty : 'a t
+val cons : 'a -> 'a t -> 'a t
+
+(** {5 List operations} *)
+
+val is_empty : 'a t -> bool
+val length : 'a t -> int
+val map : ('a -> 'b) -> 'a t -> 'b t
+val fold_left : ('a -> 'b -> 'a) -> 'a -> 'b t -> 'a
+val fold_right : ('a -> 'b -> 'b) -> 'a t -> 'b -> 'b
+val hd : 'a t -> 'a
+val tl : 'a t -> 'a t
+
+val skipn : int -> 'a t -> 'a t
+
+(** {5 Indexing operations} *)
+
+val get : 'a t -> int -> 'a
diff --git a/kernel/cClosure.ml b/kernel/cClosure.ml
index 11616da7b..b1181157e 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 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 3c86129fe..93dc2f9a7 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 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 4ef89f8c0..3ef15105a 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 fef530c87..335ca1057 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