feat: Add proof continue and root extraction

Pantograph/Goal.lean

@@ -31,14 +31,20 @@ protected def GoalState.create (expr: Expr): M GoalState := do
   --let expr ← instantiateMVars expr
   let goal := (← Meta.mkFreshExprMVar expr (kind := MetavarKind.synthetic) (userName := .anonymous))
   let savedStateMonad: Elab.Tactic.TacticM Elab.Tactic.SavedState := MonadBacktrack.saveState
-  let savedState ← savedStateMonad { elaborator := .anonymous } |>.run' { goals := [goal.mvarId!]}
+  let root := goal.mvarId!
+  let savedState ← savedStateMonad { elaborator := .anonymous } |>.run' { goals := [root]}
   return {
     savedState,
-    root := goal.mvarId!,
+    root,
-    newMVars := SSet.empty,
+    newMVars := SSet.insert .empty root,
   }
 protected def GoalState.goals (goalState: GoalState): List MVarId := goalState.savedState.tactic.goals
 
+private def GoalState.mctx (state: GoalState): MetavarContext :=
+  state.savedState.term.meta.meta.mctx
+private def GoalState.mvars (state: GoalState): SSet MVarId :=
+  state.mctx.decls.foldl (init := .empty) fun acc k _ => acc.insert k
+
 def executeTactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: Syntax) :
     M (Except (Array String) (Elab.Tactic.SavedState × List MVarId)):= do
   let tacticM (stx: Syntax):  Elab.Tactic.TacticM (Except (Array String) (Elab.Tactic.SavedState × List MVarId)) := do
@@ -93,13 +99,13 @@ protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String
     let prevMCtx := state.savedState.term.meta.meta.mctx
     -- Generate a list of mvarIds that exist in the parent state; Also test the
     -- assertion that the types have not changed on any mvars.
-    let newMVars := (← nextMCtx.decls.foldlM (fun acc mvarId mvarDecl => do
+    let newMVars ← nextMCtx.decls.foldlM (fun acc mvarId mvarDecl => do
       if let .some prevMVarDecl := prevMCtx.decls.find? mvarId then
         assert! prevMVarDecl.type == mvarDecl.type
         return acc
       else
-        return mvarId :: acc
+        return acc.insert mvarId
-      ) []).toSSet
+      ) SSet.empty
     let nextState: GoalState := {
       savedState := nextSavedState
       root := state.root,
@@ -115,38 +121,70 @@ protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String
       | .none => throwError s!"Parent mvar id does not exist {nextGoal.name}"
     return .success nextState goals.toArray
 
+/-- After finishing one branch of a proof (`graftee`), pick up from the point where the proof was left off (`target`) -/
+protected def GoalState.continue (target: GoalState) (graftee: GoalState): Except String GoalState :=
+  if target.root != graftee.root then
+    .error s!"Roots of two continued goal states do not match: {target.root.name} != {graftee.root.name}"
+  -- Ensure goals are not dangling
+  else if ¬ (target.goals.all (λ goal => graftee.mvars.contains goal)) then
+    .error s!"Some goals in target are not present in the graftee"
+  else
+    -- Set goals to the goals that have not been assigned yet, similar to the `focus` tactic.
+    let unassigned := target.goals.filter (λ goal =>
+      let mctx := graftee.mctx
+      ¬(mctx.eAssignment.contains goal || mctx.dAssignment.contains goal))
+    .ok {
+      savedState := {
+        term := graftee.savedState.term,
+        tactic := { goals := unassigned },
+      },
+      root := target.root,
+      newMVars := graftee.newMVars,
+    }
+
+protected def GoalState.rootExpr (goalState: GoalState): Option Expr :=
+  goalState.mctx.eAssignment.find? goalState.root |>.filter (λ e => ¬ e.hasMVar)
+
 -- Diagnostics functions
 
 /-- Print the metavariables in a readable format -/
-protected def GoalState.print (goalState: GoalState) (options: Protocol.GoalPrint := {}): Elab.TermElabM Unit := do
+protected def GoalState.print (goalState: GoalState) (options: Protocol.GoalPrint := {}): M Unit := do
   let savedState := goalState.savedState
   savedState.term.restore
   let goals := savedState.tactic.goals
   let mctx ← getMCtx
+  let root := goalState.root
+  -- Print the root
+  match mctx.decls.find? root with
+  | .some decl => printMVar ">" root decl
+  | .none => IO.println s!">{root.name}: ??"
   goals.forM (fun mvarId => do
-    let pref := "⊢"
+    if mvarId != root then
       match mctx.decls.find? mvarId with
-    | .some decl => printMVar pref mvarId decl
+      | .some decl => printMVar "⊢" mvarId decl
-    | .none => IO.println s!"{pref}{mvarId.name}: ??"
+      | .none => IO.println s!"⊢{mvarId.name}: ??"
   )
   let goals := goals.toSSet
   mctx.decls.forM (fun mvarId decl => do
-    if goals.contains mvarId then
+    if goals.contains mvarId || mvarId == root then
       pure ()
+    -- Always print the root goal
     else if mvarId == goalState.root then
       printMVar (pref := ">") mvarId decl
-    else if ¬(goalState.newMVars.contains mvarId) then
+    -- Print the remainig ones that users don't see in Lean
-      printMVar (pref := " ") mvarId decl
     else if options.printNonVisible then
-      printMVar (pref := "~") mvarId decl
+      let pref := if goalState.newMVars.contains mvarId then "~" else " "
+      printMVar pref mvarId decl
     else
-      IO.println s!" {mvarId.name}{userNameToString decl.userName}"
+      pure ()
+      --IO.println s!" {mvarId.name}{userNameToString decl.userName}"
   )
   where
     printMVar (pref: String) (mvarId: MVarId) (decl: MetavarDecl): Elab.TermElabM Unit := do
       if options.printContext then
         decl.lctx.fvarIdToDecl.forM printFVar
-      IO.println s!"{pref}{mvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type} {← serialize_expression_ast decl.type}"
+      let type_sexp ← serialize_expression_ast (← instantiateMVars decl.type) (sanitize := false)
+      IO.println s!"{pref}{mvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type} {type_sexp}"
       if options.printValue then
         if let Option.some value := (← getMCtx).eAssignment.find? mvarId then
           IO.println s!"  = {← Meta.ppExpr value}"

Pantograph/Protocol.lean

@@ -167,7 +167,8 @@ structure GoalDeleteResult where
 structure GoalPrint where
   printContext: Bool := true
   printValue: Bool := true
-  printNonVisible: Bool := true
+  printNewMVars: Bool := false
+  printNonVisible: Bool := false
 
 
 end Pantograph.Protocol

Pantograph/Serial.lean

@@ -45,9 +45,11 @@ def type_expr_to_bound (expr: Expr): MetaM Protocol.BoundExpression := do
       return (toString (← fvar.fvarId!.getUserName), toString (← Meta.ppExpr (← fvar.fvarId!.getType)))
     return { binders, target := toString (← Meta.ppExpr body) }
 
-private def name_to_ast: Lean.Name → String
+private def name_to_ast (name: Lean.Name) (sanitize: Bool := true): String := match name with
-  | .anonymous
+  | .anonymous => ":anon"
-  | .num _ _ => ":anon"
+  | .num n i => match sanitize with
+    | false => s!"{toString n} {i}"
+    | true => ":anon"
   | n@(.str _ _) => toString n
 
 private def level_depth: Level → Nat
@@ -100,71 +102,73 @@ def serialize_sort_level_ast (level: Level): String :=
 /--
  Completely serializes an expression tree. Json not used due to compactness
 -/
-def serialize_expression_ast (expr: Expr): MetaM String := do
+def serialize_expression_ast (expr: Expr) (sanitize: Bool := true): MetaM String := do
-  match expr with
+  return self expr
+  where
+  self (e: Expr): String :=
+    match e with
     | .bvar deBruijnIndex =>
       -- This is very common so the index alone is shown. Literals are handled below.
       -- The raw de Bruijn index should never appear in an unbound setting. In
       -- Lean these are handled using a `#` prefix.
-    return s!"{deBruijnIndex}"
+      s!"{deBruijnIndex}"
     | .fvar fvarId =>
-    let name := name_to_ast fvarId.name
+      let name := of_name fvarId.name
-    return s!"(:fv {name})"
+      s!"(:fv {name})"
     | .mvar mvarId =>
-    let name := name_to_ast mvarId.name
+      let name := of_name mvarId.name
-    return s!"(:mv {name})"
+      s!"(:mv {name})"
     | .sort level =>
       let level := serialize_sort_level_ast level
-    return s!"(:sort {level})"
+      s!"(:sort {level})"
     | .const declName _ =>
       -- The universe level of the const expression is elided since it should be
       -- inferrable from surrounding expression
-    return s!"(:c {declName})"
+      s!"(:c {declName})"
     | .app fn arg =>
-    let fn' ← serialize_expression_ast fn
+      let fn' := self fn
-    let arg' ← serialize_expression_ast arg
+      let arg' := self arg
-    return s!"({fn'} {arg'})"
+      s!"({fn'} {arg'})"
     | .lam binderName binderType body binderInfo =>
-    let binderName' := name_to_ast binderName
+      let binderName' := of_name binderName
-    let binderType' ← serialize_expression_ast binderType
+      let binderType' := self binderType
-    let body' ← serialize_expression_ast body
+      let body' := self body
       let binderInfo' := binder_info_to_ast binderInfo
-    return s!"(:lambda {binderName'} {binderType'} {body'}{binderInfo'})"
+      s!"(:lambda {binderName'} {binderType'} {body'}{binderInfo'})"
     | .forallE binderName binderType body binderInfo =>
-    let binderName' := name_to_ast binderName
+      let binderName' := of_name binderName
-    let binderType' ← serialize_expression_ast binderType
+      let binderType' := self binderType
-    let body' ← serialize_expression_ast body
+      let body' := self body
       let binderInfo' := binder_info_to_ast binderInfo
-    return s!"(:forall {binderName'} {binderType'} {body'}{binderInfo'})"
+      s!"(:forall {binderName'} {binderType'} {body'}{binderInfo'})"
     | .letE name type value body _ =>
       -- Dependent boolean flag diacarded
       let name' := name_to_ast name
-    let type' ← serialize_expression_ast type
+      let type' := self type
-    let value' ← serialize_expression_ast value
+      let value' := self value
-    let body' ← serialize_expression_ast body
+      let body' := self body
-    return s!"(:let {name'} {type'} {value'} {body'})"
+      s!"(:let {name'} {type'} {value'} {body'})"
     | .lit v =>
       -- To not burden the downstream parser who needs to handle this, the literal
       -- is wrapped in a :lit sexp.
       let v' := match v with
         | .natVal val => toString val
         | .strVal val => s!"\"{val}\""
-    return s!"(:lit {v'})"
+      s!"(:lit {v'})"
-  | .mdata _ expr =>
+    | .mdata _ inner =>
       -- NOTE: Equivalent to expr itself, but mdata influences the prettyprinter
       -- It may become necessary to incorporate the metadata.
-    return (← serialize_expression_ast expr)
+      self inner
     | .proj typeName idx struct =>
-    let struct' ← serialize_expression_ast struct
+      let struct' := self struct
-    return s!"(:proj {typeName} {idx} {struct'})"
+      s!"(:proj {typeName} {idx} {struct'})"
-
-  where
   -- Elides all unhygenic names
   binder_info_to_ast : Lean.BinderInfo → String
     | .default => ""
     | .implicit => " :implicit"
     | .strictImplicit => " :strictImplicit"
     | .instImplicit => " :instImplicit"
+  of_name (name: Name) := name_to_ast name sanitize
 
 def serialize_expression (options: Protocol.Options) (e: Expr): MetaM Protocol.Expression := do
   let pp := toString (← Meta.ppExpr e)

Test/Proofs.lean

@@ -189,26 +189,36 @@ def proof_or_comm: TestM Unit := do
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  IO.println "=====  1  ====="
-  state1.print
-  IO.println "=====  2  ====="
-  state2.print
-  IO.println "===== 4_1 ====="
-  state4_1.print
   let (state3_2, _goal) ← match ← state2.execute (goalId := 1) (tactic := "apply Or.inl") with
     | .success state #[goal] => pure (state, goal)
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  IO.println "===== 3_2 ====="
-  state3_2.print
   let state4_2 ← match ← state3_2.execute (goalId := 0) (tactic := "assumption") with
     | .success state #[] => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  IO.println "===== 4_2 ====="
+
-  state4_2.print
+  -- Ensure the proof can continue from `state4_2`.
+  let state2b ← match state2.continue state4_2 with
+    | .error msg => do
+      addTest $ assertUnreachable $ msg
+      return ()
+    | .ok state => pure state
+  addTest $ LSpec.test "state2b" (state2b.goals == [state2.goals.get! 0])
+  let (state3_1, _goal) ← match ← state2b.execute (goalId := 0) (tactic := "apply Or.inr") with
+    | .success state #[goal] => pure (state, goal)
+    | other => do
+      addTest $ assertUnreachable $ other.toString
+      return ()
+  let state4_1 ← match ← state3_1.execute (goalId := 0) (tactic := "assumption") with
+    | .success state #[] => pure state
+    | other => do
+      addTest $ assertUnreachable $ other.toString
+      return ()
+  IO.println "=====  4_1  ====="
+  state4_1.print ({ printNonVisible := false })
 
   return ()
   where