Use weak ref to memoize Evarutil.is_ground_env

1 files changed, 5 insertions, 4 deletions

diff --git a/engine/evarutil.ml b/engine/evarutil.ml
index eea7e38f87..0db9f498b9 100644
--- a/engine/evarutil.ml
+++ b/engine/evarutil.ml
@@ -127,10 +127,11 @@ let is_ground_env evd env =
 (* Memoization is safe since evar_map and environ are applicative
    structures *)
 let memo f =
-  let m = ref None in
-  fun x y -> match !m with
-  | Some (x', y', r) when x == x' && y == y' -> r
-  | _ -> let r = f x y in m := Some (x, y, r); r
+  let module E = Ephemeron.K2 in
+  let m = E.create () in
+  fun x y -> match E.get_key1 m, E.get_key2 m with
+  | Some x', Some y' when x == x' && y == y' -> Option.get (E.get_data m)
+  | _ -> let r = f x y in E.set_key1 m x; E.set_key2 m y; E.set_data m r; r
 
 let is_ground_env = memo is_ground_env