From ebe2cf272e41f15579a92ad45feabe551911e5f5 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 8 Sep 2020 10:48:47 -0700
Subject: [PATCH] feat: add `MonadCacheT`

---
 src/Lean/Util/MonadCache.lean  | 14 +++++++++++++
 tests/lean/run/monadCache.lean | 38 ++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 tests/lean/run/monadCache.lean

diff --git a/src/Lean/Util/MonadCache.lean b/src/Lean/Util/MonadCache.lean
index a4f96c7888..8181901c6a 100644
--- a/src/Lean/Util/MonadCache.lean
+++ b/src/Lean/Util/MonadCache.lean
@@ -52,6 +52,20 @@ instance {α β : Type} {m : Type → Type} [HasBeq α] [Hashable α] [Monad m]
 
 end MonadHashMapCacheAdapter
 
+abbrev MonadCacheT {ω} (α β : Type) (m : Type → Type) [STWorld ω m] [HasBeq α] [Hashable α] := StateRefT (HashMap α β) m
+
+namespace MonadCacheT
+
+instance {ω} (α β : Type) (m : Type → Type) [STWorld ω m] [HasBeq α] [Hashable α] [MonadLiftT (ST ω) m] [Monad m]
+    : MonadHashMapCacheAdapter α β (MonadCacheT α β m) :=
+{ getCache    := get,
+  modifyCache := modify }
+
+@[inline] def run {ω α β m σ} [STWorld ω m] [HasBeq α] [Hashable α] [MonadLiftT (ST ω) m] [Monad m] (x : MonadCacheT α β m σ) : m σ :=
+x.run' Std.mkHashMap
+
+end MonadCacheT
+
 /-- Auxiliary structure for "adding" a `HashMap` to a state object. -/
 structure WithHashMapCache (α β σ : Type) [HasBeq α] [Hashable α] :=
 (state : σ)
diff --git a/tests/lean/run/monadCache.lean b/tests/lean/run/monadCache.lean
new file mode 100644
index 0000000000..7132edf098
--- /dev/null
+++ b/tests/lean/run/monadCache.lean
@@ -0,0 +1,38 @@
+import Lean
+open Lean
+
+def mkTower : Nat → Expr
+| 0   => mkConst `a
+| n+1 => mkApp2 (mkConst `f) (mkTower n) (mkTower n)
+
+partial def depth : Expr → MonadCacheT Expr Nat CoreM Nat
+| e =>
+  checkCache e fun e =>
+    match e with
+    | Expr.const c [] _ => pure 1
+    | Expr.app f a _    => do d₁ ← depth f; d₂ ← depth a; pure $ Nat.max d₁ d₂ + 1
+    | _                 => pure 0
+
+#eval (depth (mkTower 100)).run
+
+partial def visit : Expr → MonadCacheT Expr Expr CoreM Expr
+| e =>
+  checkCache e fun e =>
+    match e with
+    | Expr.const `a [] _ => pure $ mkConst `b
+    | Expr.app f a _     => e.updateApp! <$> visit f <*> visit a
+    | _                  => pure e
+
+#eval (visit (mkTower 4)).run
+
+#eval do e ← (visit (mkTower 100)).run; (depth e).run
+
+partial def visitNoCache : Expr → CoreM Expr
+| e =>
+  match e with
+  | Expr.const `a [] _ => pure $ mkConst `b
+  | Expr.app f a _     => e.updateApp! <$> visitNoCache f <*> visitNoCache a
+  | _                  => pure e
+
+-- The following is super slow
+-- #eval do e ← visitNoCache (mkTower 30); (depth e).run