perf: make the small-object allocation fast path a leaf

src/include/lean/lean.h

@@ -494,6 +494,11 @@ Requires `sz` to be a positive multiple of `LEAN_OBJECT_SIZE_DELTA` of at most
 `MI_SMALL_SIZE_MAX`. Initializes `m_cs_sz`.
 */
 LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz);
+/*
+As above, but leaving the entire header to the caller, so that a caller writing all four header
+fields can do so with a single store.
+*/
+LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_raw(unsigned sz);
 #endif
 
 #ifndef __cplusplus
@@ -796,9 +801,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 they
+       merge into one. `lean_set_st_header` cannot merge, as it must preserve the `m_cs_sz` that
+       `lean_alloc_small_object_core` has already written. */
+    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/mimalloc.cpp

@@ -36,17 +36,64 @@ extern "C" void lean_mi_theap_cache_init(void) {
     lean_g_tls.mi_theap_default = mi_theap_get_default();
 }
 
+/* The callers guarantee `sz > 0 && sz % LEAN_OBJECT_SIZE_DELTA == 0 && sz <= MI_SMALL_SIZE_MAX`.
+   The bound is also a safety condition, as a larger `sz` would index `pages_free_direct` out of
+   bounds. Restating both lets the compiler fold the bin index into a plain byte offset and emit
+   the `m_cs_sz` store below without a mask. */
+static inline mi_page_t * lean_small_page(mi_theap_t * theap, unsigned sz) {
+    lean_assert(sz > 0 && (sz % MI_INTPTR_SIZE) == 0 && sz <= MI_SMALL_SIZE_MAX);
+#if defined(__GNUC__) || defined(__clang__)
+    if ((sz % MI_INTPTR_SIZE) != 0 || sz == 0 || sz > MI_SMALL_SIZE_MAX) __builtin_unreachable();
+#endif
+    return _mi_theap_get_free_small_page(theap, sz);
+}
+
+/* Wraps mimalloc's generic allocation routine. This and `lean_alloc_small_generic_core` are out
+   of line so that the matching entry points below can tail-jump to them and stay leaf functions;
+   a call that returned would cost them a prologue and epilogue on every allocation, to serve a
+   path they almost never take. Only this routine can fail, so the OOM check belongs here:
+   `mi_theap_malloc_small` would force it to the call site, where it would prevent the tail
+   jump. */
+static mi_decl_noinline lean_object * lean_alloc_small_generic_raw(mi_theap_t * theap, size_t sz) {
+    void * mem = _mi_malloc_generic(theap, sz, 0, NULL);
+    if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory();
+    return (lean_object *)mem;
+}
+
+/* Adds the `m_cs_sz` store, as `lean_alloc_small_object_core` does to
+   `lean_alloc_small_object_raw`. Repeating it here leaves that entry point nothing to do after the
+   call; otherwise `sz` would have to survive it, the call could not be a tail jump, and it would
+   need a stack frame. */
+static mi_decl_noinline lean_object * lean_alloc_small_generic_core(mi_theap_t * theap, size_t sz) {
+    lean_object * o = lean_alloc_small_generic_raw(theap, sz);
+    o->m_cs_sz = sz;
+    return o;
+}
+
 extern "C" LEAN_EXPORT LEAN_ATTR_MALLOC lean_object * lean_alloc_small_object_core(unsigned sz) {
     lean_runtime_tls * tls = &lean_g_tls;
     tls->heartbeat++;
     lean_assert(sz > 0 && sz % LEAN_OBJECT_SIZE_DELTA == 0 && sz <= MI_SMALL_SIZE_MAX);
     /* Feeding the cached theap into mimalloc saves the load of mimalloc's own thread-local: the
        heartbeat update and the theap read share one TLS address computation. */
-    void * mem = mi_theap_malloc_small(tls->mi_theap_default, sz);
-    if (LEAN_UNLIKELY(mem == NULL)) lean_internal_panic_out_of_memory();
-    lean_object * o = (lean_object *)mem;
+    mi_theap_t * const theap = tls->mi_theap_default;
+    mi_page_t * const page = lean_small_page(theap, sz);
+    if (LEAN_UNLIKELY(page->free == NULL)) return lean_alloc_small_generic_core(theap, sz);
+    /* No OOM check: only the generic path can fail, and the test above diverts to it; see
+       `lean_alloc_small_generic_raw` for why the check belongs there. */
+    lean_object * o = (lean_object *)mi_page_malloc_zero(theap, page, sz, false, NULL);
     /* `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_runtime_tls * tls = &lean_g_tls;
+    tls->heartbeat++;
+    mi_theap_t * const theap = tls->mi_theap_default;
+    mi_page_t * const page = lean_small_page(theap, sz);
+    if (LEAN_UNLIKELY(page->free == NULL)) return lean_alloc_small_generic_raw(theap, sz);
+    /* Cannot be `NULL`; see `lean_alloc_small_object_core`. */
+    return (lean_object *)mi_page_malloc_zero(theap, page, sz, false, NULL);
+}