feat: lazy initialization of closed terms (#12044)

This PR implements lazy initialization of closed terms. Previous work
has already made sure that ~70% of the closed terms occurring in core
can be statically initialized from the binary. With this the remaining
ones are initialized lazily instead of at startup.

For this we implement a small statically initializable lock that goes
with each term. When trying to access the term we quickly check a flag
to say whether it has already been initialized. If not we take the lock
and initialize it, otherwise we dereference the pointer and fetch the
value.

.github/workflows/build-template.yml

@@ -49,7 +49,7 @@ jobs:
       LSAN_OPTIONS: max_leaks=10
       # somehow MinGW clang64 (or cmake?) defaults to `g++` even though it doesn't exist
       CXX: c++
-      MACOSX_DEPLOYMENT_TARGET: 10.15
+      MACOSX_DEPLOYMENT_TARGET: 11.0
     steps:
       - name: Install Nix
         uses: DeterminateSystems/nix-installer-action@main

src/Init/System/ST.lean

@@ -19,7 +19,7 @@ opaque Void.nonemptyType (σ : Type) : NonemptyType.{0}
 instance Void.instNonempty : Nonempty (Void σ) :=
   by exact (Void.nonemptyType σ).property
 
-@[extern "lean_void_mk"]
+@[extern "lean_void_mk", never_extract]
 opaque Void.mk (x : σ) : Void σ
 
 structure ST.Out (σ : Type) (α : Type) where

src/Lean/Compiler/IR/EmitC.lean

@@ -253,17 +253,27 @@ where
   mkHeader {α : Type} [ToString α] (csSz : α) (other : Nat) (tag : Nat) : String :=
     s!"\{.m_rc = 0, .m_cs_sz = {csSz}, .m_other = {other}, .m_tag = {tag}}"
 
+def toTokenName (cppBaseName : String) : String :=
+  s!"{cppBaseName}_once"
+
+def emitFnClosedDecl (decl : Decl) (cppBaseName : String) : M Unit := do
+  emitLn s!"static lean_once_cell_t {toTokenName cppBaseName} = LEAN_ONCE_CELL_INITIALIZER;"
+  emitLn s!"static {toCType decl.resultType} {cppBaseName};"
+
 def emitFnDeclAux (decl : Decl) (cppBaseName : String) (isExternal : Bool) : M Unit := do
   let ps := decl.params
   let env ← getEnv
 
   if isSimpleGroundDecl env decl.name then
     emitGroundDecl decl cppBaseName
+  else if isClosedTermName env decl.name then
+    emitFnClosedDecl decl cppBaseName
   else
     if ps.isEmpty then
-      if isExternal then emit "extern "
-      else if isClosedTermName env decl.name then emit "static "
-      else emit "LEAN_EXPORT "
+      if isExternal then
+        emit "extern "
+      else
+        emit "LEAN_EXPORT "
     else
       if !isExternal then emit "LEAN_EXPORT "
     emit (toCType decl.resultType ++ " " ++ cppBaseName)
@@ -588,18 +598,35 @@ def emitExternCall (f : FunId) (ps : Array Param) (extData : ExternAttrData) (ys
   | some (ExternEntry.inline _ pat)     => do emit (expandExternPattern pat (toStringArgs ys)); emitLn ";"
   | _ => throw s!"failed to emit extern application '{f}'"
 
-def emitLeanFunReference (f : FunId) : M Unit := do
-  if isSimpleGroundDecl (← getEnv) f then
+def emitLeanFunReference (t : IRType) (f : FunId) : M Unit := do
+  let env ← getEnv
+  if isSimpleGroundDecl env f then
     emit s!"((lean_object*)({← toCName f}))"
+  else if isClosedTermName env f then
+    emitClosedTermRead t f
   else
     emitCName f
+where
+  emitClosedTermRead (t : IRType) (f : FunId) : M Unit := do
+    let fn ←
+      match t with
+      | .float => pure "lean_float_once"
+      | .float32 => pure "lean_float32_once"
+      | .uint8 => pure "lean_uint8_once"
+      | .uint16 => pure "lean_uint16_once"
+      | .uint32 => pure "lean_uint32_once"
+      | .uint64 => pure "lean_uint64_once"
+      | .usize => pure "lean_usize_once"
+      | .object | .tobject | .tagged | .void => pure "lean_obj_once"
+      | _ => throw s!"failed to emit closed term read for '{f}'"
+    emit s!"{fn}(&{← toCName f}, &{toTokenName (← toCName f)}, {← toCInitName f})"
 
-def emitFullApp (z : VarId) (f : FunId) (ys : Array Arg) : M Unit := do
+def emitFullApp (z : VarId) (t : IRType) (f : FunId) (ys : Array Arg) : M Unit := do
   emitLhs z
   let decl ← getDecl f
   match decl with
   | .fdecl (xs := ps) .. | .extern (xs := ps) (ext := { entries := [.opaque], .. }) .. =>
-    emitLeanFunReference f
+    emitLeanFunReference t f
     if ys.size > 0 then
       let (ys, _) := ys.zip ps |>.filter (fun (_, p) => !p.ty.isVoid) |>.unzip
       emit "("; emitArgs ys; emit ")"
@@ -676,7 +703,7 @@ def emitVDecl (z : VarId) (t : IRType) (v : Expr) : M Unit :=
   | Expr.proj i x       => emitProj z i x
   | Expr.uproj i x      => emitUProj z i x
   | Expr.sproj n o x    => emitSProj z t n o x
-  | Expr.fap c ys       => emitFullApp z c ys
+  | Expr.fap c ys       => emitFullApp z t c ys
   | Expr.pap c ys       => emitPartialApp z c ys
   | Expr.ap x ys        => emitApp z x ys
   | Expr.box t x        => emitBox z x t
@@ -831,7 +858,7 @@ def emitDeclAux (d : Decl) : M Unit := do
             emit (toCType x.ty); emit " "; emit x.x
         emit ")"
       else
-        emit ("_init_" ++ baseName ++ "()")
+        emit ("_init_" ++ baseName ++ "(void)")
       emitLn " {";
       if xs.size > closureMaxArgs && isBoxedName d.name then
         xs.size.forM fun i _ => do
@@ -888,7 +915,7 @@ def emitDeclInit (d : Decl) : M Unit := do
       if getBuiltinInitFnNameFor? env d.name |>.isSome then
         emit "}"
     | _ =>
-      if !isSimpleGroundDecl env d.name then
+      if !isClosedTermName env d.name && !isSimpleGroundDecl env d.name then
         emitCName n; emit " = "; emitCInitName n; emitLn "();"; emitMarkPersistent d n
 
 def emitInitFn : M Unit := do

src/include/lean/lean.h

@@ -3155,6 +3155,85 @@ static inline lean_obj_res lean_manual_get_root(lean_obj_arg _unit) {
 #define LEAN_SCALAR_PTR_LITERAL(b1, b2, b3, b4, b5, b6, b7, b8) (lean_object*)((uint64_t)b1 | ((uint64_t)b2 << 8) | ((uint64_t)b3 << 16) | ((uint64_t)b4 << 24) | ((uint64_t)b5 << 32) | ((uint64_t)b6 << 40) | ((uint64_t)b7 << 48) | ((uint64_t)b8 << 56))
 #endif
 
+typedef struct {
+    _Atomic(int) state;
+    _Atomic(int) lock;
+} lean_once_cell_t;
+
+#define LEAN_ONCE_CELL_INITIALIZER { .state = 0, .lock = 0 }
+
+LEAN_EXPORT lean_object* lean_obj_once_cold(lean_object** loc, lean_once_cell_t* tok, lean_object* (*init)(void));
+
+static inline lean_object* lean_obj_once(lean_object** loc, lean_once_cell_t* tok, lean_object* (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_obj_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT uint8_t lean_uint8_once_cold(uint8_t* loc, lean_once_cell_t* tok, uint8_t (*init)(void));
+
+static inline uint8_t lean_uint8_once(uint8_t* loc, lean_once_cell_t* tok, uint8_t (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_uint8_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT uint16_t lean_uint16_once_cold(uint16_t* loc, lean_once_cell_t* tok, uint16_t (*init)(void));
+
+static inline uint16_t lean_uint16_once(uint16_t* loc, lean_once_cell_t* tok, uint16_t (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_uint16_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT uint32_t lean_uint32_once_cold(uint32_t* loc, lean_once_cell_t* tok, uint32_t (*init)(void));
+
+static inline uint32_t lean_uint32_once(uint32_t* loc, lean_once_cell_t* tok, uint32_t (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_uint32_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT uint64_t lean_uint64_once_cold(uint64_t* loc, lean_once_cell_t* tok, uint64_t (*init)(void));
+
+static inline uint64_t lean_uint64_once(uint64_t* loc, lean_once_cell_t* tok, uint64_t (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_uint64_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT size_t lean_usize_once_cold(size_t* loc, lean_once_cell_t* tok, size_t (*init)(void));
+
+static inline size_t lean_usize_once(size_t* loc, lean_once_cell_t* tok, size_t (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_usize_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT float lean_float32_once_cold(float* loc, lean_once_cell_t* tok, float (*init)(void));
+
+static inline float lean_float32_once(float* loc, lean_once_cell_t* tok, float (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_float32_once_cold(loc, tok, init);
+}
+
+LEAN_EXPORT double lean_float_once_cold(double* loc, lean_once_cell_t* tok, double (*init)(void));
+
+static inline double lean_float_once(double* loc, lean_once_cell_t* tok, double (*init)(void)) {
+    if (LEAN_LIKELY(tok->state == 1)) {
+        return *loc;
+    }
+    return lean_float_once_cold(loc, tok, init);
+}
+
 #ifdef __cplusplus
 }
 #endif

src/runtime/object.cpp

@@ -2760,4 +2760,102 @@ void finalize_object() {
     delete g_ext_classes;
     delete g_ext_classes_mutex;
 }
+
+void lock_simple_atomic(std::atomic<int>& lock) {
+    while (true) {
+        lock.wait(1);
+        int should = 0;
+        if (lock.compare_exchange_strong(should, 1)) {
+            break;
+        }
+    }
+}
+
+void unlock_simple_atomic(std::atomic<int>& lock) {
+    lock.store(0);
+    lock.notify_one();
+}
+
+extern "C" LEAN_EXPORT lean_object* lean_obj_once_cold(lean_object** loc, lean_once_cell_t* tok, lean_object* (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        lean_mark_persistent(*loc);
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT uint8_t lean_uint8_once_cold(uint8_t* loc, lean_once_cell_t* tok, uint8_t (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT uint16_t lean_uint16_once_cold(uint16_t* loc, lean_once_cell_t* tok, uint16_t (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT uint32_t lean_uint32_once_cold(uint32_t* loc, lean_once_cell_t* tok, uint32_t (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT uint64_t lean_uint64_once_cold(uint64_t* loc, lean_once_cell_t* tok, uint64_t (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT size_t lean_usize_once_cold(size_t* loc, lean_once_cell_t* tok, size_t (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT float lean_float32_once_cold(float* loc, lean_once_cell_t* tok, float (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+extern "C" LEAN_EXPORT double lean_float_once_cold(double* loc, lean_once_cell_t* tok, double (*init)(void)) {
+    lock_simple_atomic(tok->lock);
+    if (tok->state.load() != 1) {
+        *loc = init();
+        tok->state.store(1);
+    }
+    unlock_simple_atomic(tok->lock);
+    return *loc;
+}
+
+
 }