hakmem/core/box/prewarm_box.c

// prewarm_box.c - Box Prewarm Implementation
#include <stdio.h>
#include <stdlib.h>
#include "../hakmem_tiny.h"       // MUST BE FIRST: Base types
#include "../tiny_tls.h"            // TinyTLSSlab type definition
#include "../hakmem_tiny_config.h" // TINY_NUM_CLASSES
#include "../hakmem_tiny_superslab.h" // SuperSlab
#include "../hakmem_tiny_integrity.h" // HAK_CHECK_CLASS_IDX
#include "prewarm_box.h"
#include "capacity_box.h"         // box_cap_init(), box_cap_avail()
#include "carve_push_box.h"       // box_carve_and_push()

// External declarations
extern __thread TinyTLSSlab g_tls_slabs[TINY_NUM_CLASSES];
extern __thread TinyTLSSLL g_tls_sll[TINY_NUM_CLASSES];
extern SuperSlab* superslab_refill(int class_idx);

// ============================================================================
// Box Prewarm API Implementation
// ============================================================================

int box_prewarm_tls(int class_idx, int count) {
    // PRIORITY 1: Bounds check
    HAK_CHECK_CLASS_IDX(class_idx, "box_prewarm_tls");

    if (count <= 0) return 0;

    // Step 1: Ensure capacity system is initialized
    // This is critical to prevent the double-free bug
    box_cap_init();

    // Step 2: Check available capacity
    uint32_t avail = box_cap_avail(class_idx);
    if (avail == 0) {
        // TLS SLL already at capacity
        return 0;
    }

    // Limit count to available capacity
    uint32_t want = (uint32_t)count;
    if (want > avail) {
        want = avail;
    }

    // Step 3: Ensure SuperSlab is available
    TinyTLSSlab* tls = &g_tls_slabs[class_idx];
    if (!tls->ss) {
        // Try to allocate SuperSlab
        if (superslab_refill(class_idx) == NULL) {
            #if !HAKMEM_BUILD_RELEASE
            fprintf(stderr, "[BOX_PREWARM] Failed to allocate SuperSlab for class %d\n",
                    class_idx);
            #endif
            return 0;
        }
        // Reload tls pointer after superslab_refill
        tls = &g_tls_slabs[class_idx];
    }

    // Step 4: Atomically carve and push blocks
    // This uses Box Carve-Push which guarantees no orphaned blocks
    uint32_t pushed = box_carve_and_push(class_idx, want);

    #if !HAKMEM_BUILD_RELEASE
    if (pushed < want) {
        fprintf(stderr, "[BOX_PREWARM] Partial prewarm: requested=%u pushed=%u class=%d\n",
                want, pushed, class_idx);
    }
    #endif

    return (int)pushed;
}

int box_prewarm_needed(int class_idx, int target_count) {
    // PRIORITY 1: Bounds check
    HAK_CHECK_CLASS_IDX(class_idx, "box_prewarm_needed");

    if (target_count <= 0) return 0;

    // Check current count
    uint32_t current = g_tls_sll[class_idx].count;
    if (current >= (uint32_t)target_count) {
        // Already at or above target
        return 0;
    }

    // Return how many more blocks needed
    return (target_count - (int)current);
}