import Lean
import Mathlib.Tactic

open Lean Meta Elab Command Tactic

Inspired by

• Lean.Elab.Tactic.focus
• Branch in lean4game, see also its doc
• `tada` on Zulip`
• Lean auto grader

elab "proofs " 
n: TSyntax `num
n:
num: Parser.Parser
num : 
tactic: Parser.Category
tactic => do
  
withMainContext: {α : Type} → TacticM α → TacticM α
withMainContext do
    let 
mvarId: MVarId
mvarId :: 
mvarIds: List MVarId
mvarIds ← 
getUnsolvedGoals: TacticM (List MVarId)
getUnsolvedGoals | 
throwNoGoalsToBeSolved: {α : Type} → TacticM α
throwNoGoalsToBeSolved
    
setGoals: List MVarId → TacticM Unit
setGoals $ 
List.replicate: {α : Type} → ℕ → α → List α
List.replicate (
n: TSyntax `num
n.
1: {ks : SyntaxNodeKinds} → TSyntax ks → Syntax
1.
toNat: Syntax → ℕ
toNat) 
mvarId: MVarId
mvarId ++ 
mvarIds: List MVarId
mvarIds

elab (name := 
proof: ParserDescr
proof) "proof" 
_n: Option (TSyntax `num)
_n:(
num: Parser.Parser
num)? 
_desc: Option (TSyntax Name.anonymous)
_desc:
interpolatedStr: Parser.Parser → Parser.Parser
interpolatedStr(
term: Parser.Category
term)? ":" 
t: TSyntax `Lean.Parser.Tactic.tacticSeq
t:
tacticSeq: Parser.Parser
tacticSeq : 
tactic: Parser.Category
tactic => do
  
withMainContext: {α : Type} → TacticM α → TacticM α
withMainContext do
    let 
mvarId: MVarId
mvarId :: 
mvarIds: List MVarId
mvarIds ← 
getUnsolvedGoals: TacticM (List MVarId)
getUnsolvedGoals | 
throwNoGoalsToBeSolved: {α : Type} → TacticM α
throwNoGoalsToBeSolved

    -- focus on the first goal
    
setGoals: List MVarId → TacticM Unit
setGoals [
mvarId: MVarId
mvarId]
    -- save the non-proven state for recover after the proof
    let 
b: Tactic.SavedState
b ← 
saveState: {s : outParam Type} → {m : Type → Type} → [self : MonadBacktrack s m] → m s
saveState

    -- run the proof
    let 
a: Unit
a ← 
evalTactic: Syntax → TacticM Unit
evalTactic 
t: TSyntax `Lean.Parser.Tactic.tacticSeq
t

    -- detect sorry in proof
    if let 
some: {α : Type ?u.673} → α → Option α
some 
proof: Expr
proof := 
(← getMCtx): MetavarContext
(← 
getMCtx: {m : Type → Type} → [self : MonadMCtx m] → m MetavarContext
getMCtx
(← getMCtx): MetavarContext
).
eAssignment: MetavarContext → PersistentHashMap MVarId Expr
eAssignment.
find?: {α β : Type} → {x : BEq α} → {x_1 : Hashable α} → PersistentHashMap α β → α → Option β
find? 
mvarId: MVarId
mvarId then
      if 
proof: Expr
proof.
isSorry: Expr → Bool
isSorry then
        
logWarning: {m : Type → Type} →
  [inst : Monad m] → [inst : MonadLog m] → [inst : AddMessageContext m] → [inst : MonadOptions m] → MessageData → m Unit
logWarning
logWarning "proof uses 'sorry'": TacticM Unit
 
"proof uses 'sorry'": String
"proof uses 'sorry'"

    -- detect unsolved goals
    let 
gs: List MVarId
gs ← 
getUnsolvedGoals: TacticM (List MVarId)
getUnsolvedGoals
    if !
gs: List MVarId
gs.
isEmpty: {α : Type} → List α → Bool
isEmpty then
      
Term.reportUnsolvedGoals: List MVarId → MetaM Unit
Term.reportUnsolvedGoals
Term.reportUnsolvedGoals gs: TacticM Unit
 
gs: List MVarId
gs

    -- if it's not the last alternative proof, restore the non-proven state and restore the remaining goals
    -- otherwise the remaining goal will be proven automatically
    if !
mvarIds: List MVarId
mvarIds.
isEmpty: {α : Type} → List α → Bool
isEmpty then
      let 
msgs: MessageLog
msgs ← 
Core.getMessageLog: CoreM MessageLog
Core.getMessageLog
      
b: Tactic.SavedState
b.
restore: Tactic.SavedState → optParam Bool false → TacticM Unit
restore
      
setGoals: List MVarId → TacticM Unit
setGoals 
mvarIds: List MVarId
mvarIds
      
Core.setMessageLog: MessageLog → CoreM Unit
Core.setMessageLog
Core.setMessageLog msgs: TacticM Unit
 
msgs: MessageLog
msgs
    -- else use the last proof as the proof
    -- TODO this doesn't guarantee that the whole theorem is proven if any of the proofs are valid
    -- it also doesn't guarantee that the whole theorem is proven only if all of the proofs are valid

    
pure: {f : Type → Type} → [self : Pure f] → {α : Type} → α → f α
pure 
a: Unit
a

elab (name := 
QED: ParserDescr
QED) "QED" : 
tactic: Parser.Category
tactic => do
  let 
gs: List MVarId
gs ← 
getUnsolvedGoals: TacticM (List MVarId)
getUnsolvedGoals
  if 
gs: List MVarId
gs.
isEmpty: {α : Type} → List α → Bool
isEmpty then
    
logInfo: {m : Type → Type} →
  [inst : Monad m] → [inst : MonadLog m] → [inst : AddMessageContext m] → [inst : MonadOptions m] → MessageData → m Unit
logInfo 
" 🎉 Goals accomplished": String
" 🎉 Goals accomplished"
  else
    
logError: {m : Type → Type} →
  [inst : Monad m] → [inst : MonadLog m] → [inst : AddMessageContext m] → [inst : MonadOptions m] → MessageData → m Unit
logError 
"☝️ Unresolved goals": String
"☝️ Unresolved goals"

elab (name := 
WIP: ParserDescr
WIP) "WIP" : 
tactic: Parser.Category
tactic => do
  let 
gs: List MVarId
gs ← 
getUnsolvedGoals: TacticM (List MVarId)
getUnsolvedGoals
  if 
gs: List MVarId
gs.
isEmpty: {α : Type} → List α → Bool
isEmpty then
    
logError: {m : Type → Type} →
  [inst : Monad m] → [inst : MonadLog m] → [inst : AddMessageContext m] → [inst : MonadOptions m] → MessageData → m Unit
logError 
" 🎉 Goals accomplished, mark it with QED": String
" 🎉 Goals accomplished, mark it with QED"
  else
    
logWarning: {m : Type → Type} →
  [inst : Monad m] → [inst : MonadLog m] → [inst : AddMessageContext m] → [inst : MonadOptions m] → MessageData → m Unit
logWarning 
"☝️ Unresolved goals": String
"☝️ Unresolved goals"

example: ∀ (a b c : ℕ), a + b * c = c * b + a
example (
a: ℕ
a 
b: ℕ
b 
c: ℕ
c : 
ℕ: Type
ℕ) : 
a: ℕ
a + 
b: ℕ
b * 
c: ℕ
c = 
c: ℕ
c * 
b: ℕ
b + 
a: ℕ
a := by
Goals accomplished! 🐙
  proofs 2
a, b, c: ℕ
a + b * c = c * b + a
a, b, c: ℕ
a + b * c = c * b + a
  proof 1:
a, b, c: ℕ
a + b * c = c * b + a
a, b, c: ℕ
a + b * c = c * b + a
    rw [
a, b, c: ℕ
a + b * c = c * b + a
Nat.add_comm: ∀ (n m : ℕ), n + m = m + n
Nat.add_comm,
a, b, c: ℕ
b * c + a = c * b + a @
Nat.add_right_cancel_iff: ∀ {m k n : ℕ}, m + n = k + n ↔ m = k
Nat.add_right_cancel_iff,
a, b, c: ℕ
b * c = c * b 
Nat.mul_comm: ∀ (n m : ℕ), n * m = m * n
Nat.mul_comm
a, b, c: ℕ
c * b = c * b]
Goals accomplished! 🐙
  proof 2:
a, b, c: ℕ
a + b * c = c * b + a
    ring
Goals accomplished! 🐙

-- to see the colorful effect, comment out #guard_msgs in

warning: proof uses 'sorry'

error: unsolved goals ⊢ 2 + 1 = 3

#guard_msgs
Error: ❌ Docstring on `#guard_msgs` does not match generated message:

warning: proof uses 'sorry'
---
error: unsolved goals
⊢ 2 + 1 = 3 in
theorem 
ex: 1 + 2 = 3
ex : 
1: ℕ
1 + 
2: ℕ
2 = 
3: ℕ
3 := by
Goals accomplished! 🐙
  proofs 5
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
  proof "trivial":
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
    rfl
Goals accomplished! 🐙
  proof "automatic":
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
    simp
Goals accomplished! 🐙
  proof "cheating":
Warning: proof uses 'sorry'
1 + 2 = 3
1 + 2 = 3
1 + 2 = 3
    sorry
Warning: proof uses 'sorry'
Goals accomplished! 🐙
  proof "WIP":
Error: unsolved goals
⊢ 2 + 1 = 3
1 + 2 = 3
1 + 2 = 3
    rewrite [
Error: unsolved goals
⊢ 2 + 1 = 3
1 + 2 = 3
Nat.add_comm: ∀ (n m : ℕ), n + m = m + n
Nat.add_comm
Error: unsolved goals
⊢ 2 + 1 = 3
2 + 1 = 3]
Error: unsolved goals
⊢ 2 + 1 = 3
2 + 1 = 3
  proof "verbose":
1 + 2 = 3
    have 
h: 1 + 2 = 3
h : 
1: ℕ
1 + 
2: ℕ
2 = 
3: ℕ
3 :=
1 + 2 = 3 by
Goals accomplished! 🐙 rfl
Goals accomplished! 🐙
    exact 
h: 1 + 2 = 3
h
Goals accomplished! 🐙