feat: JSON representation of InfoTree in output (#31)

REPL.lean

@@ -1,4 +1,4 @@
 import REPL.Frontend
-import REPL.InfoTree
+import REPL.Lean.InfoTree
 import REPL.JSON
-import REPL.Main
+import REPL.Main

REPL/JSON.lean

@@ -19,6 +19,11 @@ structure Command where
   env : Option Nat
   cmd : String
   allTactics : Option Bool := none
+  /--
+  Should be "full", "tactics", "original", or "substantive".
+  Anything else is ignored.
+  -/
+  infotree : Option String
 deriving ToJson, FromJson
 
 /--
@@ -110,6 +115,7 @@ structure CommandResponse where
   messages : List Message := []
   sorries : List Sorry := []
   tactics : List Tactic := []
+  infotree : Option Json := none
 deriving FromJson
 
 def Json.nonemptyList [ToJson α] (k : String) : List α → List (String × Json)
@@ -121,7 +127,8 @@ instance : ToJson CommandResponse where
     [("env", r.env)],
     Json.nonemptyList "messages" r.messages,
     Json.nonemptyList "sorries" r.sorries,
-    Json.nonemptyList "tactics" r.tactics
+    Json.nonemptyList "tactics" r.tactics,
+    match r.infotree with | some j => [("infotree", j)] | none => []
   ]
 
 /--

REPL/InfoTree.lean → REPL/Lean/InfoTree.lean

@@ -11,6 +11,18 @@ Additional functions to deal with `InfoTree`.
 
 open Lean Elab Meta
 
+namespace Lean.FileMap
+
+/-- Extract the range of a `Syntax` expressed as lines and columns. -/
+-- Extracted from the private declaration `Lean.Elab.formatStxRange`,
+-- in `Lean.Elab.InfoTree.Main`.
+def stxRange (fileMap : FileMap) (stx : Syntax) : Position × Position :=
+  let pos    := stx.getPos?.getD 0
+  let endPos := stx.getTailPos?.getD pos
+  (fileMap.toPosition pos, fileMap.toPosition endPos)
+
+end Lean.FileMap
+
 namespace Lean.Syntax
 
 /-- Check if a `Syntax` is an explicit invocation of the `sorry` tactic. -/
@@ -190,5 +202,60 @@ def tactics (t : InfoTree) : List (ContextInfo × Syntax × List MVarId × Posit
   (t.findTacticNodes.map fun ⟨i, ctx⟩ =>
     ({ ctx with mctx := i.mctxBefore }, i.stx, i.goalsBefore, stxRange ctx.fileMap i.stx))
 
+end Lean.Elab.InfoTree
+
+namespace Lean.Elab.TacticInfo
+
+/-- Return the range of the tactic, as a pair of file positions. -/
+def range (info : TacticInfo) (ctx : ContextInfo) : Position × Position := ctx.fileMap.stxRange info.stx
+
+/-- Pretty print a tactic. -/
+def pp (info : TacticInfo) (ctx : ContextInfo) : IO Format :=
+  ctx.runMetaM {} try
+    Lean.PrettyPrinter.ppTactic ⟨info.stx⟩
+  catch _ =>
+    pure "<failed to pretty print>"
+
+open Meta
+
+/-- Run a tactic on the goals stored in a `TacticInfo`. -/
+def runMetaMGoalsBefore (info : TacticInfo) (ctx : ContextInfo) (x : List MVarId → MetaM α) : IO α := do
+  ctx.runMetaM {} <| Meta.withMCtx info.mctxBefore <| x info.goalsBefore
+
+/-- Run a tactic on the after goals stored in a `TacticInfo`. -/
+def runMetaMGoalsAfter (info : TacticInfo) (ctx : ContextInfo) (x : List MVarId → MetaM α) : IO α := do
+  ctx.runMetaM {} <| Meta.withMCtx info.mctxAfter <| x info.goalsAfter
+
+/-- Run a tactic on the main goal stored in a `TacticInfo`. -/
+def runMetaM (info : TacticInfo) (ctx : ContextInfo) (x : MVarId → MetaM α) : IO α := do
+  match info.goalsBefore.head? with
+  | none => throw <| IO.userError s!"No goals at {← info.pp ctx}"
+  | some g => info.runMetaMGoalsBefore ctx fun _ => do g.withContext <| x g
+
+def mainGoal (info : TacticInfo) (ctx : ContextInfo) : IO Expr :=
+  info.runMetaM ctx (fun g => do instantiateMVars (← g.getType))
+
+def formatMainGoal (info : TacticInfo) (ctx : ContextInfo) : IO Format :=
+  info.runMetaM ctx (fun g => do ppExpr (← instantiateMVars (← g.getType)))
+
+def goalState (info : TacticInfo) (ctx : ContextInfo) : IO (List Format) := do
+  info.runMetaMGoalsBefore ctx (fun gs => gs.mapM fun g => do Meta.ppGoal g)
+
+def goalStateAfter (info : TacticInfo) (ctx : ContextInfo) : IO (List Format) := do
+  info.runMetaMGoalsAfter ctx (fun gs => gs.mapM fun g => do Meta.ppGoal g)
+
+def ppExpr (info : TacticInfo) (ctx : ContextInfo) (e : Expr) : IO Format :=
+  info.runMetaM ctx (fun _ => do Meta.ppExpr (← instantiateMVars e))
+
+end Lean.Elab.TacticInfo
+
+namespace Lean.Elab.InfoTree
+
+/--
+Finds all tactic invocations in an `InfoTree`,
+ignoring structuring tactics (e.g. `by`, `;`, multiline tactics, parenthesized tactics).
+-/
+def substantiveTactics (t : InfoTree) : List (TacticInfo × ContextInfo) :=
+  t.findTacticNodes.filter fun i => i.1.isSubstantive
 
 end Lean.Elab.InfoTree

REPL/Lean/InfoTree/ToJson.lean

@@ -0,0 +1,114 @@
+import REPL.Lean.InfoTree
+import REPL.Lean.ContextInfo
+
+/-!
+# Exporting an `InfoTree` as Json
+
+-/
+
+namespace Lean.Elab
+
+structure InfoTreeNode (α : Type) where
+  kind : String
+  node : Option α
+  children : List Json
+deriving ToJson
+
+deriving instance ToJson for Lean.Position
+
+structure Syntax.Range where
+  synthetic : Bool
+  start : Lean.Position
+  finish : Lean.Position
+deriving ToJson
+
+structure Syntax.Json where
+  pp : Option String
+  -- raw : String
+  range : Range
+deriving ToJson
+
+def _root_.Lean.Syntax.toRange (stx : Syntax) (ctx : ContextInfo) : Syntax.Range :=
+  let pos    := stx.getPos?.getD 0
+  let endPos := stx.getTailPos?.getD pos
+  { start := ctx.fileMap.toPosition pos
+    finish := ctx.fileMap.toPosition endPos
+    synthetic := match stx.getHeadInfo with
+    | .original .. => false
+    | _ => true }
+
+def _root_.Lean.Syntax.toJson (stx : Syntax) (ctx : ContextInfo) (lctx : LocalContext) : IO Syntax.Json := do
+  return {
+    pp := match (← ctx.ppSyntax lctx stx).pretty with
+      | "failed to pretty print term (use 'set_option pp.rawOnError true' for raw representation)" => none
+      | pp => some pp
+    -- raw := toString stx
+    range := stx.toRange ctx }
+
+structure TacticInfo.Json where
+  name : Option Name
+  stx : Syntax.Json
+  goalsBefore : List String
+  goalsAfter : List String
+deriving ToJson
+
+-- Note: this is not responsible for converting the children to Json.
+def TacticInfo.toJson (i : TacticInfo) (ctx : ContextInfo) : IO TacticInfo.Json := do
+  return {
+    name := i.name?
+    stx :=
+    { pp := Format.pretty (← i.pp ctx),
+      -- raw := toString i.info.stx,
+      range := i.stx.toRange ctx },
+    goalsBefore := (← i.goalState ctx).map Format.pretty,
+    goalsAfter := (← i.goalStateAfter ctx).map Format.pretty }
+
+structure CommandInfo.Json where
+  elaborator : Option Name
+  stx : Syntax.Json
+deriving ToJson
+
+def CommandInfo.toJson (info : CommandInfo) (ctx : ContextInfo) : IO CommandInfo.Json := do
+  return {
+    elaborator := match info.elaborator with | .anonymous => none | n => some n,
+    stx := ← info.stx.toJson ctx {} }
+
+structure TermInfo.Json where
+  elaborator : Option Name
+  stx : Syntax.Json
+  expectedType? : Option String
+  expr : String
+  isBinder : Bool
+deriving ToJson
+
+def TermInfo.toJson (info : TermInfo) (ctx : ContextInfo) : IO TermInfo.Json := do
+  return {
+    elaborator := match info.elaborator with | .anonymous => none | n => some n,
+    stx := ← info.stx.toJson ctx info.lctx,
+    expectedType? := ← info.expectedType?.mapM fun ty => do
+      pure (← ctx.ppExpr info.lctx ty).pretty
+    expr := (← ctx.ppExpr info.lctx info.expr).pretty
+    isBinder := info.isBinder }
+
+structure InfoTree.HoleJson where
+  goalState : String
+deriving ToJson
+
+partial def InfoTree.toJson (t : InfoTree) (ctx? : Option ContextInfo) : IO Json := do
+  match t with
+  | .context i t => t.toJson i
+  | .node info children =>
+    if let some ctx := ctx? then
+      let node : Option Json ← match info with
+      | .ofTermInfo    info => some <$> (do pure <| Lean.toJson (← info.toJson ctx))
+      | .ofCommandInfo info => some <$> (do pure <| Lean.toJson (← info.toJson ctx))
+      | .ofTacticInfo  info => some <$> (do pure <| Lean.toJson (← info.toJson ctx))
+      | _                   => pure none
+      return Lean.toJson (InfoTreeNode.mk info.kind node (← children.toList.mapM fun t' => t'.toJson ctx))
+    else throw <| IO.userError "No `ContextInfo` available."
+  | .hole mvarId =>
+    if let some ctx := ctx? then
+     return Lean.toJson (InfoTree.HoleJson.mk (← ctx.runMetaM {} (do Meta.ppGoal mvarId)).pretty)
+    else throw <| IO.userError "No `ContextInfo` available."
+
+end Lean.Elab

REPL/Main.lean

@@ -5,12 +5,12 @@ Authors: Scott Morrison
 -/
 import REPL.JSON
 import REPL.Frontend
-import REPL.InfoTree
 import REPL.Util.Path
 import REPL.Util.TermUnsafe
 import REPL.Lean.ContextInfo
 import REPL.Lean.Environment
-import REPL.InfoTree
+import REPL.Lean.InfoTree
+import REPL.Lean.InfoTree.ToJson
 import REPL.Snapshots
 
 /-!
@@ -216,11 +216,22 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
     cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
       { fileName := "", fileMap := default, tacticCache? := none } }
   let env ← recordCommandSnapshot cmdSnapshot
+  let jsonTrees := match s.infotree with
+  | some "full" => trees
+  | some "tactics" => trees.bind InfoTree.retainTacticInfo
+  | some "original" => trees.bind InfoTree.retainTacticInfo |>.bind InfoTree.retainOriginal
+  | some "substantive" => trees.bind InfoTree.retainTacticInfo |>.bind InfoTree.retainSubstantive
+  | _ => []
+  let infotree := if jsonTrees.isEmpty then
+    none
+  else
+    some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
   return .inl
     { env,
       messages,
       sorries,
-      tactics }
+      tactics
+      infotree }
 
 /--
 Run a single tactic, returning the id of the new proof statement, and the new goals.

test/infotree.expected.out

@@ -0,0 +1,20 @@
+{"messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 7},
+   "endPos": {"line": 1, "column": 8},
+   "data": "unused variable `h` [linter.unusedVariables]"}],
+ "infotree":
+ [{"node":
+   {"stx":
+    {"range":
+     {"synthetic": false,
+      "start": {"line": 1, "column": 28},
+      "finish": {"line": 1, "column": 39}},
+     "pp": "constructor"},
+    "name": "Lean.Parser.Tactic.constructor",
+    "goalsBefore": ["h : Int\n⊢ Nat"],
+    "goalsAfter": []},
+   "kind": "TacticInfo",
+   "children": []}],
+ "env": 0}
+

test/infotree.in

@@ -0,0 +1,2 @@
+{"cmd": "def f (h : Int) : Nat := by constructor", "infotree": "substantive"}
+