From 14bd9b319274ef386598d0a0859fc5a7103207f3 Mon Sep 17 00:00:00 2001
From: Damien Doligez <damien.doligez@inria.fr>
Date: Wed, 5 Feb 2025 00:38:27 +0100
Subject: [PATCH] Memoize level computation to avoid exponential behavior

Signed-off-by: Damien Doligez <damien.doligez@inria.fr>
---
 src/expr.mli           |  5 +++--
 src/expr/e_anon.ml     |  3 ++-
 src/expr/e_levels.ml   | 49 ++++++++++++++++++++++++++++++++----------
 src/expr/e_levels.mli  |  3 ++-
 src/expr/e_tla_norm.ml |  2 ++
 src/util/deque.ml      | 17 +++++++++++++++
 src/util/deque.mli     |  2 ++
 7 files changed, 66 insertions(+), 15 deletions(-)

diff --git a/src/expr.mli b/src/expr.mli
index ba2be4f12..aba40140b 100644
--- a/src/expr.mli
+++ b/src/expr.mli
@@ -875,8 +875,6 @@ module Levels: sig
     val kind_to_level:
         kind -> int
 
-    class virtual ['s] level_computation:
-        ['s] Visit.map
     class virtual ['s] _rm_expr_level:
         ['s] Visit.map
 
@@ -884,6 +882,9 @@ module Levels: sig
         ctx -> expr -> expr
     val rm_expr_level:
         ctx -> expr -> expr
+
+    val newcache:
+        expr -> expr
 end
 
 
diff --git a/src/expr/e_anon.ml b/src/expr/e_anon.ml
index e3e24f6c7..8c8f74fec 100644
--- a/src/expr/e_anon.ml
+++ b/src/expr/e_anon.ml
@@ -304,7 +304,8 @@ end
 
 class anon = object (self : 'self)
   inherit anon_sg as super
-  method expr scx e = Elab.desugar self#expr super#expr scx e
+  method expr scx e =
+    E_levels.newcache (Elab.desugar self#expr super#expr scx e)
 end
 
 
diff --git a/src/expr/e_levels.ml b/src/expr/e_levels.ml
index 97dd69aea..dd5921e15 100644
--- a/src/expr/e_levels.ml
+++ b/src/expr/e_levels.ml
@@ -57,6 +57,12 @@ type level_info =  (* Stores information for computing the
 
 let exprlevel : level_info pfuncs =
     Property.make "Expr.Levels.exprlevel"
+type elcache =
+  | ELCache_full of hyp Deque.dq * expr_ * level_info
+  | ELCache_empty
+let exprlevel_cache : elcache ref pfuncs =
+    Property.make "Expr.Levels.exprlevel_cache"
+
 let assert_is_level (level: int) =
     assert ((level >= 0) && (level <= 3))
     (* The value 4 signifies that the notion of level is undefined for
@@ -124,6 +130,23 @@ let kind_to_level (kind: kind): int =
     | Temporal -> 3
     | Unknown -> assert false
 
+let has_cached_level cx e =
+  match Property.get e exprlevel_cache with
+  | exception Not_found -> false
+  | {contents = ELCache_full (cx2, ecore, l)}
+    when ecore == e.core && Deque.equal (==) cx2 cx -> true
+  | _ -> false
+
+let get_cached_level e =
+  match Property.get e exprlevel_cache with
+  | {contents = ELCache_full (_, _, l)} -> l
+  | _ -> assert false
+  | exception Not_found -> assert false
+
+let set_cached_level e cx l =
+  match Property.get e exprlevel_cache with
+  | r -> r := ELCache_full (cx, e.core, l)
+  | exception Not_found -> ()
 
 (*
 The level property is stored in the syntax tree nodes.
@@ -142,8 +165,15 @@ class virtual ['s] level_computation = object (self : 'self)
     method expr ((_, cx) as scx) e =
         if (has_level e) then
             e
+        else if has_cached_level cx e then
+            assign e exprlevel (get_cached_level e)
         else
         begin
+        let assign e tag l =
+          assert (tag == exprlevel);
+          set_cached_level e cx l;
+          assign e tag l
+        in
         let max_args_level scx args = (
             let es_ = List.map (self#expr scx) args in
             let es_levels = List.map get_level es_ in
@@ -402,7 +432,7 @@ class virtual ['s] level_computation = object (self : 'self)
             assign e exprlevel level_info
         | Bang (e, sels) -> assert false  (* subexpression references are
                 expanded before expanding `ENABLED` and `\cdot`. *)
-        | Lambda (vs, e) ->
+        | Lambda (vs, e2) ->
             (* Note: 2nd-order operators (signified via `shp` below) are
             handled by the declarations in `e_scx` below and the call to
             the method `self#hyp` in the pattern case `Apply (op, es)` below.
@@ -415,7 +445,7 @@ class virtual ['s] level_computation = object (self : 'self)
                         @@ v)
                     vs) in
             (* expression `e` *)
-            let e_ = self#expr e_scx e in
+            let e_ = self#expr e_scx e2 in
             let e_level = get_level e_ in
             (*
             if e_level > 2 then
@@ -962,7 +992,7 @@ class virtual ['s] level_computation = object (self : 'self)
             *)
 
     method hyp scx h =
-        let level_info = begin match h.core with
+        let level_info h = begin match h.core with
             (* declared variable *)
             | Flex _ -> make_level_info 1
             (* declared operator of any arity *)
@@ -1003,18 +1033,14 @@ class virtual ['s] level_computation = object (self : 'self)
                 assert false  (* not implemented *)
             (* defined operator of any arity *)
             | Defn (df, _, _, _) ->
-                let df_ = self#defn scx df in
-                get_level_info df_
+                get_level_info df
             | Fact (e, _, _) ->
-                let e_ = self#expr scx e in
-                begin
-                if (has_level e_) then
-                    get_level_info e_
+                if (has_level e) then
+                    get_level_info e
                 else
                     (* For example, the notion of expression level is
                     undefined for theorems. *)
                     make_undefined_level_info
-                end
             end in
         (* Hidden definitions need to be visited to
         have level information available for computing the
@@ -1041,11 +1067,12 @@ class virtual ['s] level_computation = object (self : 'self)
                 let h = Fact (self#expr scx e, vis, tm) @@ h in
                 (adj scx h, h)
             in
-        let e_ = assign e_ exprlevel level_info in
+        let e_ = assign e_ exprlevel (level_info e_) in
         (scx_, e_)
 
 end
 
+let newcache e = Property.assign e exprlevel_cache (ref ELCache_empty)
 
 let compute_level cx expr =
     let visitor = object
diff --git a/src/expr/e_levels.mli b/src/expr/e_levels.mli
index 026b4d6bf..22d234a60 100644
--- a/src/expr/e_levels.mli
+++ b/src/expr/e_levels.mli
@@ -22,8 +22,9 @@ val assert_has_correct_level: expr -> unit
 val kind_to_level: kind -> int
 
 
-class virtual ['s] level_computation : ['s] E_visit.map
 class virtual ['s] _rm_expr_level : ['s] E_visit.map_visible_hyp
 
 val compute_level: ctx -> expr -> expr
 val rm_expr_level: ctx -> expr -> expr
+
+val newcache: expr -> expr
diff --git a/src/expr/e_tla_norm.ml b/src/expr/e_tla_norm.ml
index 10b574795..abeebfb30 100644
--- a/src/expr/e_tla_norm.ml
+++ b/src/expr/e_tla_norm.ml
@@ -10,6 +10,8 @@ open Property
 open E_t
 module B = Builtin
 
+let ( @@ ) core e = E_levels.newcache (core @@ e)
+
 let rec tuple_flatten e = match e.core with
   | Tuple es ->
       List.concat (List.map tuple_flatten es)
diff --git a/src/util/deque.ml b/src/util/deque.ml
index d05d269c2..4a16282d3 100644
--- a/src/util/deque.ml
+++ b/src/util/deque.ml
@@ -206,3 +206,20 @@ let alter ?(backwards=false) q n alt_fn =
         end
   in
   spin 0 q
+
+(* Compare deques `q1` and `q2`, using `cmp` to compare their elements. *)
+let equal cmp q1 q2 =
+  let rec diff l1 l2 =
+    match l1, l2 with
+    | x1 :: ll1, x2 :: ll2 when cmp x1 x2 -> diff ll1 ll2
+    | [], _ -> Some l2
+    | _, [] -> Some l1
+    | _ -> None
+  in
+  q1.flen + q1.rlen = q2.flen + q2.rlen && begin
+    let f = diff q1.front q2.front in
+    let r = diff q1.rear q2.rear in
+    match f, r with
+    | Some lf, Some lr -> List.for_all2 cmp lf (List.rev lr)
+    | _, _ -> false
+  end
diff --git a/src/util/deque.mli b/src/util/deque.mli
index 643fc6e01..0fff87198 100644
--- a/src/util/deque.mli
+++ b/src/util/deque.mli
@@ -33,3 +33,5 @@ val prepend_list : 'a list -> 'a dq -> 'a dq
 
 val of_list : 'a list -> 'a dq
 val to_list : 'a dq -> 'a list
+
+val equal : ('a -> 'a -> bool) -> 'a dq -> 'a dq -> bool