Game Boy `malloc` & `free`

; Untested, but I think it should work.
; FIXED: Free didn't check for ends of heap



INCLUDE "stdlib/hardware.inc"

HEAP_SIZE = 2048


SECTION "Heap", WRAMX, BANK[1]

; Block header format:
;  2 bytes BE - Size; bit 15 is set whenever the block is used
;  2 bytes LE - Pointer to next
;  2 bytes LE - Pointer to prev


Heap:
    ds HEAP_SIZE
.end


SECTION "Heap management functions", ROM0

; Inits the heap.
; @return a  0
; @destroy hl
InitHeap::
    ld hl, Heap
    IF HEAP_SIZE & $8000
        FAIL "Heap size ({HEAP_SIZE}) must not be over $8000. How did you do that anyways?!?"
    ENDC
    ; Size
    ld a, HIGH(HEAP_SIZE - 6) ; Bit 7 is reset
    ld [hli], a
    ld a, LOW(HEAP_SIZE - 6)
    ld [hli], a
    xor a
    ; Next ptr
    ld [hli], a
    ld [hli], a
    ; Prev ptr
    ld [hli], a
    ld [hli], a
    ret

; Allocates a block in memory
; @param bc Requested size
; @return Z Reset if a block is found
; @return hl Pointer to allocated memory
; @destroy a de
Malloc:
    xor a
    ld de, 0 ; Set "found" ptr to NULL
    ld hl, Heap

.searchHeap
    ; Check sign of difference between block size and requested size
    ld a, [hli]
    bit 7, a
    ; hl must point to second byte of designated block on exit
    jr nz, .notThisBlock
    cp b
    jr c, .notThisBlock
    jr nz, .largeEnough
    ; High bytes are equal, check low bytes
    ld a, [hl]
    cp c
    jr c, .notThisBlock
    jr nz, .largeEnough

    ; Block is *exactly* the right size, use it immediately
    dec hl
    set 7, [hl]
    ld de, 6
    add hl, de
    ret

.largeEnough
    ; Check if block is larger than current block (so we avoid splitting in small fragments)
    ld a, d
    or e
    jr z, .useThisBlock ; If no block selected yet, use this one
    dec de
    ld a, [de]
    inc de
    cpl
    inc a
    dec hl
    add a, [hl] ; a = [hl] - [de]
    inc hl
    jr c, .notThisBlock
    ld a, [de]
    cp [hl]
    jr nc, .notThisBlock ; If blocks are of same size, don't pick them
.useThisBlock
    ld e, l
    ld d, h
.notThisBlock
    inc hl
    ld a, [hli]
    ld h, [hl]
    ld l, a
    or h
    jr nz, .searchHeap ; Keep searching if we didn't reach the end of the heap

    ld a, d
    or e
    ret z ; Return if no block has been found

    ; Modify the target block (split, change size & pointers...)
    ; Don't split the target block if the resulting block would be too small
    ld h, d
    ld l, e
    ld a, [hld] ; Read low
    sub 6 + 1 ; Subtract header size (+ 1)
    sub c
    jr nc, .sizeOK
    ld a, [hli]
    sbc b
    jr c, .dontSplit ; Note: Z is reset
    ; Tbh I doubt 1-byte blocks are any use, but I'm not sure of any better solution
    ; I know that *-fit mallocs are bad, but I'm not sure we can afford a more complex one...
    ; Now, modify the current block's header, and create a new block

    ; Add header size to requested size, simplifies later calculations
    ld a, c
    add a, 6
    ld c, a
    adc a, b
    sub c
    ld b, a
    ; Go to high (2nd) byte of new block's prev ptr (hl points to 2nd byte of size)
    ld a, c
    add a, 2
    add a, l
    ld l, a
    ld a, b
    add a, h
    ld h, a
    ; Write prev ptr
    dec de ; de points to first byte of block
    ld a, d
    ld [hld], a
    ld a, e
    ld [hld], a
    ; Copy next ptr
    inc de
    inc de ; de points to "next" ptr
    inc de
    ld a, [de]
    ld [hld], a
    dec de
    ld a, [de]
    ld [hld], a
    ; Write new size
    dec de
    ld a, c
    ld a, [de]
    sub c
    ld [hld], a
    dec de
    ld a, [de]
    sbc b
    ld [hli], a
    ; Subtract header size to get real size
    ld a, c
    sub 6
    ld c, a
    adc a, b
    sub c
    ; Write new size
    or $80 ; Mark block as used
    ld [de], a
    inc de
    ld a, c
    ld [de], a
    ; Z reset from `or $80`
.dontSplit
    ld de, 6 - 1
    add hl, de ; Preserves Z
    ret

; Frees the given block
; @param hl Pointer to the block to be freed
; @destroy a bc de hl
Free::
    ld de, -6
    add hl, de
    bit 7, [hl]
    ret nz ; Prevent double frees
    ; Get ptr to next ptr in `de`
    ld a, l
    add a, 2
    ld e, a
    adc a, h
    sub e
    ld d, a
    ; Get ptr to next block
    ld a, [de]
    ld l, a
    inc de
    ld a, [de]
    ld h, a
    or l
    jr z, .dontMergeWithNext
    ld a, [hli]
    bit 7, a ; 1 cycle slower than `bit 7, [hl]`, but saves one `inc hl`
    jr nz, .dontMergeWithNext
    ; Next block is free, need to merge with them
    and $7F
    ld b, a ; Save size for later
    ld a, [hli]
    add a, 6 ; Add header size while we have this in `a`
    ld c, a
    adc a, b
    sub c
    ld b, a
    ; Copy next ptr
    inc hl
    ld a, [hld]
    ld [de], a
    dec de
    ld a, [hl]
    ld [de], a
    dec de
    ; Add sizes
    ld a, [de]
    add a, c
    ld [de], a
    dec de
    ld a, [de]
    adc a, b
    ld [de], a
    ; Make `de` point to next ptr's 2nd byte
    inc de
    inc de
    inc de
.dontMergeWithNext
    inc de
    ; Get prev ptr
    ld a, [de]
    inc de
    ld l, a
    ld a, [de]
    ld h, a
    or l
    ret z
    ld a, [hli]
    bit 7, a
    ret nz ; End here if we aren't gonna merge
    and $7F
    ld b, a ; Save size for later
    ld a, [hli]
    add a, 6 ; Add header size
    ld c, a
    adc a, b
    sub c
    ld b, a
    ; Copy `next` ptr
    dec de
    dec de
    inc hl
    ld a, [de]
    dec de
    ld [hld], a
    ld a, [de]
    dec de
    ld [hld], a
    ld a, [de]
    dec de
    add a, c
    ld [hld], a
    ld a, [de]
    adc a, b
    ld [hl], a
    ret