Declaration information extraction

REPL/Extract/Declaration.lean

@@ -0,0 +1,231 @@
+import REPL.JSON
+
+open Lean Elab System Parser
+
+namespace REPL
+
+/-! Extract information about declarations from info trees
+Inspired by <https://github.com/frenzymath/jixia>
+-/
+
+/-- See `Lean.Parser.Command.declModifiers` and `Lean.Elab.elabModifiers` -/
+def getModifiers (stx : Syntax) (ctx: ContextInfo): DeclModifiers :=
+  match stx with
+  | .node _ ``Command.declModifiers _ =>
+    { docString := stx[0].getOptional?.map (fun stx =>
+        { content := stx.prettyPrint.pretty, range := stx.toRange ctx }),
+      visibility := (stx[2].getOptional?.map (·.prettyPrint.pretty.trim)).getD "regular",
+      computeKind := (stx[4].getOptional?.map (·.prettyPrint.pretty.trim)).getD "regular",
+      isProtected := !stx[3].isNone,
+      isUnsafe := !stx[5].isNone,
+      recKind := (stx[6].getOptional?.map (·.prettyPrint.pretty.trim)).getD "default",
+      attributes := stx[1].getArgs.toList.flatMap parseAttributes, }
+  | _ => {}
+  where
+    /-- Parse attribute instances from a Term.attributes syntax node
+    See `Lean.Parser.Term.attributes`.
+    -/
+    parseAttributes (attrSyntax : Syntax) : List String :=
+      match attrSyntax with
+      | .node _ ``Term.attributes args =>
+        -- args[0] is "@[", args[1] is the attribute list, args[2] is "]"
+        if args.size > 1 then
+          args[1]!.getArgs.toList.flatMap fun inst =>
+            match inst with
+            | .node _ ``Term.attrInstance _ => [inst.prettyPrint.pretty.trim]
+            | _ => []
+        else []
+      | _ => []
+
+partial def getIdents (stx : Syntax) : Array Name :=
+  match stx with
+  | .node _ _ args => args.flatMap getIdents
+  | .ident _ _ id _ => #[id]
+  | _ => #[]
+
+/-- See `Lean.Elab.Term.toBinderViews` -/
+def toBinderViews (stx : Syntax) : Array BinderView :=
+  let k := stx.getKind
+  if stx.isIdent || k == ``Term.hole then
+    -- binderIdent
+    #[{ id := (expandBinderIdent stx), type := mkHole stx, binderInfo := "default" }]
+  else if k == ``Lean.Parser.Term.explicitBinder then
+    -- `(` binderIdent+ binderType (binderDefault <|> binderTactic)? `)`
+    let ids := getBinderIds stx[1]
+    let type        := stx[2]
+    -- let optModifier := stx[3]
+    ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "default" }
+  else if k == ``Lean.Parser.Term.implicitBinder then
+    -- `{` binderIdent+ binderType `}`
+    let ids := getBinderIds stx[1]
+    let type := stx[2]
+    ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "implicit" }
+  else if k == ``Lean.Parser.Term.strictImplicitBinder then
+    -- `⦃` binderIdent+ binderType `⦄`
+    let ids := getBinderIds stx[1]
+    let type := stx[2]
+    ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "strictImplicit" }
+  else if k == ``Lean.Parser.Term.instBinder then
+    -- `[` optIdent type `]`
+    let id := expandOptIdent stx[1]
+    let type := stx[2]
+    #[ { id := id, type := type, binderInfo := "instImplicit" } ]
+  else
+    #[]
+  where
+    getBinderIds (ids : Syntax) : Array Syntax :=
+      ids.getArgs.map fun (id : Syntax) =>
+        let k := id.getKind
+        if k == identKind || k == `Lean.Parser.Term.hole then id
+        else Syntax.missing
+    expandBinderType (ref : Syntax) (stx : Syntax) : Syntax :=
+      if stx.getNumArgs == 0 then mkHole ref else stx[1]
+    expandBinderIdent (stx : Syntax) : Syntax :=
+      match stx with
+      | `(_) => Syntax.missing
+      | _    => stx
+    expandOptIdent (stx : Syntax) : Syntax :=
+      if stx.isNone then Syntax.missing else stx[0]
+
+def getScope (ctx : ContextInfo) (state : Command.State) : ScopeInfo :=
+  let scope := state.scopes.head!
+  {
+    varDecls := scope.varDecls.map fun stx => s!"variable {stx.raw.prettyPrint.pretty}",
+    includeVars := scope.includedVars.toArray.map fun name => name.eraseMacroScopes,
+    omitVars := scope.omittedVars.toArray.map fun name => name.eraseMacroScopes,
+    levelNames := scope.levelNames.toArray,
+    currNamespace := ctx.currNamespace,
+    openDecl := ctx.openDecls,
+  }
+
+partial def extractDeclarationInfo (cmdInfo : CommandInfo) (infoTree : InfoTree) (ctx : ContextInfo)
+  (prevState : Command.State) : DeclarationInfo :=
+  let stx := cmdInfo.stx
+  let modifiers := getModifiers stx[0] ctx
+  let decl := stx[1]
+  let kind := decl.getKind
+  let kindStr := match kind with
+    | .str _ s => s
+    | _ => kind.toString
+
+  -- See `Lean.Elab.DefView`
+  let (signature, id, binders, type?, value?) := match kind with
+    | ``Command.abbrev
+    | ``Command.definition =>
+      let (binders, type) := expandOptDeclSig decl[2]
+      (decl[2], decl[1], binders, type, some decl[3])
+    | ``Command.theorem =>
+      let (binders, type) := expandDeclSig decl[2]
+      (decl[2], decl[1], binders, some type, some decl[3])
+    | ``Command.opaque =>
+      let (binders, type) := expandDeclSig decl[2]
+      (decl[2], decl[1], binders, some type, decl[3].getOptional?)
+    | ``Command.axiom =>
+      let (binders, type) := expandDeclSig decl[2]
+      (decl[2], decl[1], binders, some type, none)
+    | ``Command.inductive
+    | ``Command.classInductive =>
+      let (binders, type) := expandOptDeclSig decl[2]
+      (decl[2], decl[1], binders, type, some decl[4])
+    | ``Command.instance =>
+      let (binders, type) := expandDeclSig decl[4]
+      let declId := match stx[3].getOptional? with
+        | some declId => declId
+        | none        => Syntax.missing -- TODO: improve this
+      (decl[4], declId, binders, some type, some decl[5])
+    | ``Command.example =>
+      let id              := mkIdentFrom stx[0] `_example (canonical := true)
+      let declId          := mkNode ``Parser.Command.declId #[id, mkNullNode]
+      let (binders, type) := expandOptDeclSig decl[1]
+      (decl[1], declId, binders, type, some decl[2])
+    | ``Command.structure =>
+      let signature := Syntax.node2 .none ``Command.optDeclSig decl[2] decl[4]
+      let (binders, type) := (decl[2], some decl[4])
+      (signature, decl[1], binders, type, none)
+    | _ => unreachable!
+
+  let name := id[0].getId
+  let fullName := match getFullname infoTree name with -- TODO: could be better
+    | [] => name
+    | a :: _  => a
+
+  let binders : Option DeclBinders := match binders.getArgs with
+    | #[] => none
+    | _ => some { pp := binders.prettyPrint.pretty,
+                  groups := binders.getArgs.map (·.prettyPrint.pretty),
+                  map := binders.getArgs.flatMap toBinderViews,
+                  range := binders.toRange ctx }
+
+  -- let a := prevState.env.constants.find! decl[1].getId
+  -- a.getUsedConstantsAsSet
+
+  let extractConstants (stx : Syntax) : Array Name := -- TODO: improve this
+    let idents := ((getIdents stx).foldl
+      (init := NameSet.empty) fun acc name => acc.insert name).toArray
+    idents
+    -- idents.filter prevState.env.constants.contains
+
+  {
+    pp := stx.prettyPrint.pretty,
+    name,
+    fullName,
+    kind := kindStr,
+    modifiers,
+    scope := getScope ctx prevState,
+    signature := { pp := signature.prettyPrint.pretty,
+                    constants := extractConstants signature,
+                    range := signature.toRange ctx },
+    binders,
+    type := type?.map fun t =>
+      { pp := t.prettyPrint.pretty, constants := extractConstants t, range := t.toRange ctx },
+    value := value?.map fun v =>
+      { pp := v.prettyPrint.pretty, constants := extractConstants v, range := v.toRange ctx },
+    range := stx.toRange ctx
+  }
+where
+  getFullname (infoTree : InfoTree) (shortName : Name) : List Name :=
+    match infoTree with
+    | .context _ t => getFullname t shortName
+    | .node i ts  =>
+      match i with
+        | .ofTermInfo ti => ti.expr.constName?.toList.filter (fun n =>
+              match shortName.componentsRev with
+              | [] => true
+              | h :: _ => match n.componentsRev with
+                | [] => false
+                | h' :: _ => h == h'
+            )
+        | _ => ts.toList.flatMap (getFullname · shortName)
+    | _ => []
+
+open Lean.Parser in
+/-- Extract all declarations from InfoTrees -/
+partial def extractCmdDeclarationInfos (trees : List InfoTree) (prevState : Command.State) :
+    List DeclarationInfo :=
+  let allDecls := trees.map (extractFromTree · none prevState)
+  allDecls.flatten
+where
+  extractFromTree (t : InfoTree) (ctx? : Option ContextInfo) (prevState : Command.State) :
+      List DeclarationInfo :=
+    match t with
+    | .context ctx t => extractFromTree t (ctx.mergeIntoOuter? ctx?) prevState
+    | .node i ts =>
+      match i with
+      | .ofCommandInfo cmdInfo =>
+        match ctx? with
+        | some ctx =>
+          if cmdInfo.stx.getKind == ``Command.declaration then
+            [extractDeclarationInfo cmdInfo t ctx prevState]
+          else
+            ts.toList.flatMap (extractFromTree · ctx? prevState)
+        | none => ts.toList.flatMap (extractFromTree · ctx? prevState)
+      | _ => ts.toList.flatMap (extractFromTree · ctx? prevState)
+    | _ => []
+
+def extractAllDeclarationInfos (treeList : List (IncrementalState × Option InfoTree)) (prevState : Command.State) : List DeclarationInfo :=
+  match treeList with
+  | [] => []
+  | (state, infoTree?) :: rest =>
+    let decls := extractCmdDeclarationInfos infoTree?.toList prevState
+    let restDecls := extractAllDeclarationInfos rest state.commandState
+    decls ++ restDecls

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 (·.elabSnap.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
+    let result : (IncrementalState × Option InfoTree) :=
+      ({ commandState := { snap.elabSnap.resultSnap.get.cmdState with messages, infoState.trees := trees }
+         , parserState := snap.parserState
+         , cmdPos := snap.parserState.pos
+         , commands := commands.map (·.stx)
+         , inputCtx := inputCtx
+         , initialSnap := initialSnap
+       }, snap.elabSnap.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 := toString header.raw -- do not contain comments and whitespace
+    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)