From 23f41fddb30e4e932ecc57f1ac2ccdc122171f72 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 31 Mar 2022 19:53:03 -0700
Subject: [PATCH] feat: basic tactic cache

TODO: move `IO.Ref` to command
---
 src/Init/Notation.lean          |  2 +
 src/Lean/Elab/Tactic.lean       |  1 +
 src/Lean/Elab/Tactic/Cache.lean | 66 +++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)
 create mode 100644 src/Lean/Elab/Tactic/Cache.lean

diff --git a/src/Init/Notation.lean b/src/Init/Notation.lean
index 1b9cb24617..a6e6f90671 100644
--- a/src/Init/Notation.lean
+++ b/src/Init/Notation.lean
@@ -461,6 +461,8 @@ macro "unhygienic " t:tacticSeq : tactic => `(set_option tactic.hygienic false i
 
 syntax (name := fail) "fail " (str)? : tactic
 
+syntax (name := checkpoint) "checkpoint " tacticSeq : tactic
+
 end Tactic
 
 namespace Attr
diff --git a/src/Lean/Elab/Tactic.lean b/src/Lean/Elab/Tactic.lean
index d59b8766a9..63a49bd0ed 100644
--- a/src/Lean/Elab/Tactic.lean
+++ b/src/Lean/Elab/Tactic.lean
@@ -19,3 +19,4 @@ import Lean.Elab.Tactic.Conv
 import Lean.Elab.Tactic.Delta
 import Lean.Elab.Tactic.Meta
 import Lean.Elab.Tactic.Unfold
+import Lean.Elab.Tactic.Cache
diff --git a/src/Lean/Elab/Tactic/Cache.lean b/src/Lean/Elab/Tactic/Cache.lean
new file mode 100644
index 0000000000..108788341f
--- /dev/null
+++ b/src/Lean/Elab/Tactic/Cache.lean
@@ -0,0 +1,66 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Elab.Tactic.Basic
+
+namespace Lean.Elab.Tactic
+open Meta
+
+private def equivMVarDecl (d₁ d₂ : MetavarDecl) : Bool :=
+  d₁.type == d₂.type &&
+  d₁.lctx.decls.size == d₂.lctx.decls.size &&
+  Id.run do
+     for i in [:d₁.lctx.decls.size] do
+       match d₁.lctx.decls[i], d₂.lctx.decls[i] with
+       | some localDecl₁, some localDecl₂ => if localDecl₁.type != localDecl₂.type then return false
+       | none, none => pure ()
+       | _, _ => return false
+     return true
+
+structure Snapshot where
+  core   : Core.State
+  meta   : Meta.State
+  term   : Term.State
+  tactic : Tactic.State
+  stx    : Syntax
+  deriving Inhabited
+
+abbrev Cache := Std.HashMap MVarId Snapshot
+
+def Cache.find? (c : Cache) (mvarId : MVarId) (stx : Syntax) : TacticM (Option Snapshot) := do
+  let some s := Std.HashMap.find? c mvarId | return none
+  let mvarDecl ← getMVarDecl mvarId
+  let some mvarDeclOld := s.meta.mctx.findDecl? mvarId | return none
+  if equivMVarDecl mvarDecl mvarDeclOld && stx == s.stx then
+    return some s
+  else
+    return none
+
+/-
+  Very naive cache for tactic state. TODO: move it to a LSP snapshot
+-/
+builtin_initialize cacheRef : IO.Ref Cache ← IO.mkRef {}
+
+@[builtinTactic checkpoint] def evalCheckpoint : Tactic := fun stx =>
+  focus do
+    let mvarId ← getMainGoal
+    match (← (← cacheRef.get).find? mvarId stx[1]) with
+    | some s =>
+      set s.core
+      set s.meta
+      set s.term
+      set s.tactic
+    | none =>
+      evalTactic stx[1]
+      let s := {
+        stx    := stx[1]
+        core   := (← getThe Core.State)
+        meta   := (← getThe Meta.State)
+        term   := (← getThe Term.State)
+        tactic := (← get)
+      }
+      cacheRef.modify (·.insert mvarId s)
+
+end Lean.Elab.Tactic