Command to get theorem type

REPL/Frontend.lean

@@ -4,11 +4,44 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Scott Morrison
 -/
 import Lean.Elab.Frontend
+import Std.Data.HashMap
 
-open Lean Elab
+open Lean Elab Language
 
 namespace Lean.Elab.IO
 
+partial def IO.processCommandsIncrementally' (inputCtx : Parser.InputContext)
+    (parserState : Parser.ModuleParserState) (commandState : Command.State)
+    (old? : Option IncrementalState) :
+    BaseIO (List (IncrementalState × Option InfoTree)) := do
+  let task ← Language.Lean.processCommands inputCtx parserState commandState
+    (old?.map fun old => (old.inputCtx, old.initialSnap))
+  return go task.get task #[]
+where
+  go initialSnap t commands :=
+    let snap := t.get
+    let commands := commands.push snap
+    -- Opting into reuse also enables incremental reporting, so make sure to collect messages from
+    -- all snapshots
+    let messages := toSnapshotTree initialSnap
+      |>.getAll.map (·.diagnostics.msgLog)
+      |>.foldl (· ++ ·) {}
+    -- In contrast to messages, we should collect info trees only from the top-level command
+    -- snapshots as they subsume any info trees reported incrementally by their children.
+    let trees := commands.map (·.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
+    let result : (IncrementalState × Option InfoTree) :=
+      ({ commandState := { snap.resultSnap.get.cmdState with messages, infoState.trees := trees }
+         , parserState := snap.parserState
+         , cmdPos := snap.parserState.pos
+         , commands := commands.map (·.stx)
+         , inputCtx := inputCtx
+         , initialSnap := initialSnap
+       }, snap.infoTreeSnap.get.infoTree?)
+    if let some next := snap.nextCmdSnap? then
+      result :: go initialSnap next.task commands
+    else
+      [result]
+
 /--
 Wrapper for `IO.processCommands` that enables info states, and returns
 * the new command state
@@ -17,41 +50,48 @@ Wrapper for `IO.processCommands` that enables info states, and returns
 -/
 def processCommandsWithInfoTrees
     (inputCtx : Parser.InputContext) (parserState : Parser.ModuleParserState)
-    (commandState : Command.State) : IO (Command.State × List Message × List InfoTree) := do
+    (commandState : Command.State) (incrementalState? : Option IncrementalState := none)
+    : IO (List (IncrementalState × Option InfoTree) × List Message) := do
   let commandState := { commandState with infoState.enabled := true }
-  let s ← IO.processCommands inputCtx parserState commandState <&> Frontend.State.commandState
-  pure (s, s.messages.toList, s.infoState.trees.toList)
+  let incStates ← IO.processCommandsIncrementally' inputCtx parserState commandState incrementalState?
+  pure (incStates, (incStates.getLast?.map (·.1.commandState.messages.toList)).getD {})
 
 /--
 Process some text input, with or without an existing command state.
-If there is no existing environment, we parse the input for headers (e.g. import statements),
-and create a new environment.
-Otherwise, we add to the existing environment.
+Supports header caching to avoid reprocessing the same imports repeatedly.
 
 Returns:
-1. The header-only command state (only useful when cmdState? is none)
-2. The resulting command state after processing the entire input
-3. List of messages
-4. List of info trees
+1. The resulting command state after processing the entire input
+2. List of messages
+3. List of info trees along with Command.State from the incremental processing
+4. Updated header cache mapping import keys to Command.State
 -/
 def processInput (input : String) (cmdState? : Option Command.State)
+    (incrementalState? : Option IncrementalState := none)
+    (headerCache : Std.HashMap String Command.State)
     (opts : Options := {}) (fileName : Option String := none) :
-    IO (Command.State × Command.State × List Message × List InfoTree) := unsafe do
+    IO (Command.State × List (IncrementalState × Option InfoTree) × List Message × (Std.HashMap String Command.State)) := unsafe do
   Lean.initSearchPath (← Lean.findSysroot)
   enableInitializersExecution
   let fileName   := fileName.getD "<input>"
   let inputCtx   := Parser.mkInputContext input fileName
-
   match cmdState? with
   | none => do
-    -- Split the processing into two phases to prevent self-reference in proofs in tactic mode
     let (header, parserState, messages) ← Parser.parseHeader inputCtx
-    let (env, messages) ← processHeader header opts messages inputCtx
-    let headerOnlyState := Command.mkState env messages opts
-    let (cmdState, messages, trees) ← processCommandsWithInfoTrees inputCtx parserState headerOnlyState
-    return (headerOnlyState, cmdState, messages, trees)
-
+    let importKey := (input.take parserState.pos.byteIdx).trim
+    match headerCache.get? importKey with
+    | some cachedHeaderState => do
+      -- Header is cached, use it as the base command state
+      let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState cachedHeaderState incrementalState?
+      return (cachedHeaderState, incStates, messages, headerCache)
+    | none => do
+      -- Header not cached, process it and cache the result
+      let (env, messages) ← processHeader header opts messages inputCtx
+      let headerOnlyState := Command.mkState env messages opts
+      let headerCache := headerCache.insert importKey headerOnlyState
+      let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState headerOnlyState incrementalState?
+      return (headerOnlyState, incStates, messages, headerCache)
   | some cmdStateBefore => do
     let parserState : Parser.ModuleParserState := {}
-    let (cmdStateAfter, messages, trees) ← processCommandsWithInfoTrees inputCtx parserState cmdStateBefore
-    return (cmdStateBefore, cmdStateAfter, messages, trees)
+    let (incStates, messages) ← processCommandsWithInfoTrees inputCtx parserState cmdStateBefore incrementalState?
+    return (cmdStateBefore, incStates, messages, headerCache)

REPL/JSON.lean

@@ -13,7 +13,10 @@ namespace REPL
 
 structure CommandOptions where
   allTactics : Option Bool := none
+  declTypes: Option Bool := none
+  namespaces: Option Bool := none
   rootGoals : Option Bool := none
+  conclusions : Option Bool := none
   /--
   Should be "full", "tactics", "original", or "substantive".
   Anything else is ignored.
@@ -46,7 +49,7 @@ deriving ToJson, FromJson
 structure Pos where
   line : Nat
   column : Nat
-deriving ToJson, FromJson
+deriving ToJson, FromJson, BEq, Hashable
 
 /-- Severity of a message. -/
 inductive Severity
@@ -116,6 +119,45 @@ def Tactic.of (goals tactic : String) (pos endPos : Lean.Position) (proofState :
     proofState,
     usedConstants }
 
+structure DeclType where
+  pos : Pos
+  endPos : Pos
+  type : String
+  pp : String
+  conclusion : String
+deriving ToJson, FromJson
+
+/-- Construct the JSON representation of a Declaration type. -/
+def DeclType.of (type pp conclusion : String) (pos endPos : Lean.Position) : DeclType :=
+  { pos := ⟨pos.line, pos.column⟩,
+    endPos := ⟨endPos.line, endPos.column⟩,
+    type,
+    pp
+    conclusion }
+
+structure Namespace where
+  pos : Pos
+  endPos : Pos
+  currentNamespace : String
+  openDecls: List String
+  pp : String
+deriving ToJson, FromJson
+
+def Namespace.of (currentNamespace pp : String) (openDecls : List String) (pos endPos : Lean.Position) : Namespace :=
+  { pos := ⟨pos.line, pos.column⟩,
+    endPos := ⟨endPos.line, endPos.column⟩,
+    currentNamespace,
+    openDecls,
+    pp }
+structure Conclusion where
+  pos : Pos
+  endPos : Pos
+  pp: String
+deriving FromJson, ToJson
+
+def Conclusion.of (pp : String) (pos endPos : Lean.Position) : Conclusion :=
+  {pos := ⟨pos.line, pos.column⟩, endPos := ⟨endPos.line, endPos.column⟩, pp := pp}
+
 /--
 A response to a Lean command.
 `env` can be used in later calls, to build on the stored environment.
@@ -125,6 +167,9 @@ structure CommandResponse where
   messages : List Message := []
   sorries : List Sorry := []
   tactics : List Tactic := []
+  decls: List DeclType := []
+  namespaces: List Namespace := []
+  conclusions : List Conclusion := []
   infotree : Option Json := none
 deriving FromJson
 
@@ -138,6 +183,9 @@ instance : ToJson CommandResponse where
     Json.nonemptyList "messages" r.messages,
     Json.nonemptyList "sorries" r.sorries,
     Json.nonemptyList "tactics" r.tactics,
+    Json.nonemptyList "decls" r.decls,
+    Json.nonemptyList "namespaces" r.namespaces,
+    Json.nonemptyList "conclusions" r.conclusions,
     match r.infotree with | some j => [("infotree", j)] | none => []
   ]
 

REPL/Lean/InfoTree.lean

@@ -130,14 +130,17 @@ Keep `.node` nodes and `.hole` nodes satisfying predicates.
 
 Returns a `List InfoTree`, although in most situations this will be a singleton.
 -/
-partial def filter (p : Info → Bool) (m : MVarId → Bool := fun _ => false) :
+partial def filter (p : Info → Bool) (m : MVarId → Bool := fun _ => false) (stop: Info → Bool := fun _ => false) :
     InfoTree → List InfoTree
-  | .context ctx tree => tree.filter p m |>.map (.context ctx)
+  | .context ctx tree => tree.filter p m stop |>.map (.context ctx)
   | .node info children =>
     if p info then
-      [.node info (children.toList.map (filter p m)).flatten.toPArray']
+      if stop info then
+        [.node info children]
+      else
+        [.node info (children.toList.map (filter p m stop)).flatten.toPArray']
     else
-      (children.toList.map (filter p m)).flatten
+      if stop info then children.toList else (children.toList.map (filter p m stop)).flatten
   | .hole mvar => if m mvar then [.hole mvar] else []
 
 /-- Discard all nodes besides `.context` nodes and `TacticInfo` nodes. -/
@@ -190,6 +193,35 @@ def findTacticNodes (t : InfoTree) : List (TacticInfo × ContextInfo × List MVa
   | (.ofTacticInfo i, some ctx, rootGoals) => (i, ctx, rootGoals)
   | _ => none
 
+/-- Returns all `TermInfo` nodes for a given `InfoTree`. -/
+partial def findTermNodes (t : InfoTree) (ctx? : Option ContextInfo := none) :
+  List (TermInfo × ContextInfo) :=
+  let infoCtx := t.findAllInfo ctx? fun (info: Info) =>
+    match info with
+    | .ofTermInfo _ => true
+    | _ => false
+  infoCtx.flatMap fun ⟨info, ctx?⟩ =>
+   match info with
+   | .ofTermInfo i => match ctx? with
+    | some ctx => [(i, ctx)]
+    | _ => []
+   | _ => []
+
+
+/-- Returns all `CommandInfo` nodes for a given `InfoTree`. -/
+partial def findCommandNodes (t : InfoTree) (ctx? : Option ContextInfo := none) :
+  List (CommandInfo × ContextInfo) :=
+  let infoCtx := t.findAllInfo ctx? fun (info: Info) =>
+    match info with
+    | .ofCommandInfo _ => true
+    | _ => false
+  infoCtx.flatMap fun ⟨info, ctx?⟩ =>
+   match info with
+   | .ofCommandInfo i => match ctx? with
+    | some ctx => [(i, ctx)]
+    | _ => []
+   | _ => []
+
 /-- Returns the root goals for a given `InfoTree`. -/
 partial def findRootGoals (t : InfoTree) (ctx? : Option ContextInfo := none) :
   List (TacticInfo × ContextInfo × List MVarId) :=
@@ -258,6 +290,75 @@ def tactics (t : InfoTree) : List (ContextInfo × Syntax × List MVarId × List
       range.snd,
       i.getUsedConstantsAsSet.toArray )
 
+def declType (t : InfoTree) : Option (ContextInfo × Expr × Syntax × LocalContext × Position × Position) :=
+  let terms: List (TermInfo × ContextInfo) := t.findTermNodes
+  match terms.getLast? with
+  | some ⟨i, ctx⟩ =>
+    let range := stxRange ctx.fileMap i.stx
+    (ctx, i.expr, i.stx, i.lctx, range.fst, range.snd)
+  | _ => none
+
+def namespaces (t : InfoTree) : List (CommandInfo × ContextInfo × Position × Position) :=
+  let nodes := t.findCommandNodes |>.filter fun ⟨i, _⟩ =>
+    match i.elaborator with
+    | ``Lean.Elab.Command.elabOpen => true
+    | ``Lean.Elab.Command.elabSection => true
+    | ``Lean.Elab.Command.elabEnd => true
+    | ``Lean.Elab.Command.elabNamespace => true
+    | _ => false
+  nodes.map fun ⟨i, ctx⟩ =>
+    let range := stxRange ctx.fileMap i.stx
+    (i, ctx, range.fst, range.snd)
+
+-- Get the TermInfo with the largest position in the file
+partial def maxTermNode (terms: List (TermInfo × Option ContextInfo)) : Option (TermInfo × ContextInfo) :=
+  match terms.head? with
+  | some ⟨ti, ctx?⟩ =>
+    let largestTail := maxTermNode terms.tail
+    match largestTail with
+    | some ⟨tailinfo, tailctx⟩ =>
+      match ctx? with
+      | some ctx =>
+        let range := stxRange ctx.fileMap ti.stx
+        let tailRange := stxRange tailctx.fileMap tailinfo.stx
+        if tailRange.snd.line > range.snd.line then some ⟨tailinfo, tailctx⟩
+        else
+          if tailRange.snd.column > range.snd.column then some ⟨tailinfo, tailctx⟩
+          else some ⟨ti, ctx⟩
+      | none => some ⟨tailinfo, tailctx⟩
+    | none => match ctx? with
+      | some ctx => some ⟨ti, ctx⟩
+      | none => none
+  | none => none
+
+def conclusion (t : InfoTree) : IO (List (TermInfo × ContextInfo)) := do
+  let only_declarations := fun i => match i with
+    | .ofCommandInfo m => if m.elaborator == ``Lean.Elab.Command.elabDeclaration then
+      match m.stx.getArgs.toList.getLast? with
+      | some stx => stx.getKind == ``Lean.Parser.Command.theorem
+      | none => false
+      else false
+    | _ => false
+  let only_term := fun i => match i with
+    | .ofTermInfo _ => true
+    | _ => false
+
+  let trees := t.filter (fun i => only_declarations i) (stop := only_declarations)
+  let terms_info : List (List (Info × Option ContextInfo)) := trees.map <| fun t => t.findAllInfo none only_term (stop := fun i => Bool.not <| only_declarations i)
+
+  let terms: List (List (TermInfo × Option ContextInfo)) := terms_info.map <| fun t => t.flatMap (fun ⟨i, ctx?⟩ =>
+    match i with
+    | .ofTermInfo ti => match ti.stx.getHeadInfo with
+      | .original .. => [⟨ti, ctx?⟩]
+      | _ => []
+    | _ => [])
+
+  let optional_conclusions := terms.map <| fun ts => maxTermNode ts
+  pure (optional_conclusions.flatMap <| fun c =>
+    match c with
+    | some s => [s]
+    | _ => [])
+
 def rootGoals (t : InfoTree) : List (ContextInfo × List MVarId × Position) :=
   t.findRootGoals.map fun ⟨i, ctx, rootGoals⟩ =>
     let range := stxRange ctx.fileMap i.stx