Skip to content

Commit e743bce

Browse files
augustepoirouxaugustepoiroux
committed
Add initial code for extracting information about declarations
1 parent fd0236e commit e743bce

File tree

6 files changed

+1187
-2
lines changed

6 files changed

+1187
-2
lines changed

‎REPL/Extract/Declaration.lean‎

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
import REPL.JSON
2+
3+
open Lean Elab System Parser
4+
5+
namespace REPL
6+
7+
/-! Extract information about declarations from info trees
8+
Inspired by <https://github.com/frenzymath/jixia>
9+
-/
10+
11+
/-- See `Lean.Parser.Command.declModifiers` and `Lean.Elab.elabModifiers` -/
12+
def getModifiers (stx : Syntax) (ctx: ContextInfo): DeclModifiers :=
13+
match stx with
14+
| .node _ ``Command.declModifiers args =>
15+
{ docString := stx[0].getOptional?.map (fun stx =>
16+
{ content := stx.prettyPrint.pretty, range := stx.toRange ctx }),
17+
visibility := (stx[2].getOptional?.map (·.prettyPrint.pretty.trim)).getD "regular",
18+
computeKind := (stx[3].getOptional?.map (·.prettyPrint.pretty.trim)).getD "regular",
19+
recKind := (stx[5].getOptional?.map (·.prettyPrint.pretty.trim)).getD "default",
20+
isUnsafe := !stx[4].isNone,
21+
attributes := stx[1].getArgs.toList.flatMap parseAttributes, }
22+
| _ => {}
23+
where
24+
/-- Parse attribute instances from a Term.attributes syntax node
25+
See `Lean.Parser.Term.attributes`.
26+
-/
27+
parseAttributes (attrSyntax : Syntax) : List String :=
28+
match attrSyntax with
29+
| .node _ ``Term.attributes args =>
30+
-- args[0] is "@[", args[1] is the attribute list, args[2] is "]"
31+
if args.size > 1 then
32+
args[1]!.getArgs.toList.flatMap fun inst =>
33+
match inst with
34+
| .node _ ``Term.attrInstance _ => [inst.prettyPrint.pretty.trim]
35+
| _ => []
36+
else []
37+
| _ => []
38+
39+
partial def getIdents (stx : Syntax) : Array Name :=
40+
match stx with
41+
| .node _ _ args => args.flatMap getIdents
42+
| .ident _ _ id _ => #[id]
43+
| _ => #[]
44+
45+
/-- See `Lean.Elab.Term.toBinderViews` -/
46+
def toBinderViews (stx : Syntax) : Array BinderView :=
47+
let k := stx.getKind
48+
if stx.isIdent || k == ``Term.hole then
49+
-- binderIdent
50+
#[{ id := (expandBinderIdent stx), type := mkHole stx, binderInfo := "default" }]
51+
else if k == ``Lean.Parser.Term.explicitBinder then
52+
-- `(` binderIdent+ binderType (binderDefault <|> binderTactic)? `)`
53+
let ids := getBinderIds stx[1]
54+
let type := stx[2]
55+
let optModifier := stx[3]
56+
ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "default" }
57+
else if k == ``Lean.Parser.Term.implicitBinder then
58+
-- `{` binderIdent+ binderType `}`
59+
let ids := getBinderIds stx[1]
60+
let type := stx[2]
61+
ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "implicit" }
62+
else if k == ``Lean.Parser.Term.strictImplicitBinder then
63+
-- `⦃` binderIdent+ binderType `⦄`
64+
let ids := getBinderIds stx[1]
65+
let type := stx[2]
66+
ids.map fun id => { id := (expandBinderIdent id), type := (expandBinderType id type), binderInfo := "strictImplicit" }
67+
else if k == ``Lean.Parser.Term.instBinder then
68+
-- `[` optIdent type `]`
69+
let id := expandOptIdent stx[1]
70+
let type := stx[2]
71+
#[ { id := id, type := type, binderInfo := "instImplicit" } ]
72+
else
73+
#[]
74+
where
75+
getBinderIds (ids : Syntax) : Array Syntax :=
76+
ids.getArgs.map fun (id : Syntax) =>
77+
let k := id.getKind
78+
if k == identKind || k == `Lean.Parser.Term.hole then id
79+
else Syntax.missing
80+
expandBinderType (ref : Syntax) (stx : Syntax) : Syntax :=
81+
if stx.getNumArgs == 0 then mkHole ref else stx[1]
82+
expandBinderIdent (stx : Syntax) : Syntax :=
83+
match stx with
84+
| `(_) => Syntax.missing
85+
| _ => stx
86+
expandOptIdent (stx : Syntax) : Syntax :=
87+
if stx.isNone then Syntax.missing else stx[0]
88+
89+
def getScope (ctx : ContextInfo) (state : Command.State) : ScopeInfo :=
90+
let scope := state.scopes.head!
91+
{
92+
varDecls := scope.varDecls.map fun stx => s!"variable {stx.raw.prettyPrint.pretty}",
93+
includeVars := scope.includedVars.toArray.map fun name => name.eraseMacroScopes,
94+
omitVars := scope.omittedVars.toArray.map fun name => name.eraseMacroScopes,
95+
levelNames := scope.levelNames.toArray,
96+
currNamespace := ctx.currNamespace,
97+
openDecl := ctx.openDecls,
98+
}
99+
100+
partial def extractDeclarationInfo (cmdInfo : CommandInfo) (infoTree : InfoTree) (ctx : ContextInfo)
101+
(prevState : Command.State) : DeclarationInfo :=
102+
let stx := cmdInfo.stx
103+
let modifiers := getModifiers stx[0] ctx
104+
let decl := stx[1]
105+
let kind := decl.getKind
106+
let kindStr := match kind with
107+
| .str _ s => s
108+
| _ => kind.toString
109+
110+
-- See `Lean.Elab.DefView`
111+
let (signature, id, binders, type?, value?) := match kind with
112+
| ``Command.abbrev
113+
| ``Command.definition =>
114+
let (binders, type) := expandOptDeclSig decl[2]
115+
(decl[2], decl[1], binders, type, some decl[3])
116+
| ``Command.theorem =>
117+
let (binders, type) := expandDeclSig decl[2]
118+
(decl[2], decl[1], binders, some type, some decl[3])
119+
| ``Command.opaque =>
120+
let (binders, type) := expandDeclSig decl[2]
121+
(decl[2], decl[1], binders, some type, decl[3].getOptional?)
122+
| ``Command.axiom =>
123+
let (binders, type) := expandDeclSig decl[2]
124+
(decl[2], decl[1], binders, some type, none)
125+
| ``Command.inductive
126+
| ``Command.classInductive =>
127+
let (binders, type) := expandOptDeclSig decl[2]
128+
(decl[2], decl[1], binders, type, some decl[4])
129+
| ``Command.instance =>
130+
let (binders, type) := expandDeclSig decl[4]
131+
let declId := match stx[3].getOptional? with
132+
| some declId => declId
133+
| none => Syntax.missing -- TODO: improve this
134+
(decl[4], declId, binders, some type, some decl[5])
135+
| ``Command.example =>
136+
let id := mkIdentFrom stx[0] `_example (canonical := true)
137+
let declId := mkNode ``Parser.Command.declId #[id, mkNullNode]
138+
let (binders, type) := expandOptDeclSig decl[1]
139+
(decl[1], declId, binders, type, some decl[2])
140+
| ``Command.structure =>
141+
let signature := Syntax.node2 .none ``Command.optDeclSig decl[2] decl[4]
142+
let (binders, type) := (decl[2], some decl[4])
143+
(signature, decl[1], binders, type, none)
144+
| _ => unreachable!
145+
146+
let name := id[0].getId
147+
let fullName := match getFullname infoTree name with -- TODO: could be better
148+
| [] => name
149+
| a :: _ => a
150+
151+
let binderViews := binders.getArgs.flatMap toBinderViews
152+
let binders : Option DeclBinders := match binders.getArgs with
153+
| #[] => none
154+
| _ => some { pp := binders.prettyPrint.pretty,
155+
groups := binders.getArgs.map (·.prettyPrint.pretty),
156+
map := binders.getArgs.flatMap toBinderViews,
157+
range := binders.toRange ctx }
158+
159+
let a := prevState.env.constants.find! decl[1].getId
160+
-- a.getUsedConstantsAsSet
161+
162+
let extractConstants (stx : Syntax) : Array Name := -- TODO: improve this
163+
let idents := ((getIdents stx).foldl
164+
(init := NameSet.empty) fun acc name => acc.insert name).toArray
165+
idents
166+
-- idents.filter prevState.env.constants.contains
167+
168+
{
169+
pp := stx.prettyPrint.pretty,
170+
name,
171+
fullName,
172+
kind := kindStr,
173+
modifiers,
174+
scope := getScope ctx prevState,
175+
signature := { pp := signature.prettyPrint.pretty,
176+
constants := extractConstants signature,
177+
range := signature.toRange ctx },
178+
binders,
179+
type := type?.map fun t =>
180+
{ pp := t.prettyPrint.pretty, constants := extractConstants t, range := t.toRange ctx },
181+
value := value?.map fun v =>
182+
{ pp := v.prettyPrint.pretty, constants := extractConstants v, range := v.toRange ctx },
183+
range := stx.toRange ctx
184+
}
185+
where
186+
getFullname (infoTree : InfoTree) (shortName : Name) : List Name :=
187+
match infoTree with
188+
| .context _ t => getFullname t shortName
189+
| .node i ts =>
190+
match i with
191+
| .ofTermInfo ti => ti.expr.constName?.toList.filter (fun n =>
192+
match shortName.componentsRev with
193+
| [] => true
194+
| h :: _ => match n.componentsRev with
195+
| [] => false
196+
| h' :: _ => h == h'
197+
)
198+
| _ => ts.toList.flatMap (getFullname · shortName)
199+
| _ => []
200+
201+
open Lean.Parser in
202+
/-- Extract all declarations from InfoTrees -/
203+
partial def extractCmdDeclarationInfos (trees : List InfoTree) (prevState : Command.State) :
204+
List DeclarationInfo :=
205+
let allDecls := trees.map (extractFromTree · none prevState)
206+
allDecls.flatten
207+
where
208+
extractFromTree (t : InfoTree) (ctx? : Option ContextInfo) (prevState : Command.State) :
209+
List DeclarationInfo :=
210+
match t with
211+
| .context ctx t => extractFromTree t (ctx.mergeIntoOuter? ctx?) prevState
212+
| .node i ts =>
213+
match i with
214+
| .ofCommandInfo cmdInfo =>
215+
match ctx? with
216+
| some ctx =>
217+
if cmdInfo.stx.getKind == ``Command.declaration then
218+
[extractDeclarationInfo cmdInfo t ctx prevState]
219+
else
220+
ts.toList.flatMap (extractFromTree · ctx? prevState)
221+
| none => ts.toList.flatMap (extractFromTree · ctx? prevState)
222+
| _ => ts.toList.flatMap (extractFromTree · ctx? prevState)
223+
| _ => []
224+
225+
def extractAllDeclarationInfos (treeList : List (IncrementalState × Option InfoTree)) (prevState : Command.State) : List DeclarationInfo :=
226+
match treeList with
227+
| [] => []
228+
| (state, infoTree?) :: rest =>
229+
let decls := extractCmdDeclarationInfos infoTree?.toList prevState
230+
let restDecls := extractAllDeclarationInfos rest state.commandState
231+
decls ++ restDecls

‎REPL/JSON.lean‎

Lines changed: 89 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,15 @@ Authors: Scott Morrison
66
import Lean.Data.Json
77
import Lean.Message
88
import Lean.Elab.InfoTree.Main
9+
import REPL.Lean.InfoTree.ToJson
910

1011
open Lean Elab InfoTree
1112

1213
namespace REPL
1314

1415
structure CommandOptions where
1516
allTactics : Option Bool := none
17+
declarations: Option Bool := none
1618
rootGoals : Option Bool := none
1719
/--
1820
Should be "full", "tactics", "original", or "substantive".
@@ -117,6 +119,91 @@ def Tactic.of (goals tactic : String) (pos endPos : Lean.Position) (proofState :
117119
proofState,
118120
usedConstants }
119121

122+
structure DocString where
123+
content : String
124+
range : Syntax.Range
125+
deriving ToJson
126+
127+
/-- See `Lean.Elab.Modifiers` -/
128+
structure DeclModifiers where
129+
docString : Option DocString := none
130+
visibility : String := "regular" -- "regular", "private", "protected", or "public"
131+
computeKind : String := "regular" -- "regular", "meta", or "noncomputable"
132+
recKind : String := "default" -- "default", "partial", or "nonrec"
133+
isUnsafe : Bool := false
134+
attributes : List String := []
135+
deriving ToJson
136+
137+
structure DeclSignature where
138+
pp : String
139+
constants : Array Name
140+
range : Syntax.Range
141+
deriving ToJson
142+
143+
structure BinderView where
144+
id : Syntax
145+
type : Syntax
146+
binderInfo : String
147+
148+
instance : ToJson BinderView where
149+
toJson (bv : BinderView) : Json :=
150+
Json.mkObj [
151+
("id", toString bv.id.prettyPrint),
152+
("type", toString bv.type.prettyPrint),
153+
("binderInfo", bv.binderInfo)
154+
]
155+
156+
structure DeclBinders where
157+
pp : String
158+
groups : Array String
159+
map : Array BinderView
160+
range : Syntax.Range
161+
deriving ToJson
162+
163+
structure DeclType where
164+
pp : String
165+
constants : Array Name
166+
range : Syntax.Range
167+
deriving ToJson
168+
169+
structure DeclValue where
170+
pp : String
171+
constants : Array Name
172+
range : Syntax.Range
173+
deriving ToJson
174+
175+
local instance : ToJson OpenDecl where
176+
toJson
177+
| .simple ns except => json%{
178+
simple: { «namespace»: $ns, except: $except }
179+
}
180+
| .explicit id declName => json%{
181+
rename: { name: $declName, as: $id }
182+
}
183+
184+
structure ScopeInfo where
185+
varDecls : Array String
186+
includeVars : Array Name
187+
omitVars : Array Name
188+
levelNames : Array Name
189+
currNamespace : Name
190+
openDecl : List OpenDecl
191+
deriving ToJson
192+
193+
structure DeclarationInfo where
194+
pp: String
195+
range : Syntax.Range
196+
scope : ScopeInfo
197+
name : Name
198+
fullName : Name
199+
kind : String
200+
modifiers : DeclModifiers
201+
signature: DeclSignature
202+
binders : Option DeclBinders := none
203+
type : Option DeclType := none
204+
value : Option DeclValue := none
205+
deriving ToJson
206+
120207
/--
121208
A response to a Lean command.
122209
`env` can be used in later calls, to build on the stored environment.
@@ -126,8 +213,8 @@ structure CommandResponse where
126213
messages : List Message := []
127214
sorries : List Sorry := []
128215
tactics : List Tactic := []
216+
declarations: List DeclarationInfo := []
129217
infotree : Option Json := none
130-
deriving FromJson
131218

132219
def Json.nonemptyList [ToJson α] (k : String) : List α → List (String × Json)
133220
| [] => []
@@ -139,6 +226,7 @@ instance : ToJson CommandResponse where
139226
Json.nonemptyList "messages" r.messages,
140227
Json.nonemptyList "sorries" r.sorries,
141228
Json.nonemptyList "tactics" r.tactics,
229+
Json.nonemptyList "declarations" r.declarations,
142230
match r.infotree with | some j => [("infotree", j)] | none => []
143231
]
144232

‎REPL/Main.lean‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ Authors: Scott Morrison
66
import REPL.JSON
77
import REPL.Frontend
88
import REPL.Util.Path
9+
import REPL.Extract.Declaration
910
import REPL.Lean.ContextInfo
1011
import REPL.Lean.Environment
1112
import REPL.Lean.InfoTree
@@ -390,6 +391,9 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
390391
let tactics ← match s.allTactics with
391392
| some true => tacticsCmd incStates initialCmdState.env
392393
| _ => pure []
394+
let declarations := match s.declarations with
395+
| some true => extractAllDeclarationInfos incStates initialCmdState
396+
| _ => []
393397
let cmdSnapshot :=
394398
{ cmdState
395399
cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
@@ -415,7 +419,8 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
415419
{ env,
416420
messages,
417421
sorries,
418-
tactics
422+
tactics,
423+
declarations,
419424
infotree }
420425

421426
def processFile (s : File) : M IO (CommandResponse ⊕ Error) := do

0 commit comments

Comments
 (0)