More efficient implementation for substitutions.

We use a variant of skewed lists enriched over a monoid, whose purpose is
to count the number of lifts added to the substitution. This makes addition
O(1) and lookup O(log n).

1 files changed, 128 insertions, 91 deletions

diff --git a/kernel/esubst.ml b/kernel/esubst.ml
index 558ff96486..97acd698e8 100644
--- a/kernel/esubst.ml
+++ b/kernel/esubst.ml
@@ -60,94 +60,131 @@ let rec is_lift_id = function
 (*  Substitutions    *)
 (*********************)
 
-(* (bounded) explicit substitutions of type 'a *)
-type 'a subs =
-  | ESID of int            (* ESID(n)    = %n END   bounded identity *)
-  | CONS of 'a array * 'a subs
-                           (* CONS([|t1..tn|],S)  =
-                              (S.t1...tn)    parallel substitution
-                              beware of the order *)
-  | SHIFT of int * 'a subs (* SHIFT(n,S) = (^n o S) terms in S are relocated *)
-                           (*                        with n vars *)
-  | LIFT of int * 'a subs  (* LIFT(n,S) = (%n S) stands for ((^n o S).n...1) *)
-
-(* operations of subs: collapses constructors when possible.
- * Needn't be recursive if we always use these functions
- *)
-
-let subs_id i = ESID i
-
-let subs_cons(x,s) = if Int.equal (Array.length x) 0 then s else CONS(x,s)
-
-let subs_liftn n = function
-  | ESID p        -> ESID (p+n) (* bounded identity lifted extends by p *)
-  | LIFT (p,lenv) -> LIFT (p+n, lenv)
-  | lenv          -> LIFT (n,lenv)
-
-let subs_lift a = subs_liftn 1 a
-let subs_liftn n a = if Int.equal n 0 then a else subs_liftn n a
-
-let subs_shft = function
-  | (0, s)            -> s
-  | (n, SHIFT (k,s1)) -> SHIFT (k+n, s1)
-  | (n, s)            -> SHIFT (n,s)
-let subs_shft s = if Int.equal (fst s) 0 then snd s else subs_shft s
-
-(* Tests whether a substitution is equal to the identity *)
-let rec is_subs_id = function
-    ESID _     -> true
-  | LIFT(_,s)  -> is_subs_id s
-  | SHIFT(0,s) -> is_subs_id s
-  | CONS(x,s)  -> Int.equal (Array.length x) 0 && is_subs_id s
-  | _          -> false
-
-(* Expands de Bruijn k in the explicit substitution subs
- * lams accumulates de shifts to perform when retrieving the i-th value
- * the rules used are the following:
- *
- *    [id]k       --> k
- *    [S.t]1      --> t
- *    [S.t]k      --> [S](k-1)  if k > 1
- *    [^n o S] k  --> [^n]([S]k)
- *    [(%n S)] k  --> k         if k <= n
- *    [(%n S)] k  --> [^n]([S](k-n))
- *
- * the result is (Inr (k+lams,p)) when the variable is just relocated
- * where p is None if the variable points inside subs and Some(k) if the
- * variable points k bindings beyond subs.
- *)
-let rec exp_rel lams k subs =
-  match subs with
-    | CONS (def,_) when k <= Array.length def
-                           -> Inl(lams,def.(Array.length def - k))
-    | CONS (v,l)           -> exp_rel lams (k - Array.length v) l
-    | LIFT (n,_) when k<=n -> Inr(lams+k,None)
-    | LIFT (n,l)           -> exp_rel (n+lams) (k-n) l
-    | SHIFT (n,s)          -> exp_rel (n+lams) k s
-    | ESID n when k<=n     -> Inr(lams+k,None)
-    | ESID n               -> Inr(lams+k,Some (k-n))
-
-let expand_rel k subs = exp_rel 0 k subs
-
-let rec subs_map f = function
-| ESID _ as s -> s
-| CONS (x, s) -> CONS (Array.map f x, subs_map f s)
-| SHIFT (n, s) -> SHIFT (n, subs_map f s)
-| LIFT (n, s) -> LIFT (n, subs_map f s)
-
-let rec lift_subst mk_cl s1 s2 = match s1 with
-| ELID -> subs_map (fun c -> mk_cl ELID c) s2
-| ELSHFT(s, k) -> subs_shft(k, lift_subst mk_cl s s2)
-| ELLFT (k, s) ->
-  match s2 with
-  | CONS(x,s') ->
-      CONS(CArray.Fun1.map mk_cl s1 x, lift_subst mk_cl s1 s')
-  | ESID n -> lift_subst mk_cl s (ESID (n + k))
-  | SHIFT(k',s') ->
-      if k<k'
-      then subs_shft(k, lift_subst mk_cl s (subs_shft(k'-k, s')))
-      else subs_shft(k', lift_subst mk_cl (el_liftn (k-k') s) s')
-  | LIFT(k',s') ->
-      if k<k'
-      then subs_liftn k (lift_subst mk_cl s (subs_liftn (k'-k) s'))
-      else subs_liftn k' (lift_subst mk_cl (el_liftn (k-k') s) s')
+(* Variant of skewed lists enriched w.r.t. a monoid. See the Range module.
+
+  In addition to the indexed data, every node contains a monoid element, in our
+  case, integers. It corresponds to the number of partial lifts to apply when
+  reaching this subtree. The total lift is obtained by summing all the partial
+  lifts encountered in the tree traversal. For efficiency, we also cache the
+  sum of partial lifts of the whole subtree as the last argument of the [Node]
+  constructor. *)
+
+type mon = int
+let mul n m = n + m
+let one = 0
+
+type 'a or_var = Arg of 'a | Var of int
+
+type 'a tree =
+| Leaf of mon * 'a or_var
+| Node of mon * 'a or_var * 'a tree * 'a tree * mon
+
+type 'a subs = Nil of mon | Cons of int * 'a tree * 'a subs
+
+let eval = function
+| Leaf (w, _) -> w
+| Node (w1, _, _, _, w2) -> mul w1 w2
+
+let leaf x = Leaf (one, x)
+let node x t1 t2 = Node (one, x, t1, t2, mul (eval t1) (eval t2))
+
+let rec tree_get h w t i = match t with
+| Leaf (w', x) ->
+  let w = mul w w' in
+  if i = 0 then w, Inl x else assert false
+| Node (w', x, t1, t2, _) ->
+  let w = mul w w' in
+  if i = 0 then w, Inl x
+  else
+    let h = h / 2 in
+    if i <= h then tree_get h w t1 (i - 1)
+    else tree_get h (mul w (eval t1)) t2 (i - h - 1)
+
+let rec get w l i = match l with
+| Nil w' -> mul w w', Inr i
+| Cons (h, t, rem) ->
+  if i < h then tree_get h w t i else get (mul (eval t) w) rem (i - h)
+
+let get l i = get one l i
+
+let tree_write w = function
+| Leaf (w', x) -> Leaf (mul w w', x)
+| Node (w', x, t1, t2, wt) -> Node (mul w w', x, t1, t2, wt)
+
+let write w l = match l with
+| Nil w' -> Nil (mul w w')
+| Cons (h, t, rem) -> Cons (h, tree_write w t, rem)
+
+let empty = Nil one
+
+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 expand_rel n s =
+  let k, v = get s (n - 1) in
+  match v with
+  | Inl (Arg v) -> Inl (k, v)
+  | Inl (Var i) -> Inr (k + i, None)
+  | Inr i -> Inr (k + i + 1, Some (i + 1))
+
+let is_subs_id = function
+| Nil w -> Int.equal w 0
+| Cons (_, _, _) -> false
+
+let subs_cons (v, s) =
+  Array.fold_left (fun accu x -> cons (Arg x) accu) s v
+
+let rec push_vars i s =
+  if Int.equal i 0 then s
+  else push_vars (pred i) (cons (Var i) s)
+
+let subs_liftn n s =
+  if Int.equal n 0 then s
+  else
+    let s = write n s in
+    push_vars n s
+
+let subs_lift s =
+  cons (Var 1) (write 1 s)
+
+let subs_id n =
+  if Int.equal n 0 then empty
+  else push_vars n empty
+
+let subs_shft (n, s) = write n s
+
+let rec pop n i e =
+  if Int.equal n 0 then i, e
+  else match e with
+  | ELID -> i, e
+  | ELLFT (k, e) ->
+    if k <= n then pop (n - k) i e
+    else i, ELLFT (k - n, e)
+  | ELSHFT (e, k) -> pop n (i + k) e
+
+let apply mk e = function
+| Var i -> Var (reloc_rel i e)
+| Arg v -> Arg (mk e v)
+
+let rec tree_map mk e = function
+| Leaf (w, x) ->
+  let (n, e) = pop w 0 e in
+  Leaf (w + n, apply mk e x), e
+| Node (w, x, t1, t2, _) ->
+  let (n, e) = pop w 0 e in
+  let x = apply mk e x in
+  let t1, e = tree_map mk e t1 in
+  let t2, e = tree_map mk e t2 in
+  Node (w + n, x, t1, t2, mul (eval t1) (eval t2)), e
+
+let rec lift_subst mk e s = match s with
+| Nil w ->
+  let (n, _) = pop w 0 e in
+  Nil (w + n)
+| Cons (h, t, rem) ->
+  let t, e = tree_map mk e t in
+  let rem = lift_subst mk e rem in
+  Cons (h, t, rem)