From 3add9553823fa2465235730f2fe21b42eb724eef Mon Sep 17 00:00:00 2001
From: Gabriel Ebner <gebner@gebner.org>
Date: Fri, 16 Sep 2022 17:12:48 +0200
Subject: [PATCH] fix: preserve separators in evalSepByIndentTactic

---
 src/Lean/Elab/Tactic/BuiltinTactic.lean       | 89 +++++++++++++------
 tests/lean/interactive/plainGoal.lean         |  2 +
 .../interactive/plainGoal.lean.expected.out   | 56 +++++++-----
 3 files changed, 95 insertions(+), 52 deletions(-)

diff --git a/src/Lean/Elab/Tactic/BuiltinTactic.lean b/src/Lean/Elab/Tactic/BuiltinTactic.lean
index 9c5ce5aa70..9c2b567828 100644
--- a/src/Lean/Elab/Tactic/BuiltinTactic.lean
+++ b/src/Lean/Elab/Tactic/BuiltinTactic.lean
@@ -35,42 +35,75 @@ open Parser.Tactic
 @[builtinTactic paren] def evalParen : Tactic := fun stx =>
   evalTactic stx[1]
 
-def isCheckpointableTactic (arg : Syntax.Tactic) : TacticM Bool := do
+def isCheckpointableTactic (arg : Syntax) : TacticM Bool := do
   -- TODO: make it parametric
-  let kind := arg.1.getKind
+  let kind := arg.getKind
   return kind == ``Lean.Parser.Tactic.save
 
-partial def addCheckpoints (args : Array Syntax.Tactic) : TacticM (Array Syntax.Tactic) := do
-  -- if (← readThe Term.Context).tacticCache? |>.isSome then
-  if (← args.anyM fun arg => isCheckpointableTactic arg) then
-    let argsNew ← go 0 #[] #[]
-    return argsNew
-  return args
-where
-  push (acc : Array Syntax.Tactic) (result : Array Syntax.Tactic) : Array Syntax.Tactic :=
-    if acc.isEmpty then
-      result
-    else
-      result.push <| Unhygienic.run do withRef acc.back do
-        `(tactic| checkpoint $[$acc]*)
+/--
+Takes a `sepByIndent tactic "; "`, and inserts `checkpoint` blocks for `save` tactics.
 
-  go (i : Nat) (acc : Array Syntax.Tactic) (result : Array Syntax.Tactic) : TacticM (Array Syntax.Tactic) := do
-    if h : i < args.size then
-      let arg := args.get ⟨i, h⟩
-      if (← isCheckpointableTactic arg) then
-        -- `argNew` is `arg` as singleton sequence. The idea is to make sure it does not satisfy `isCheckpointableTactic` anymore
-        let argNew := ⟨mkNullNode #[arg]⟩ -- HACK
-        let acc := acc.push argNew
-        go (i+1) #[] (push acc result)
-      else
-        go (i+1) (acc.push arg) result
+Input:
+```
+  a
+  b
+  save
+  c
+  d
+  save
+  e
+```
+
+Output:
+```
+  checkpoint
+    a
+    b
+    save
+  checkpoint
+    c
+    d
+    save
+  e
+```
+-/
+-- Note that we need to preserve the separators to show the right goals after semicolons.
+def addCheckpoints (stx : Syntax) : TacticM Syntax := do
+  -- if (← readThe Term.Context).tacticCache? |>.isSome then
+  if !(← stx.getSepArgs.anyM isCheckpointableTactic) then return stx
+  let mut currentCheckpointBlock := #[]
+  let mut output := #[]
+  for i in [:stx.getArgs.size / 2] do
+    let tac := stx[2*i]
+    let sep? := stx.getArgs[2*i+1]?
+    if (← isCheckpointableTactic tac) then
+      let checkpoint : Syntax :=
+        mkNode ``checkpoint #[
+          mkAtomFrom tac "checkpoint",
+          mkNode ``tacticSeq #[
+            mkNode ``tacticSeq1Indented #[
+              -- HACK: null node is not a valid tactic, but prevents infinite loop
+              mkNullNode (currentCheckpointBlock.push tac)
+            ]
+          ]
+        ]
+      currentCheckpointBlock := #[]
+      output := output.push checkpoint
+      if let some sep := sep? then output := output.push sep
     else
-      return result ++ acc
+      currentCheckpointBlock := currentCheckpointBlock.push tac
+      if let some sep := sep? then currentCheckpointBlock := currentCheckpointBlock.push sep
+  output := output ++ currentCheckpointBlock
+  return stx.setArgs output
 
 /-- Evaluate `sepByIndent tactic "; " -/
 def evalSepByIndentTactic (stx : Syntax) : TacticM Unit := do
-  for seqElem in (← addCheckpoints (stx.getSepArgs.map (⟨·⟩))) do
-    evalTactic seqElem
+  let stx ← addCheckpoints stx
+  for arg in stx.getArgs, i in [:stx.getArgs.size] do
+    if i % 2 == 0 then
+      evalTactic arg
+    else
+      saveTacticInfoForToken arg
 
 @[builtinTactic tacticSeq1Indented] def evalTacticSeq1Indented : Tactic := fun stx =>
   evalSepByIndentTactic stx[0]
diff --git a/tests/lean/interactive/plainGoal.lean b/tests/lean/interactive/plainGoal.lean
index 7a9591cc90..dfa8945e00 100644
--- a/tests/lean/interactive/plainGoal.lean
+++ b/tests/lean/interactive/plainGoal.lean
@@ -21,6 +21,8 @@ example : 0 + n = n := by
 example : α → α := by
   intro a; apply a
        --^ $/lean/plainGoal
+        --^ $/lean/plainGoal
+         --^ $/lean/plainGoal
 
 example (h1 : n = m) (h2 : m = 0) : 0 = n := by
   rw [h1, h2]
diff --git a/tests/lean/interactive/plainGoal.lean.expected.out b/tests/lean/interactive/plainGoal.lean.expected.out
index fdf1b7ca19..f0862658db 100644
--- a/tests/lean/interactive/plainGoal.lean.expected.out
+++ b/tests/lean/interactive/plainGoal.lean.expected.out
@@ -37,64 +37,72 @@
 {"rendered": "```lean\nα : Sort ?u\na : α\n⊢ α\n```",
  "goals": ["α : Sort ?u\na : α\n⊢ α"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 25, "character": 3}}
+ "position": {"line": 21, "character": 10}}
+{"rendered": "```lean\nα : Sort ?u\na : α\n⊢ α\n```",
+ "goals": ["α : Sort ?u\na : α\n⊢ α"]}
+{"textDocument": {"uri": "file://plainGoal.lean"},
+ "position": {"line": 21, "character": 11}}
+{"rendered": "```lean\nα : Sort ?u\na : α\n⊢ α\n```",
+ "goals": ["α : Sort ?u\na : α\n⊢ α"]}
+{"textDocument": {"uri": "file://plainGoal.lean"},
+ "position": {"line": 27, "character": 3}}
 {"rendered": "```lean\nn m : Nat\nh1 : n = m\nh2 : m = 0\n⊢ 0 = n\n```",
  "goals": ["n m : Nat\nh1 : n = m\nh2 : m = 0\n⊢ 0 = n"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 25, "character": 9}}
+ "position": {"line": 27, "character": 9}}
 {"rendered": "```lean\nn m : Nat\nh1 : n = m\nh2 : m = 0\n⊢ 0 = m\n```",
  "goals": ["n m : Nat\nh1 : n = m\nh2 : m = 0\n⊢ 0 = m"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 25, "character": 13}}
+ "position": {"line": 27, "character": 13}}
 {"rendered": "no goals", "goals": []}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 32, "character": 3}}
+ "position": {"line": 34, "character": 3}}
 {"rendered": "```lean\ncase zero\n⊢ 0 + Nat.zero = Nat.zero\n```",
  "goals": ["case zero\n⊢ 0 + Nat.zero = Nat.zero"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 38, "character": 3}}
+ "position": {"line": 40, "character": 3}}
 {"rendered":
  "```lean\ncase zero\n⊢ 0 + Nat.zero = Nat.zero\n```\n---\n```lean\ncase succ\nn✝ : Nat\n: 0 + n✝ = n✝\n⊢ 0 + Nat.succ n✝ = Nat.succ n✝\n```",
  "goals":
  ["case zero\n⊢ 0 + Nat.zero = Nat.zero",
   "case succ\nn✝ : Nat\n: 0 + n✝ = n✝\n⊢ 0 + Nat.succ n✝ = Nat.succ n✝"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 42, "character": 3}}
+ "position": {"line": 44, "character": 3}}
 {"rendered":
  "```lean\ncase zero\n⊢ 0 + Nat.zero = Nat.zero\n```\n---\n```lean\ncase succ\nn✝ : Nat\n⊢ 0 + Nat.succ n✝ = Nat.succ n✝\n```",
  "goals":
  ["case zero\n⊢ 0 + Nat.zero = Nat.zero",
   "case succ\nn✝ : Nat\n⊢ 0 + Nat.succ n✝ = Nat.succ n✝"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 46, "character": 3}}
+ "position": {"line": 48, "character": 3}}
 {"rendered": "```lean\na b : Nat\n⊢ a = b\n```",
  "goals": ["a b : Nat\n⊢ a = b"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 49, "character": 20}}
+ "position": {"line": 51, "character": 20}}
 {"rendered": "```lean\nα : Sort ?u\n⊢ α → α\n```",
  "goals": ["α : Sort ?u\n⊢ α → α"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 53, "character": 3}}
+ "position": {"line": 55, "character": 3}}
 {"rendered":
  "```lean\nα : Sort ?u\np : α → Prop\na b : α\ninst✝ : DecidablePred p\nh : ∀ {p : α → Prop} [inst : DecidablePred p], p a → p b\n⊢ p a\n```",
  "goals":
  ["α : Sort ?u\np : α → Prop\na b : α\ninst✝ : DecidablePred p\nh : ∀ {p : α → Prop} [inst : DecidablePred p], p a → p b\n⊢ p a"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 59, "character": 3}}
+ "position": {"line": 61, "character": 3}}
 {"rendered": "```lean\ncase left\n⊢ True\n```", "goals": ["case left\n⊢ True"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 61, "character": 3}}
+ "position": {"line": 63, "character": 3}}
 {"rendered": "```lean\ncase right\n⊢ False\n```",
  "goals": ["case right\n⊢ False"]}
-{"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 66, "character": 3}}
-{"rendered": "```lean\ncase left\n⊢ True\n```", "goals": ["case left\n⊢ True"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
  "position": {"line": 68, "character": 3}}
+{"rendered": "```lean\ncase left\n⊢ True\n```", "goals": ["case left\n⊢ True"]}
+{"textDocument": {"uri": "file://plainGoal.lean"},
+ "position": {"line": 70, "character": 3}}
 {"rendered": "```lean\ncase right\n⊢ False\n```",
  "goals": ["case right\n⊢ False"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 76, "character": 29}}
+ "position": {"line": 78, "character": 29}}
 {"rendered":
  "```lean\nt a n✝ : Nat\n: t * (a + n✝) = t * a + t * n✝\n⊢ t * (a + n✝) + t = t * a + t * Nat.succ n✝\n```\n---\n```lean\nt a n✝ : Nat\n: t * (a + n✝) = t * a + t * n✝\n⊢ t * (a + n✝) + t = t * a + (t * n✝ + t)\n```\n---\n```lean\nt a n✝ : Nat\n: t * (a + n✝) = t * a + t * n✝\n⊢ t * (a + n✝) + t = t * a + (t * n✝ + t)\n```",
  "goals":
@@ -102,37 +110,37 @@
   "t a n✝ : Nat\n: t * (a + n✝) = t * a + t * n✝\n⊢ t * (a + n✝) + t = t * a + (t * n✝ + t)",
   "t a n✝ : Nat\n: t * (a + n✝) = t * a + t * n✝\n⊢ t * (a + n✝) + t = t * a + (t * n✝ + t)"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 80, "character": 53}}
+ "position": {"line": 82, "character": 53}}
 {"rendered":
  "```lean\ncase nil\nα : Type ?u\nbs cs : List α\n⊢ [] ++ bs ++ cs = [] ++ (bs ++ cs)\n```",
  "goals":
  ["case nil\nα : Type ?u\nbs cs : List α\n⊢ [] ++ bs ++ cs = [] ++ (bs ++ cs)"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 80, "character": 54}}
+ "position": {"line": 82, "character": 54}}
 {"rendered":
  "```lean\ncase nil\nα : Type ?u\nbs cs : List α\n⊢ [] ++ bs ++ cs = [] ++ (bs ++ cs)\n```",
  "goals":
  ["case nil\nα : Type ?u\nbs cs : List α\n⊢ [] ++ bs ++ cs = [] ++ (bs ++ cs)"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 84, "character": 38}}
+ "position": {"line": 86, "character": 38}}
 {"rendered": "no goals", "goals": []}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 87, "character": 39}}
+ "position": {"line": 89, "character": 39}}
 null
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 91, "character": 16}}
+ "position": {"line": 93, "character": 16}}
 {"rendered":
  "```lean\ncase left\n⊢ True\n```\n---\n```lean\ncase right\n⊢ False\n```",
  "goals": ["case left\n⊢ True", "case right\n⊢ False"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 95, "character": 8}}
+ "position": {"line": 97, "character": 8}}
 {"rendered": "```lean\n| True = True\n```", "goals": ["| True = True"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 97, "character": 4}}
+ "position": {"line": 99, "character": 4}}
 {"rendered": "```lean\n| True = True\n```", "goals": ["| True = True"]}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 99, "character": 2}}
+ "position": {"line": 101, "character": 2}}
 {"rendered": "no goals", "goals": []}
 {"textDocument": {"uri": "file://plainGoal.lean"},
- "position": {"line": 103, "character": 2}}
+ "position": {"line": 105, "character": 2}}
 {"rendered": "```lean\n⊢ False\n```", "goals": ["⊢ False"]}