From 84a0cd1f5952d7bb827fcb4ae5cb6e47860dd3ff Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Mon, 19 Oct 2020 14:51:23 -0700
Subject: [PATCH] feat: add `[builtinInit]` attribute

---
 src/Lean/Compiler/IR/EmitC.lean    |  2 +-
 src/Lean/Compiler/IR/EmitUtil.lean |  2 +-
 src/Lean/Compiler/InitAttr.lean    | 44 +++++++++++++++++++++---------
 src/Lean/KeyedDeclsAttribute.lean  |  2 +-
 src/Lean/Parser/Extension.lean     |  2 +-
 5 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/src/Lean/Compiler/IR/EmitC.lean b/src/Lean/Compiler/IR/EmitC.lean
index 527e930923..05d73c7f0f 100644
--- a/src/Lean/Compiler/IR/EmitC.lean
+++ b/src/Lean/Compiler/IR/EmitC.lean
@@ -684,7 +684,7 @@ if isIOUnitInitFn env n then
   emitLn "if (lean_io_result_is_error(res)) return res;"
   emitLn "lean_dec_ref(res);"
 else if d.params.size == 0 then
-  match getInitFnNameFor env d.name with
+  match getInitFnNameFor? env d.name with
   | some initFn =>
     emit "res = "; emitCName initFn; emitLn "(lean_io_mk_world());"
     emitLn "if (lean_io_result_is_error(res)) return res;"
diff --git a/src/Lean/Compiler/IR/EmitUtil.lean b/src/Lean/Compiler/IR/EmitUtil.lean
index 11a7044497..047f9c25d3 100644
--- a/src/Lean/Compiler/IR/EmitUtil.lean
+++ b/src/Lean/Compiler/IR/EmitUtil.lean
@@ -38,7 +38,7 @@ partial def collectFnBody : FnBody → M Unit
 
 def collectInitDecl (fn : Name) : M Unit := do
 let env ← read
-match getInitFnNameFor env fn with
+match getInitFnNameFor? env fn with
 | some initFn => collect initFn
 | _           => pure ()
 
diff --git a/src/Lean/Compiler/InitAttr.lean b/src/Lean/Compiler/InitAttr.lean
index ca05b99e11..5baa1a150b 100644
--- a/src/Lean/Compiler/InitAttr.lean
+++ b/src/Lean/Compiler/InitAttr.lean
@@ -22,8 +22,8 @@ match getIOTypeArg type with
 | some type => isUnitType type
 | _ => false
 
-initialize initAttr : ParametricAttribute Name ←
-registerParametricAttribute `init "initialization procedure for global references" $ fun declName stx => do
+def registerInitAttr (attrName : Name) : IO (ParametricAttribute Name) :=
+registerParametricAttribute attrName "initialization procedure for global references" fun declName stx => do
   let decl ← getConstInfo declName
   match attrParamSyntaxToIdentifier stx with
   | some initFnName =>
@@ -40,22 +40,40 @@ registerParametricAttribute `init "initialization procedure for global reference
       else throwError "initialization function must have type `IO Unit`"
     | _ => throwError "unexpected kind of argument"
 
+initialize regularInitAttr : ParametricAttribute Name ← registerInitAttr `init
+initialize builtinInitAttr : ParametricAttribute Name ← registerInitAttr `builtinInit
+
+def getInitFnNameForCore? (env : Environment) (attr : ParametricAttribute Name) (fn : Name) : Option Name :=
+match attr.getParam env fn with
+| some Name.anonymous => none
+| some n              => some n
+| _                   => none
+
+@[export lean_get_builtin_init_fn_name_for]
+def getBuiltinInitFnNameFor? (env : Environment) (fn : Name) : Option Name :=
+getInitFnNameForCore? env builtinInitAttr fn
+
+@[export lean_get_regular_init_fn_name_for]
+def getRegularInitFnNameFor? (env : Environment) (fn : Name) : Option Name :=
+getInitFnNameForCore? env regularInitAttr fn
+
+-- Only used
 def isIOUnitInitFn (env : Environment) (fn : Name) : Bool :=
-match initAttr.getParam env fn with
-| some Name.anonymous => true
-| _ => false
+let aux (attr : ParametricAttribute Name) : Bool :=
+  match attr.getParam env fn with
+  | some Name.anonymous => true
+  | _ => false
+aux builtinInitAttr || aux regularInitAttr
 
 @[export lean_get_init_fn_name_for]
-def getInitFnNameFor (env : Environment) (fn : Name) : Option Name :=
-match initAttr.getParam env fn with
-| some Name.anonymous => none
-| some n => some n
-| _ => none
+def getInitFnNameFor? (env : Environment) (fn : Name) : Option Name :=
+getBuiltinInitFnNameFor? env fn <|> getRegularInitFnNameFor? env fn
 
 def hasInitAttr (env : Environment) (fn : Name) : Bool :=
-(getInitFnNameFor env fn).isSome
+(getInitFnNameFor? env fn).isSome
 
-def setInitAttr (env : Environment) (declName : Name) (initFnName : Name := Name.anonymous) : Except String Environment :=
-initAttr.setParam env declName initFnName
+def setBuiltinInitAttr (env : Environment) (declName : Name) (initFnName : Name := Name.anonymous) : Except String Environment :=
+-- builtinInitAttr.setParam env declName initFnName -- TODO: use this one
+regularInitAttr.setParam env declName initFnName
 
 end Lean
diff --git a/src/Lean/KeyedDeclsAttribute.lean b/src/Lean/KeyedDeclsAttribute.lean
index 84cc805e8a..6f540caf92 100644
--- a/src/Lean/KeyedDeclsAttribute.lean
+++ b/src/Lean/KeyedDeclsAttribute.lean
@@ -117,7 +117,7 @@ match env.addAndCompile {} decl with
 | Except.error e => do
   msg ← (e.toMessageData {}).toString;
   throw (IO.userError ("failed to emit registration code for builtin '" ++ toString declName ++ "': " ++ msg))
-| Except.ok env  => IO.ofExcept (setInitAttr env name)
+| Except.ok env  => IO.ofExcept (setBuiltinInitAttr env name)
 
 /- TODO: add support for scoped attributes -/
 protected unsafe def init {γ} (df : Def γ) (attrDeclName : Name) : IO (KeyedDeclsAttribute γ) := do
diff --git a/src/Lean/Parser/Extension.lean b/src/Lean/Parser/Extension.lean
index b373a13738..825dd9ae32 100644
--- a/src/Lean/Parser/Extension.lean
+++ b/src/Lean/Parser/Extension.lean
@@ -372,7 +372,7 @@ let decl := Declaration.defnDecl { name := name, lparams := [], type := type, va
 match env.addAndCompile {} decl with
 -- TODO: pretty print error
 | Except.error _ => throw (IO.userError ("failed to emit registration code for builtin parser '" ++ toString declName ++ "'"))
-| Except.ok env  => IO.ofExcept (setInitAttr env name)
+| Except.ok env  => IO.ofExcept (setBuiltinInitAttr env name)
 
 def declareLeadingBuiltinParser (env : Environment) (catName : Name) (declName : Name) (prio : Nat) : IO Environment := -- TODO: use CoreM?
 declareBuiltinParser env `Lean.Parser.addBuiltinLeadingParser catName declName prio