// block_alloc.dsa
// Fixed-size block allocator
//
// Memory layout:
// [base + 0]: free list head pointer (pointer to first free block, or 0 if none)
// [base + 4]: block size
// [base + 8]: total blocks
// [base + 12]: base address of block pool
// [base + 16+]: block pool (each block starts with a 4-byte next pointer)
//
// Usage:
// include block_alloc "./lib/memory/block_alloc.dsa"
//
// For init:
//   push num_blocks (e.g., 32)
//   push block_size (e.g., 64 bytes)
//   call block_alloc::init
//   pop block_size
//   pop num_blocks
//   ; result in spr+8 (allocator handle)
//
// For alloc:
//   push allocator_handle
//   call block_alloc::alloc
//   pop allocator_handle
//   ; result in spr+8 (pointer to block, or 0 if out of memory)
//
// For free:
//   push block_pointer
//   push allocator_handle
//   call block_alloc::free
//   pop allocator_handle
//   pop block_pointer

dw heap_start: 0x30000  // Start of our heap area

// Initialize the allocator
// Args: block_size, num_blocks
// Returns: allocator handle (pointer to metadata)
init:
    push bpr
    mov spr, bpr

    ldw bpr, rg0, 8   // block_size
    ldw bpr, rg1, 12  // num_blocks

    // Allocate metadata (16 bytes) + pool space
    ldw heap_start, rg2  // base address for this allocator
    mov rg2, rg3         // save base in rg3

    // Calculate total size needed: 16 + (block_size * num_blocks)
    // We'll use a simple multiplication by repeated addition
    mov rg0, rg4         // block_size to rg4
    mov rg1, rg5         // num_blocks to rg5
    lli 0, acc           // accumulator for multiplication

_multiply_loop:
    cmp rg5, zero
    jeq _multiply_done
    add acc, rg4, acc
    dec rg5
    jmp _multiply_loop

_multiply_done:
    // acc now contains block_size * num_blocks
    addi acc, 16         // add metadata size

    // Update heap_start for next allocation
    add rg2, acc, acc
    stw acc, heap_start

    // Now set up metadata at rg3 (base)
    // [base + 0]: free list head (will point to first block)
    // [base + 4]: block_size
    // [base + 8]: total blocks
    // [base + 12]: pool base address

    addi rg3, 16, rg6     // rg6 = pool base

    stw rg6, rg3          // store pool base as free list head initially
    stw rg0, rg3, 4       // store block_size
    stw rg1, rg3, 8       // store total blocks
    stw rg6, rg3, 12      // store pool base address

    // Now initialize the free list
    // Each block's first 4 bytes point to the next block
    // rg6 = current block pointer
    // rg0 = block_size
    // rg1 = num_blocks (counter)

    dec rg1              // we'll count down from num_blocks-1

_init_loop:
    cmp rg1, zero
    jeq _init_loop_done

    // Calculate next block address: current + block_size
    add rg6, rg0, rg7    // rg7 = next block address

    // Store next pointer at current block
    stw rg7, rg6

    // Move to next block
    mov rg7, rg6
    dec rg1
    jmp _init_loop

_init_loop_done:
    // Last block points to null (0)
    lli 0, acc
    stw acc, rg6

    // Return allocator handle (base address - 16 to get back to metadata start)
    stw rg3, bpr, 8

    mov bpr, spr
    pop bpr
    return


// Allocate a block
// Args: allocator_handle
// Returns: pointer to block (or 0 if out of memory)
alloc:
    push bpr
    mov spr, bpr

    ldw bpr, rg0, 8      // allocator handle (metadata base)

    // Load free list head
    ldw rg0, rg1         // rg1 = free list head

    // Check if free list is empty
    cmp rg1, zero
    jeq _alloc_fail

    // Free list is not empty, pop the first block
    // Load the next pointer from the block we're allocating
    ldw rg1, rg2         // rg2 = next free block

    // Update free list head to point to next block
    stw rg2, rg0

    // Return the allocated block (rg1)
    stw rg1, bpr, 8
    jmp _alloc_done

_alloc_fail:
    // No free blocks, return 0
    lli 0, acc
    stw acc, bpr, 8

_alloc_done:
    mov bpr, spr
    pop bpr
    return


// Free a block
// Args: allocator_handle, block_pointer
// Returns: nothing (but could return error code if block is invalid)
free:
    push bpr
    mov spr, bpr

    ldw bpr, rg0, 8      // allocator handle
    ldw bpr, rg6, 12     // pointer to the block pointer to free
    ldw rg6, rg1         // rg1 = block pointer to free

    // Load current free list head
    ldw rg0, rg2         // rg2 = current head

    // Set the freed block's next pointer to current head
    stw rg2, rg1

    // Update free list head to point to freed block
    stw rg1, rg0

    // Update the freed block's previous pointer to NULL
    lli 0, rg1
    stw rg1, rg6

    mov bpr, spr
    pop bpr
    return


// Debug function: get stats
// Args: allocator_handle
// Returns: nothing (but could populate a stats structure)
get_stats:
    push bpr
    mov spr, bpr

    ldw bpr, rg0, 8      // allocator handle

    // Count free blocks by traversing the free list
    ldw rg0, rg1         // rg1 = free list head
    lli 0, rg2           // rg2 = counter

count_loop:
    cmp rg1, zero
    jeq count_done

    inc rg2
    ldw rg1, rg1         // move to next block
    jmp count_loop

count_done:
    // rg2 now contains number of free blocks
    // Could store this somewhere or return it
    stw rg2, bpr, 8

    mov bpr, spr
    pop bpr
    return