just for bench

src/include/lean/lean.h

@@ -41,6 +41,15 @@ extern "C" {
 #define LEAN_NORETURN __attribute__((noreturn))
 #endif
 
+/* Marks a function whose returned pointer does not alias any other live pointer, like `malloc`. */
+#if defined(__GNUC__) || defined(__clang__)
+#define LEAN_ATTR_MALLOC __attribute__((malloc))
+#elif defined(_MSC_VER)
+#define LEAN_ATTR_MALLOC __declspec(restrict)
+#else
+#define LEAN_ATTR_MALLOC
+#endif
+
 #if defined(__GNUC__) || defined(__clang__)
 #define LEAN_UNLIKELY(x) (__builtin_expect((x), 0))
 #define LEAN_LIKELY(x) (__builtin_expect((x), 1))
@@ -475,22 +484,36 @@ static inline unsigned lean_get_slot_idx(unsigned sz) {
 
 LEAN_EXPORT void lean_inc_heartbeat(void);
 
+#ifdef LEAN_MIMALLOC
+/* Fused allocation entry points implemented in `runtime/mimalloc.cpp`: a single call covering the
+   heartbeat update and the (inlined) mimalloc fast path. Both require `sz` to be a multiple of
+   `LEAN_OBJECT_SIZE_DELTA` and at most `MI_SMALL_SIZE_MAX`. `lean_alloc_small_object_core` also
+   initializes `m_cs_sz`, while `lean_alloc_small_object_raw` leaves the entire header to the
+   caller. */
+LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz);
+LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_raw(unsigned sz);
+#endif
+
 #ifndef __cplusplus
 void * malloc(size_t);  // avoid including big `stdlib.h`
 #endif
 
 static inline lean_object * lean_alloc_small_object(unsigned sz) {
-    lean_inc_heartbeat();
 #ifdef LEAN_MIMALLOC
-    // HACK: emulate behavior of small allocator to avoid `leangz` breakage for now
-    // NOTE: `sz` is known at compile time for most callers
+    // NOTE: `sz` is known at compile time for most callers, folding the branch below
     sz = lean_align(sz, LEAN_OBJECT_SIZE_DELTA);
-    void * mem = sz <= MI_SMALL_SIZE_MAX ? mi_malloc_small(sz) : mi_malloc(sz);
+    if (LEAN_LIKELY(sz <= MI_SMALL_SIZE_MAX)) {
+        return lean_alloc_small_object_core(sz);
+    }
+    lean_inc_heartbeat();
+    void * mem = mi_malloc(sz);
     if (mem == 0) lean_internal_panic_out_of_memory();
     lean_object * o = (lean_object*)mem;
+    // see the `m_cs_sz` comment at `lean_alloc_small_object_core`
     o->m_cs_sz = sz;
     return o;
 #else
+    lean_inc_heartbeat();
     void * mem = malloc(sizeof(size_t) + sz);
     if (mem == 0) lean_internal_panic_out_of_memory();
     *(size_t*)mem = sz;
@@ -771,9 +794,37 @@ static inline uint8_t * lean_ctor_scalar_cptr(lean_object * o) {
 
 static inline lean_object * lean_alloc_ctor(unsigned tag, unsigned num_objs, unsigned scalar_sz) {
     assert(tag <= LeanMaxCtorTag && num_objs < LEAN_MAX_CTOR_FIELDS && scalar_sz < LEAN_MAX_CTOR_SCALARS_SIZE);
-    lean_object * o = lean_alloc_ctor_memory(lean_usize_add_checked(lean_usize_add_checked(sizeof(lean_ctor_object), lean_usize_mul_checked(sizeof(void*), num_objs)), scalar_sz));
+    size_t sz = lean_usize_add_checked(lean_usize_add_checked(sizeof(lean_ctor_object), lean_usize_mul_checked(sizeof(void*), num_objs)), scalar_sz);
+#ifdef LEAN_MIMALLOC
+    // NOTE: `sz` is known at compile time for most callers, folding the branches below
+    size_t sz1 = lean_align(sz, LEAN_OBJECT_SIZE_DELTA);
+    lean_object * o;
+    if (LEAN_LIKELY(sz1 <= MI_SMALL_SIZE_MAX)) {
+        o = lean_alloc_small_object_raw((unsigned)sz1);
+    } else {
+        lean_inc_heartbeat();
+        void * mem = mi_malloc(sz1);
+        if (mem == 0) lean_internal_panic_out_of_memory();
+        o = (lean_object*)mem;
+    }
+    if (sz1 > sz) {
+        /* Zero the last word so that the (sz1 - sz) uninitialized trailing bytes do not make the
+           structural comparisons in `maxsharing.cpp` and `compact.cpp` miss sharing. */
+        ((size_t*)((char*)o + sz1))[-1] = 0;
+    }
+    /* Write the full header with adjacent stores: for the constant arguments of compiled code this
+       becomes a single store, whereas `lean_set_st_header` would have to preserve `m_cs_sz` across
+       the opaque allocation call with a read-modify-write. */
+    lean_internal_set_rc(o, 1);
+    o->m_cs_sz = (unsigned)sz1;
+    o->m_other = num_objs;
+    o->m_tag   = tag;
+    return o;
+#else
+    lean_object * o = lean_alloc_ctor_memory((unsigned)sz);
     lean_set_st_header(o, tag, num_objs);
     return o;
+#endif
 }
 
 static inline b_lean_obj_res lean_ctor_get(b_lean_obj_arg o, unsigned i) {

src/runtime/CMakeLists.txt

@@ -1,22 +1,31 @@
 if(USE_MIMALLOC)
-  set(MIMALLOC_SRC ${LEAN_BINARY_DIR}/../mimalloc/src/mimalloc/src/static.c)
-  # (C flags are incomplete, compile as C++ instead like everything else). A source property rather
-  # than a target one, but it is the same for every variant below.
-  set_source_files_properties(${MIMALLOC_SRC} PROPERTIES LANGUAGE CXX)
+  # `mimalloc.cpp` #includes all of mimalloc (`static.c`, compiled as C++ like everything else)
+  # together with the runtime's fused allocation entry points so mimalloc's fast path inlines into
+  # them.
+  set(MIMALLOC_SRC ${CMAKE_CURRENT_SOURCE_DIR}/mimalloc.cpp)
   # make all symbols visible, always build with optimizations as otherwise Lean becomes too slow
   set(MIMALLOC_OPTS -DMI_SHARED_LIB -DMI_SHARED_LIB_EXPORT -O3 -DNDEBUG -DMI_WIN_NOREDIRECT -Wno-unused-function)
   if(CMAKE_CXX_COMPILER_ID MATCHES "AppleClang|Clang")
     list(APPEND MIMALLOC_OPTS -Wno-deprecated)
   endif()
+  if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+    # GCC refuses to inline mimalloc's exported allocation functions into the entry points when
+    # compiling with -fPIC unless told that no symbol will be interposed (clang's default)
+    list(APPEND MIMALLOC_OPTS -fno-semantic-interposition)
+  endif()
 
   # One target per (mitigation level, TLS model) the build needs, each carrying its own
-  # `target_compile_options`. The flags cannot live on the shared source instead: source properties
-  # are directory-scoped and CMake emits them *after* target options, so a target trying to override
-  # the mitigation level would silently lose to whatever the source property set.
+  # `target_compile_options`.
   function(add_mimalloc_variant name type secure)
     add_library(${name} ${type} ${MIMALLOC_SRC})
-    # Lean code includes it as `lean/mimalloc.h`, but `static.c` itself wants the original directory
-    target_include_directories(${name} PRIVATE ${LEAN_BINARY_DIR}/../mimalloc/src/mimalloc/include)
+    # Lean code includes it as `lean/mimalloc.h`, but `static.c` itself wants the original include
+    # directory, and `mimalloc.cpp` includes `static.c` from the source directory
+    target_include_directories(
+      ${name}
+      PRIVATE
+        ${LEAN_BINARY_DIR}/../mimalloc/src/mimalloc/include
+        ${LEAN_BINARY_DIR}/../mimalloc/src/mimalloc/src
+    )
     target_compile_options(${name} PRIVATE ${MIMALLOC_OPTS} -DMI_SECURE=${secure} ${ARGN})
   endfunction()
 
@@ -82,6 +91,10 @@ set(
   uv/signal.cpp
   openssl.cpp
 )
+if(NOT USE_MIMALLOC)
+  # with mimalloc, `object_rc.cpp` is #included into `mimalloc.cpp` instead (see there)
+  list(APPEND RUNTIME_OBJS object_rc.cpp)
+endif()
 
 add_library(leanrt_initial-exec STATIC ${RUNTIME_OBJS})
 set_target_properties(leanrt_initial-exec PROPERTIES ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})

src/runtime/alloc.cpp

@@ -15,6 +15,12 @@ Author: Leonardo de Moura
 #define LEAN_NOINLINE
 #endif
 
+#ifdef _MSC_VER
+extern "C" __declspec(thread) uint64_t lean_g_heartbeat = 0;
+#else
+extern "C" __thread uint64_t lean_g_heartbeat = 0;
+#endif
+
 namespace lean {
 
 void initialize_alloc() {
@@ -23,22 +29,20 @@ void initialize_alloc() {
 void finalize_alloc() {
 }
 
-LEAN_THREAD_VALUE(uint64_t, g_heartbeat, 0);
-
 void set_heartbeats(uint64_t count) {
-    g_heartbeat = count;
+    lean_g_heartbeat = count;
 }
 
 void add_heartbeats(uint64_t count) {
-    g_heartbeat += count;
+    lean_g_heartbeat += count;
 }
 
 extern "C" LEAN_EXPORT void lean_inc_heartbeat() {
     add_heartbeats(1);
 }
 
 uint64_t get_num_heartbeats() {
-    return g_heartbeat;
+    return lean_g_heartbeat;
 }
 
 }

src/runtime/alloc.h

@@ -9,6 +9,17 @@ Author: Leonardo de Moura
 #include <stdint.h>
 #include <lean/lean.h>
 
+/* The heartbeat counter of the current thread; incremented on every small-object allocation.
+   Non-static so that the fused allocation entry points in `mimalloc.cpp` can increment it without
+   a function call. */
+extern "C" {
+#ifdef _MSC_VER
+extern __declspec(thread) uint64_t lean_g_heartbeat;
+#else
+extern __thread uint64_t lean_g_heartbeat;
+#endif
+}
+
 namespace lean {
 LEAN_EXPORT void * alloc(size_t sz);
 LEAN_EXPORT void dealloc(void * o, size_t sz);

src/runtime/mimalloc.cpp

@@ -0,0 +1,44 @@
+/*
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Author: Julia M. Himmel
+*/
+/*
+Compiles all of mimalloc (`static.c`) together with the runtime's small-object allocation entry
+points in a single translation unit so that mimalloc's allocation fast path inlines into them:
+compiled Lean code reaches the allocator, including the heartbeat update, with a single call.
+`static.c` must come first so that mimalloc's headers configure themselves from this variant's
+`MI_*` flags (see `add_mimalloc_variant` in `CMakeLists.txt`) before `lean/mimalloc.h` is seen.
+*/
+#include <static.c>
+// mimalloc's `atomic.h` already defined `_Atomic` for C++; `lean.h` redefines it identically
+#undef _Atomic
+#include <lean/lean.h>
+#include "runtime/alloc.h"
+
+extern "C" LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz) {
+    lean_g_heartbeat++;
+    /* The callers guarantee `sz <= MI_SMALL_SIZE_MAX`, so `mi_malloc_small` applies; unlike
+       `mi_malloc` it does not have to test the size on the fast path. */
+    void * mem = mi_malloc_small(sz);
+    if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory();
+    lean_object * o = (lean_object *)mem;
+    /* `m_cs_sz` must be the exact (aligned) requested size, not mimalloc's potentially larger
+       block size: `lean_small_object_size` and `leangz` rely on it. */
+    o->m_cs_sz = sz;
+    return o;
+}
+
+extern "C" LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_raw(unsigned sz) {
+    lean_g_heartbeat++;
+    void * mem = mi_malloc_small(sz);
+    if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory();
+    /* The caller initializes the entire header, including `m_cs_sz`; leaving it out here keeps
+       `sz` dead after the allocation, so this compiles to a minimal leaf-like fast path. */
+    return (lean_object *)mem;
+}
+
+/* Big-object allocation and the RC deletion machinery, in this TU for the same reason as the
+   entry points above: `mi_malloc`/`mi_free` inline into them. */
+#include "object_rc.cpp"