From 1435fbeda5e211c493307649c1dd71b9fe8536ac Mon Sep 17 00:00:00 2001
From: David Thrane Christiansen <david@lean-fro.org>
Date: Fri, 17 May 2024 10:14:54 +0200
Subject: [PATCH] feat: proof states after "by" (#28)

Fixes #25
---
 InternalTests.lean                            |  7 +++++--
 .../SubVerso/Highlighting/Code.lean           | 19 ++++++++++++++-----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/InternalTests.lean b/InternalTests.lean
index d500c42..674cee2 100644
--- a/InternalTests.lean
+++ b/InternalTests.lean
@@ -56,7 +56,7 @@ theorem test (n : Nat) : n * 1 = n := by
 
 
 
--- We don't have #guard_msgs in all supported Lean versions, so here's a dumb replacement:
+-- We don't have #guard_msgs in all supported Lean versions, so here's a low-tech replacement:
 
 open Lean Elab Command in
 elab "#evalString" s:str e:term : command => do
@@ -68,10 +68,13 @@ elab "#evalString" s:str e:term : command => do
     let [msg] := msgs'.toList
       | throwError "Too many messages"
     if (← msg.toString) != s.getString then
-      throwErrorAt e "Expected {s.getString}, got {← msg.toString}"
+      throwErrorAt e "Expected {String.quote s.getString}, got {String.quote (← msg.toString)}"
   finally
     MonadMessageState.setMessages msgs
 
+#evalString "[[\"n * 1 = n\"]]\n"
+  (proofEx.highlighted[0].proofStates.data.filter (·.fst == "by") |>.map (·.snd.data.map (·.conclusion)))
+
 #evalString "[[some `zero], [some `succ], [none], [some `succ.succ], [none]]\n"
   (proofEx.highlighted[0].proofStates.data.filter (·.fst == "=>") |>.map (·.snd.data.map (·.name)))
 
diff --git a/src/highlighting/SubVerso/Highlighting/Code.lean b/src/highlighting/SubVerso/Highlighting/Code.lean
index 0edfc1b..1fb5986 100644
--- a/src/highlighting/SubVerso/Highlighting/Code.lean
+++ b/src/highlighting/SubVerso/Highlighting/Code.lean
@@ -608,13 +608,22 @@ partial def findTactics
   let endPosition := text.toPosition endPos
 
   -- Blacklisted tactics. TODO: make into an extensible table.
-  -- For now, the `by` keyword itself is blacklisted, as is its leading token
-  if stx.getKind ∈ [``Lean.Parser.Term.byTactic, ``Lean.Parser.Term.byTactic'] ||
-     stx matches .atom _ "by" then
-    return none
-  -- `;` is blacklisted as well - no need to highlight states identically
+  -- `;` is blacklisted  - no need to highlight states identically
   if stx matches .atom _ ";" then return none
 
+  -- Place the initial proof state after the `by` token
+  if stx matches .atom _ "by" then
+    for t in trees do
+      for i in infoIncludingSyntax t stx |>.reverse do
+        if not i.2.isOriginal then continue
+        if let (_, .ofTacticInfo tacticInfo) := i then
+          match tacticInfo.stx with
+          | `(Lean.Parser.Term.byTactic| by%$tk $tactics)
+          | `(Lean.Parser.Term.byTactic'| by%$tk $tactics) =>
+            if tk == stx then
+              return (← findTactics' ids trees tactics startPos endPos endPosition (before := true))
+          | _ => continue
+
   -- Special handling for =>: show the _before state_
   if stx matches .atom _ "=>" then
     for t in trees do