hakmem/core/box/carve_push_box.c

// carve_push_box.c - Box Carve-And-Push Implementation
#include <stdio.h>
#include <stdlib.h>
#include <stdatomic.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" // ss_active_add(), SuperSlab
#include "../hakmem_tiny_integrity.h" // HAK_CHECK_CLASS_IDX
#include "../tiny_region_id.h"      // HEADER_MAGIC, HEADER_CLASS_MASK
#include "carve_push_box.h"
#include "capacity_box.h"           // box_cap_has_room()
#include "tls_sll_box.h"            // tls_sll_push()
#include "tiny_next_ptr_box.h"      // tiny_next_write()
#include "../tiny_refill_opt.h"     // TinyRefillChain, trc_linear_carve()
#include "../tiny_box_geometry.h"   // tiny_stride_for_class(), tiny_slab_base_for_geometry()

// External declarations
extern __thread TinyTLSSlab g_tls_slabs[TINY_NUM_CLASSES];
extern __thread TinyTLSSLL g_tls_sll[TINY_NUM_CLASSES];

// ============================================================================
// Internal Helpers
// ============================================================================

// Rollback: return carved blocks to freelist
static void rollback_carved_blocks(int class_idx, TinySlabMeta* meta,
                                    void* head, uint32_t count) {
    // Walk the chain and prepend to freelist
    void* node = head;
    for (uint32_t i = 0; i < count && node; i++) {
        void* next = tiny_next_read(class_idx, node);
        // Prepend to freelist
        tiny_next_write(class_idx, node, meta->freelist);
        meta->freelist = node;
        node = next;
    }
    // Rollback metadata counters
    meta->carved = (uint16_t)((uint32_t)meta->carved - count);
    meta->used = (uint16_t)((uint32_t)meta->used - count);
}

// ============================================================================
// Box Carve-Push API Implementation
// ============================================================================

uint32_t box_carve_and_push(int class_idx, uint32_t want) {
    // PRIORITY 1: Bounds check
    HAK_CHECK_CLASS_IDX(class_idx, "box_carve_and_push");

    if (want == 0) return 0;

    // Step 1: Check TLS SLL capacity
    if (!box_cap_has_room(class_idx, want)) {
        // Not enough room in TLS SLL
        return 0;
    }

    // Step 2: Get TLS slab
    TinyTLSSlab* tls = &g_tls_slabs[class_idx];
    if (!tls->ss || !tls->meta) {
        // No SuperSlab available
        return 0;
    }

    TinySlabMeta* meta = tls->meta;

    // Step 3: Check if slab has enough uncarved blocks
    uint32_t available = (meta->capacity > meta->carved)
                         ? (meta->capacity - meta->carved) : 0;
    if (available < want) {
        // Not enough uncarved blocks
        // Note: Could try superslab_refill() here, but keeping it simple for now
        return 0;
    }

    // Step 4: Get stride and slab base
    size_t bs = tiny_stride_for_class(class_idx);
    uint8_t* slab_base = tls->slab_base ? tls->slab_base
                                        : tiny_slab_base_for_geometry(tls->ss, tls->slab_idx);

    // Step 5: Carve blocks (builds a linked chain)
    TinyRefillChain chain;
    trc_linear_carve(slab_base, bs, meta, want, class_idx, &chain);

    // Sanity check
    if (chain.count != want) {
        // Carve failed to produce expected count
        // This should not happen, but handle defensively
        #if !HAKMEM_BUILD_RELEASE
        fprintf(stderr, "[BOX_CARVE_PUSH] WARN: carved %u blocks but expected %u\n",
                chain.count, want);
        #endif
        // Rollback metadata (carved/used already updated by trc_linear_carve)
        meta->carved = (uint16_t)((uint32_t)meta->carved - chain.count);
        meta->used = (uint16_t)((uint32_t)meta->used - chain.count);
        return 0;
    }

    // Step 6: Push all blocks to TLS SLL (with rollback on failure)
    uint32_t sll_cap = box_cap_get(class_idx);
    uint32_t pushed = 0;
    void* node = chain.head;

    for (uint32_t i = 0; i < want && node; i++) {
        void* next = tiny_next_read(class_idx, node);

        if (!tls_sll_push(class_idx, node, sll_cap)) {
            // Push failed (SLL full or other error)
            // Rollback: pop all pushed blocks and return to freelist
            #if !HAKMEM_BUILD_RELEASE
            fprintf(stderr, "[BOX_CARVE_PUSH] Push failed at block %u/%u, rolling back\n",
                    i, want);
            #endif

            // Pop the blocks we just pushed
            uint32_t pop_failed = 0;
            for (uint32_t j = 0; j < pushed; j++) {
                void* popped;
                if (tls_sll_pop(class_idx, &popped)) {
                    // Return to freelist
                    tiny_next_write(class_idx, popped, meta->freelist);
                    meta->freelist = popped;
                } else {
                    // Pop failed - block remains orphaned in SLL!
                    pop_failed++;
                    #if !HAKMEM_BUILD_RELEASE
                    fprintf(stderr, "[BOX_CARVE_PUSH_ROLLBACK] Pop failed for block %u/%u (cls=%d)\n",
                            j, pushed, class_idx);
                    fprintf(stderr, "[BOX_CARVE_PUSH_ROLLBACK] Block orphaned in TLS SLL - potential double-free risk!\n");
                    #endif
                }
            }

            #if !HAKMEM_BUILD_RELEASE
            if (pop_failed > 0) {
                fprintf(stderr, "[BOX_CARVE_PUSH_ROLLBACK] WARNING: %u/%u blocks orphaned in SLL (cls=%d)\n",
                        pop_failed, pushed, class_idx);
            }
            #endif

            // Return remaining unpushed blocks to freelist
            while (node) {
                void* next_unpushed = tiny_next_read(class_idx, node);
                tiny_next_write(class_idx, node, meta->freelist);
                meta->freelist = node;
                node = next_unpushed;
            }

            // Rollback metadata counters
            meta->carved = (uint16_t)((uint32_t)meta->carved - want);
            meta->used = (uint16_t)((uint32_t)meta->used - want);

            return 0; // All-or-nothing: return 0 on failure
        }

        pushed++;
        node = next;
    }

    // Step 7: Update active counter (all blocks successfully pushed)
    ss_active_add(tls->ss, want);

    return want; // Success: all blocks pushed
}

uint32_t box_carve_and_push_with_freelist(int class_idx, uint32_t want) {
    // PRIORITY 1: Bounds check
    HAK_CHECK_CLASS_IDX(class_idx, "box_carve_and_push_with_freelist");

    if (want == 0) return 0;

    // Step 1: Check capacity
    if (!box_cap_has_room(class_idx, want)) {
        return 0;
    }

    // Step 2: Get TLS slab
    TinyTLSSlab* tls = &g_tls_slabs[class_idx];
    if (!tls->ss || !tls->meta) {
        return 0;
    }

    TinySlabMeta* meta = tls->meta;
    uint32_t sll_cap = box_cap_get(class_idx);
    uint32_t pushed = 0;

    // Step 3: Try freelist first
    while (pushed < want && meta->freelist) {
        void* p = meta->freelist;
        meta->freelist = tiny_next_read(class_idx, p);
        meta->used++;

        // CRITICAL FIX: Restore header BEFORE pushing to TLS SLL
        // Freelist blocks may have stale data at offset 0
        // (same fix as in tiny_superslab_alloc.inc.h:159-169)
#if HAKMEM_TINY_HEADER_CLASSIDX
        *(uint8_t*)p = HEADER_MAGIC | (class_idx & HEADER_CLASS_MASK);
#endif

        if (!tls_sll_push(class_idx, p, sll_cap)) {
            // Rollback
            tiny_next_write(class_idx, p, meta->freelist);
            meta->freelist = p;
            meta->used--;

            // Rollback all pushed
            for (uint32_t j = 0; j < pushed; j++) {
                void* popped;
                if (tls_sll_pop(class_idx, &popped)) {
                    tiny_next_write(class_idx, popped, meta->freelist);
                    meta->freelist = popped;
                    meta->used--;
                }
            }
            return 0;
        }

        ss_active_add(tls->ss, 1);
        pushed++;
    }

    // Step 4: If still need more, try carving
    if (pushed < want) {
        uint32_t need = want - pushed;
        uint32_t carved = box_carve_and_push(class_idx, need);

        if (carved < need) {
            // Partial failure: rollback freelist pushes
            for (uint32_t j = 0; j < pushed; j++) {
                void* popped;
                if (tls_sll_pop(class_idx, &popped)) {
                    tiny_next_write(class_idx, popped, meta->freelist);
                    meta->freelist = popped;
                    meta->used--;
                    ss_active_add(tls->ss, -1);
                }
            }
            return 0;
        }

        pushed += carved;
    }

    return pushed;
}