feat: mkElabAttribute

src/Init/Lean/Elab/Command.lean

@@ -49,7 +49,7 @@ instance CommandElabM.monadLog : MonadLog CommandElabM :=
   getFileName := do ctx ← read; pure ctx.fileName,
   logMessage  := fun msg => modify $ fun s => { messages := s.messages.add msg, .. s } }
 
-abbrev CommandElabTable := SMap SyntaxNodeKind CommandElab
+abbrev CommandElabTable := ElabFnTable CommandElab
 def mkBuiltinCommandElabTable : IO (IO.Ref CommandElabTable) := IO.mkRef {}
 @[init mkBuiltinCommandElabTable] constant builtinCommandElabTable : IO.Ref CommandElabTable := arbitrary _
 
@@ -86,17 +86,18 @@ registerAttribute {
  applicationTime := AttributeApplicationTime.afterCompilation
 }
 
-abbrev CommandElabAttribute := ElabAttribute CommandElabTable
-def mkCommandElabAttribute : IO CommandElabAttribute := mkElabAttribute `commandTerm "command" builtinCommandElabTable
+abbrev CommandElabAttribute := ElabAttribute CommandElab
+def mkCommandElabAttribute : IO CommandElabAttribute :=
+mkElabAttribute CommandElab `commandElab `Lean.Parser.Command `Lean.Elab.Command.CommandElab "command" builtinCommandElabTable
 @[init mkCommandElabAttribute] constant commandElabAttribute : CommandElabAttribute := arbitrary _
 
 def elabCommand (stx : Syntax) : CommandElabM Unit :=
 stx.ifNode
   (fun n => do
     s ← get;
-    let tables := commandElabAttribute.ext.getState s.env;
+    let table := (commandElabAttribute.ext.getState s.env).table;
     let k := n.getKind;
-    match tables.find? k with
+    match table.find? k with
     | some elab => elab n
     | none      => throwError stx ("command '" ++ toString k ++ "' has not been implemented"))
   (fun _ => throwError stx ("unexpected command"))

src/Init/Lean/Elab/Term.lean

@@ -126,7 +126,7 @@ instance TermElabM.MonadQuotation : MonadQuotation TermElabM := {
   withFreshMacroScope := @Term.withFreshMacroScope
 }
 
-abbrev TermElabTable := SMap SyntaxNodeKind TermElab
+abbrev TermElabTable := ElabFnTable TermElab
 def mkBuiltinTermElabTable : IO (IO.Ref TermElabTable) :=  IO.mkRef {}
 @[init mkBuiltinTermElabTable] constant builtinTermElabTable : IO.Ref TermElabTable := arbitrary _
 
@@ -163,8 +163,9 @@ registerAttribute {
  applicationTime := AttributeApplicationTime.afterCompilation
 }
 
-abbrev TermElabAttribute := ElabAttribute TermElabTable
-def mkTermElabAttribute : IO TermElabAttribute := mkElabAttribute `elabTerm "term" builtinTermElabTable
+abbrev TermElabAttribute := ElabAttribute TermElab
+def mkTermElabAttribute : IO TermElabAttribute :=
+mkElabAttribute TermElab `termElab `Lean.Parser.Term `Lean.Elab.Term.TermElab "term" builtinTermElabTable
 @[init mkTermElabAttribute] constant termElabAttribute : TermElabAttribute := arbitrary _
 
 /--
@@ -437,9 +438,9 @@ def elabTerm (stx : Syntax) (expectedType? : Option Expr) (catchExPostpone := tr
 withFreshMacroScope $ withNode stx $ fun node => do
   trace `Elab.step stx $ fun _ => stx;
   s ← get;
-  let tables := termElabAttribute.ext.getState s.env;
+  let table := (termElabAttribute.ext.getState s.env).table;
   let k := node.getKind;
-  match tables.find? k with
+  match table.find? k with
   | some elab =>
     catch
       (elab node expectedType?)

src/Init/Lean/Elab/Util.lean

@@ -30,40 +30,93 @@ match attrParamSyntaxToIdentifier arg with
   throw ("invalid syntax node kind '" ++ toString k ++ "'")
 | none   => throw ("syntax node kind is missing")
 
-structure ElabAttributeEntry :=
-(kind     : SyntaxNodeKind)
-(declName : Name)
+structure ElabAttributeOLeanEntry :=
+(kind      : SyntaxNodeKind)
+(constName : Name)
 
-structure ElabAttribute (σ : Type) :=
+structure ElabAttributeEntry (γ : Type) extends ElabAttributeOLeanEntry :=
+(elabFn   : γ)
+
+abbrev ElabFnTable (γ : Type) := SMap SyntaxNodeKind γ
+
+structure ElabAttributeExtensionState (γ : Type) :=
+(newEntries : List ElabAttributeOLeanEntry := [])
+(table      : ElabFnTable γ                := {})
+
+instance ElabAttributeExtensionState.inhabited (γ) : Inhabited (ElabAttributeExtensionState γ) :=
+⟨{}⟩
+
+abbrev ElabAttributeExtension (γ) := PersistentEnvExtension ElabAttributeOLeanEntry (ElabAttributeEntry γ) (ElabAttributeExtensionState γ)
+
+structure ElabAttribute (γ : Type) :=
 (attr : AttributeImpl)
-(ext  : PersistentEnvExtension ElabAttributeEntry ElabAttributeEntry σ)
+(ext  : ElabAttributeExtension γ)
 (kind : String)
 
-instance ElabAttribute.inhabited {σ} [Inhabited σ] : Inhabited (ElabAttribute σ) := ⟨{ attr := arbitrary _, ext := arbitrary _, kind := "" }⟩
+instance ElabAttribute.inhabited {γ} : Inhabited (ElabAttribute γ) := ⟨{ attr := arbitrary _, ext := arbitrary _, kind := "" }⟩
 
-/-
-This is just the basic skeleton for attributes such as `[termElab]` and `[commandElab]` attributes, and associated environment extensions.
-The state is initialized using `builtinTable`.
+private def ElabAttribute.mkInitial {γ} (builtinTableRef : IO.Ref (ElabFnTable γ)) : IO (ElabAttributeExtensionState γ) := do
+table ← builtinTableRef.get;
+pure { table := table }
 
-The current implementation just uses the bultin elaborators.
--/
-def mkElabAttribute {σ} [Inhabited σ] (attrName : Name) (kind : String) (builtinTable : IO.Ref σ) : IO (ElabAttribute σ) := do
-ext : PersistentEnvExtension ElabAttributeEntry ElabAttributeEntry σ ← registerPersistentEnvExtension {
+private def throwUnexpectedElabType {γ} (typeName : Name) (constName : Name) : ExceptT String Id γ :=
+throw ("unexpected elaborator type at '" ++ toString constName ++ "' `" ++ toString typeName ++ "` expected")
+
+private unsafe def mkElabFnOfConstantUnsafe (γ) (env : Environment) (typeName : Name) (constName : Name) : ExceptT String Id γ :=
+match env.find? constName with
+| none      => throw ("unknow constant '" ++ toString constName ++ "'")
+| some info =>
+  match info.type with
+  | Expr.const c _ _ =>
+    if c != typeName then throwUnexpectedElabType typeName constName
+    else env.evalConst γ constName
+  | _ => throwUnexpectedElabType typeName constName
+
+@[implementedBy mkElabFnOfConstantUnsafe]
+constant mkElabFnOfConstant (γ : Type) (env : Environment) (typeName : Name) (constName : Name) : ExceptT String Id γ := throw ""
+
+private def ElabAttribute.addImportedParsers {γ} (typeName : Name) (builtinTableRef : IO.Ref (ElabFnTable γ)) (env : Environment) (es : Array (Array ElabAttributeOLeanEntry))
+    : IO (ElabAttributeExtensionState γ) := do
+table ← builtinTableRef.get;
+table ← es.foldlM
+  (fun table entries =>
+    entries.foldlM
+      (fun (table : ElabFnTable γ) entry =>
+        match mkElabFnOfConstant γ env typeName entry.constName with
+        | Except.ok f     => pure $ table.insert entry.kind f
+        | Except.error ex => throw (IO.userError ex))
+      table)
+  table;
+pure { table := table }
+
+private def ElabAttribute.addExtensionEntry {γ} (s : ElabAttributeExtensionState γ) (e : ElabAttributeEntry γ) : ElabAttributeExtensionState γ :=
+{ table := s.table.insert e.kind e.elabFn, newEntries := e.toElabAttributeOLeanEntry :: s.newEntries }
+
+private def ElabAttribute.add {γ} (parserNamespace : Name) (typeName : Name) (ext : ElabAttributeExtension γ)
+    (env : Environment) (constName : Name) (arg : Syntax) (persistent : Bool) : IO Environment := do
+match mkElabFnOfConstant γ env typeName constName with
+| Except.error ex => throw (IO.userError ex)
+| Except.ok f     => do
+  kind ← IO.ofExcept $ syntaxNodeKindOfAttrParam env parserNamespace arg;
+  pure $ ext.addEntry env { kind := kind, elabFn := f, constName := constName }
+
+/- TODO: add support for scoped attributes -/
+def mkElabAttribute (γ) (attrName : Name) (parserNamespace : Name) (typeName : Name) (kind : String) (builtinTableRef : IO.Ref (ElabFnTable γ)) : IO (ElabAttribute γ) := do
+ext : ElabAttributeExtension γ ← registerPersistentEnvExtension {
   name            := attrName,
-  mkInitial       := pure (arbitrary _),
-  addImportedFn   := fun env es => do
-    table ← builtinTable.get;
-    -- TODO: populate table with `es`
-    pure table,
-  addEntryFn      := fun (s : σ) _ => s,                            -- TODO
-  exportEntriesFn := fun _ => #[],                                  -- TODO
-  statsFn         := fun _ => fmt (kind ++ " elaborator attribute") -- TODO
+  mkInitial       := ElabAttribute.mkInitial builtinTableRef,
+  addImportedFn   := ElabAttribute.addImportedParsers typeName builtinTableRef,
+  addEntryFn      := ElabAttribute.addExtensionEntry,
+  exportEntriesFn := fun s => s.newEntries.reverse.toArray,
+  statsFn         := fun s => format "number of local entries: " ++ format s.newEntries.length
 };
 let attrImpl : AttributeImpl := {
-  name  := attrName,
-  descr := kind ++ " elaborator",
-  add   := fun env decl args persistent => pure env -- TODO
+  name            := attrName,
+  descr           := kind ++ " elaborator",
+  add             := ElabAttribute.add parserNamespace typeName ext,
+  applicationTime := AttributeApplicationTime.afterCompilation
 };
+registerAttribute attrImpl;
 pure { ext := ext, attr := attrImpl, kind := kind }
 
 @[init] private def regTraceClasses : IO Unit := do

tests/playground/termParserAttr.lean

@@ -22,5 +22,11 @@ ParserDescr.node `boo
       (ParserDescr.parser `term 0)
       (ParserDescr.symbol "|]" 0)))
 
+open Lean.Elab.Term
+
+@[termElab tst] def elabTst : TermElab :=
+fun stx expected? =>
+  elabTerm (stx.getArg 1) expected?
+
 #eval run "#check [| @id.{1} Nat |]"
-#eval run "#check (| id 1 |)"
+-- #eval run "#check (| id 1 |)"