; heap.s - deterministic generated-program free-list heap (E05-S01).
;
; The caller supplies the region; this module does not choose a generated-
; program heap address. All routines run in the generated program's current
; RAM bank. See src/common/heap.inc for the block layout and API.
;
; Allocation is first-fit in address order. A free block is split when the
; remainder is at least HEAP_MIN_BLOCK. Freeing inserts in address order and
; coalesces with the following free block(s) and then with the preceding one.
; Every boundary, size, magic, and double-free check is always on.

.setcpu "6502"
.include "../common/heap.inc"

.export heap_init, heap_alloc, heap_free
.export heap_base, heap_end, heap_free_head, heap_free_total, heap_error

.segment "ZEROPAGE"
heap_ptr:       .res 2
heap_ptr2:      .res 2

.segment "CODE"

; A/X = base (A must be 0), Y = pages (1..127). C=1 with HEAP_ERR_ARG on a
; malformed region. On success the heap is one free block.
heap_init:
        sta heap_arg
        stx heap_arg+1
        sty heap_arg+2
        lda #HEAP_OK
        sta heap_error

        lda heap_arg+2
        beq @bad
        cmp #HEAP_MAX_PAGES+1
        bcs @bad
        lda heap_arg
        bne @bad

        lda heap_arg+1
        clc
        adc heap_arg+2
        bcs @bad
        cmp #$C1                    ; heap must end at or below $C000
        bcs @bad
        sta heap_end+1
        lda #$00
        sta heap_end

        lda heap_arg+1
        sta heap_base+1
        lda #$00
        sta heap_base

        lda heap_base
        sta heap_free_head
        lda heap_base+1
        sta heap_free_head+1

        lda #$00
        sta heap_free_total
        lda heap_arg+2
        sta heap_free_total+1

        lda heap_base
        sta heap_ptr
        lda heap_base+1
        sta heap_ptr+1
        ldy #0
        lda #$00
        sta (heap_ptr),y            ; size low
        iny
        lda heap_arg+2
        sta (heap_ptr),y            ; size high = pages
        iny
        lda #$00
        sta (heap_ptr),y            ; next low
        iny
        sta (heap_ptr),y            ; next high
        clc
        rts
@bad:
        lda #HEAP_ERR_ARG
        sta heap_error
        sec
        rts

; A/X = requested payload bytes. Returns A/X = payload, C=0; or A/X=0, C=1
; with heap_error. The payload is zeroed.
heap_alloc:
        sta heap_req
        stx heap_req+1
        lda #HEAP_OK
        sta heap_error

        lda heap_req
        clc
        adc #HEAP_HEADER_SIZE
        sta heap_need
        lda heap_req+1
        adc #0
        sta heap_need+1
        bcc :+
        jmp @oom
:
        lda heap_need
        and #1
        beq @rounded
        inc heap_need
        bne @rounded
        inc heap_need+1
        bne :+
        jmp @oom
:
@rounded:
        lda heap_need+1
        bpl :+
        jmp @oom
:
        bne @min_ok
        lda heap_need
        cmp #HEAP_MIN_BLOCK
        bcs @min_ok
        lda #HEAP_MIN_BLOCK
        sta heap_need
        lda #0
        sta heap_need+1
@min_ok:
        lda #0
        sta heap_prev
        sta heap_prev+1
        lda heap_free_head
        sta heap_cur
        lda heap_free_head+1
        sta heap_cur+1

@walk:
        lda heap_cur
        ora heap_cur+1
        bne :+
        jmp @oom
:
        lda heap_cur
        sta heap_ptr
        lda heap_cur+1
        sta heap_ptr+1
        ldy #0
        lda (heap_ptr),y
        sta heap_bsize
        iny
        lda (heap_ptr),y
        and #$7F
        sta heap_bsize+1

        lda heap_bsize+1
        cmp heap_need+1
        bcc @next
        bne @found
        lda heap_bsize
        cmp heap_need
        bcc @next

@found:
        lda heap_bsize
        sec
        sbc heap_need
        sta heap_rem
        lda heap_bsize+1
        sbc heap_need+1
        sta heap_rem+1
        lda heap_rem+1
        beq :+
        jmp @split
:
        lda heap_rem
        cmp #HEAP_MIN_BLOCK
        bcc :+
        jmp @split
:
        jmp @use_whole

@next:
        lda heap_cur
        sta heap_prev
        lda heap_cur+1
        sta heap_prev+1
        ldy #2
        lda (heap_ptr),y
        sta heap_cur
        iny
        lda (heap_ptr),y
        sta heap_cur+1
        jmp @walk

@use_whole:
        lda heap_cur
        sta heap_block
        lda heap_cur+1
        sta heap_block+1
        lda heap_block
        sta heap_ptr
        lda heap_block+1
        sta heap_ptr+1
        ldy #2
        lda (heap_ptr),y
        sta heap_next
        iny
        lda (heap_ptr),y
        sta heap_next+1

        lda heap_prev
        ora heap_prev+1
        beq @uw_head
        lda heap_prev
        sta heap_ptr
        lda heap_prev+1
        sta heap_ptr+1
        ldy #2
        lda heap_next
        sta (heap_ptr),y
        iny
        lda heap_next+1
        sta (heap_ptr),y
        jmp @uw_removed
@uw_head:
        lda heap_next
        sta heap_free_head
        lda heap_next+1
        sta heap_free_head+1
@uw_removed:
        lda heap_free_total
        sec
        sbc heap_bsize
        sta heap_free_total
        lda heap_free_total+1
        sbc heap_bsize+1
        sta heap_free_total+1

        jsr @mark_used
        lda heap_bsize
        sec
        sbc #HEAP_HEADER_SIZE
        sta heap_zero_len
        lda heap_bsize+1
        sbc #0
        sta heap_zero_len+1
        jmp @zero_return

@split:
        lda heap_cur
        sta heap_block
        lda heap_cur+1
        sta heap_block+1
        lda heap_block
        sta heap_ptr
        lda heap_block+1
        sta heap_ptr+1
        ldy #2
        lda (heap_ptr),y
        sta heap_next
        iny
        lda (heap_ptr),y
        sta heap_next+1

        lda heap_block
        clc
        adc heap_need
        sta heap_rem_block
        lda heap_block+1
        adc heap_need+1
        sta heap_rem_block+1

        lda heap_rem_block
        sta heap_ptr2
        lda heap_rem_block+1
        sta heap_ptr2+1
        ldy #0
        lda heap_rem
        sta (heap_ptr2),y
        iny
        lda heap_rem+1
        sta (heap_ptr2),y
        iny
        lda heap_next
        sta (heap_ptr2),y
        iny
        lda heap_next+1
        sta (heap_ptr2),y

        lda heap_prev
        ora heap_prev+1
        beq @sp_head
        lda heap_prev
        sta heap_ptr
        lda heap_prev+1
        sta heap_ptr+1
        ldy #2
        lda heap_rem_block
        sta (heap_ptr),y
        iny
        lda heap_rem_block+1
        sta (heap_ptr),y
        jmp @sp_linked
@sp_head:
        lda heap_rem_block
        sta heap_free_head
        lda heap_rem_block+1
        sta heap_free_head+1
@sp_linked:
        lda heap_need
        sta heap_bsize
        lda heap_need+1
        sta heap_bsize+1
        jsr @mark_used

        lda heap_free_total
        sec
        sbc heap_need
        sta heap_free_total
        lda heap_free_total+1
        sbc heap_need+1
        sta heap_free_total+1

        lda heap_need
        sec
        sbc #HEAP_HEADER_SIZE
        sta heap_zero_len
        lda heap_need+1
        sbc #0
        sta heap_zero_len+1
        jmp @zero_return

@mark_used:
        lda heap_block
        sta heap_ptr
        lda heap_block+1
        sta heap_ptr+1
        ldy #0
        lda heap_bsize
        sta (heap_ptr),y
        iny
        lda heap_bsize+1
        ora #$80
        sta (heap_ptr),y
        iny
        lda #$5A
        sta (heap_ptr),y
        iny
        lda #$A5
        sta (heap_ptr),y
        rts

@zero_return:
        lda heap_block
        clc
        adc #HEAP_HEADER_SIZE
        sta heap_ptr
        sta heap_ptr2
        lda heap_block+1
        adc #0
        sta heap_ptr+1
        sta heap_ptr2+1
        jsr heap_zero
        lda heap_ptr2
        ldx heap_ptr2+1
        clc
        rts

@oom:
        lda #HEAP_ERR_OOM
        sta heap_error
        lda #0
        tax
        sec
        rts

; Zero heap_zero_len bytes at heap_ptr. A=0 and is preserved.
heap_zero:
        ldx heap_zero_len+1
        beq @tail
        ldy #0
        lda #0
@page:
        sta (heap_ptr),y
        iny
        bne @page
        inc heap_ptr+1
        dex
        bne @page
@tail:
        ldy #0
        lda #0
        cpy heap_zero_len
        beq @done
@byte:
        sta (heap_ptr),y
        iny
        cpy heap_zero_len
        bcc @byte
@done:
        rts

; A/X = payload pointer. Null is a successful no-op. On success the block is
; free and coalesced; on failure C=1 and heap_error is set.
heap_free:
        sta heap_arg
        stx heap_arg+1
        lda #HEAP_OK
        sta heap_error
        lda heap_arg
        ora heap_arg+1
        bne :+
        jmp @ok
:
        lda heap_arg
        and #1
        beq :+
        jmp @corrupt
:

        lda heap_arg
        sec
        sbc #HEAP_HEADER_SIZE
        sta heap_block
        lda heap_arg+1
        sbc #0
        sta heap_block+1
        bcs @header_ok
        jmp @corrupt
@header_ok:
        lda heap_block+1
        cmp heap_base+1
        bcs :+
        jmp @corrupt
:
        bne @above_base
        lda heap_block
        cmp heap_base
        bcs :+
        jmp @corrupt
:
@above_base:
        lda heap_block+1
        cmp heap_end+1
        bcc @in_range
        beq :+
        jmp @corrupt
:
        lda heap_block
        cmp heap_end
        bcc :+
        jmp @corrupt
:
@in_range:
        lda heap_block
        sta heap_ptr
        lda heap_block+1
        sta heap_ptr+1
        ldy #0
        lda (heap_ptr),y
        sta heap_bsize
        iny
        lda (heap_ptr),y
        sta heap_bsize+1

        lda heap_bsize+1
        and #$80
        bne :+
        jmp @double
:
        lda heap_bsize+1
        and #$7F
        sta heap_bsize+1

        lda heap_bsize+1
        bne @size_hi
        lda heap_bsize
        cmp #HEAP_MIN_BLOCK
        bcs :+
        jmp @corrupt
:
@size_hi:
        lda heap_bsize
        and #1
        beq :+
        jmp @corrupt
:

        lda heap_block
        clc
        adc heap_bsize
        sta heap_tmp
        lda heap_block+1
        adc heap_bsize+1
        sta heap_tmp+1
        bcc :+
        jmp @corrupt
:
        lda heap_tmp+1
        cmp heap_end+1
        bcc @size_ok
        beq :+
        jmp @corrupt
:
        lda heap_tmp
        cmp heap_end
        beq :+
        jmp @corrupt
:
@size_ok:
        ldy #2
        lda (heap_ptr),y
        cmp #$5A
        beq :+
        jmp @corrupt
:
        iny
        lda (heap_ptr),y
        cmp #$A5
        beq :+
        jmp @corrupt
:

        ldy #0
        lda heap_bsize
        sta (heap_ptr),y
        iny
        lda heap_bsize+1
        sta (heap_ptr),y
        iny
        lda #0
        sta (heap_ptr),y
        iny
        sta (heap_ptr),y

        clc
        lda heap_free_total
        adc heap_bsize
        sta heap_free_total
        lda heap_free_total+1
        adc heap_bsize+1
        sta heap_free_total+1

        lda #0
        sta heap_prev
        sta heap_prev+1
        lda heap_free_head
        sta heap_cur
        lda heap_free_head+1
        sta heap_cur+1

@find:
        lda heap_cur
        ora heap_cur+1
        beq @place
        lda heap_cur+1
        cmp heap_block+1
        bcc @advance
        bne @place
        lda heap_cur
        cmp heap_block
        bcc @advance
        bne :+
        jmp @double
:
        jmp @place
@advance:
        lda heap_cur
        sta heap_prev
        lda heap_cur+1
        sta heap_prev+1
        lda heap_cur
        sta heap_ptr
        lda heap_cur+1
        sta heap_ptr+1
        ldy #2
        lda (heap_ptr),y
        sta heap_cur
        iny
        lda (heap_ptr),y
        sta heap_cur+1
        jmp @find

@place:
        lda heap_cur
        ora heap_cur+1
        beq @merged_next
@merge_next:
        lda heap_block
        clc
        adc heap_bsize
        sta heap_tmp
        lda heap_block+1
        adc heap_bsize+1
        sta heap_tmp+1
        lda heap_tmp+1
        cmp heap_cur+1
        bne @merged_next
        lda heap_tmp
        cmp heap_cur
        bne @merged_next

        lda heap_cur
        sta heap_ptr
        lda heap_cur+1
        sta heap_ptr+1
        ldy #0
        lda (heap_ptr),y
        clc
        adc heap_bsize
        sta heap_bsize
        iny
        lda (heap_ptr),y
        adc heap_bsize+1
        sta heap_bsize+1
        ldy #2
        lda (heap_ptr),y
        sta heap_next
        iny
        lda (heap_ptr),y
        sta heap_next+1
        lda heap_next
        sta heap_cur
        lda heap_next+1
        sta heap_cur+1
        jmp @merge_next

@merged_next:
        lda heap_block
        sta heap_ptr
        lda heap_block+1
        sta heap_ptr+1
        ldy #0
        lda heap_bsize
        sta (heap_ptr),y
        iny
        lda heap_bsize+1
        sta (heap_ptr),y
        iny
        lda heap_cur
        sta (heap_ptr),y
        iny
        lda heap_cur+1
        sta (heap_ptr),y

        lda heap_prev
        ora heap_prev+1
        beq @link_head
        lda heap_prev
        sta heap_ptr
        lda heap_prev+1
        sta heap_ptr+1
        ldy #0
        lda (heap_ptr),y
        sta heap_tmp
        iny
        lda (heap_ptr),y
        sta heap_tmp+1

        lda heap_prev
        clc
        adc heap_tmp
        sta heap_rem
        lda heap_prev+1
        adc heap_tmp+1
        sta heap_rem+1
        lda heap_rem+1
        cmp heap_block+1
        bne @link_prev_block
        lda heap_rem
        cmp heap_block
        bne @link_prev_block

        ldy #0
        lda heap_tmp
        clc
        adc heap_bsize
        sta (heap_ptr),y
        iny
        lda heap_tmp+1
        adc heap_bsize+1
        sta (heap_ptr),y
        iny
        lda heap_cur
        sta (heap_ptr),y
        iny
        lda heap_cur+1
        sta (heap_ptr),y
        jmp @ok
@link_prev_block:
        ldy #2
        lda heap_block
        sta (heap_ptr),y
        iny
        lda heap_block+1
        sta (heap_ptr),y
        jmp @ok
@link_head:
        lda heap_block
        sta heap_free_head
        lda heap_block+1
        sta heap_free_head+1
        jmp @ok

@ok:
        lda #0
        tax
        clc
        rts
@double:
        lda #HEAP_ERR_DOUBLE
        sta heap_error
        lda #0
        tax
        sec
        rts
@corrupt:
        lda #HEAP_ERR_CORRUPT
        sta heap_error
        lda #0
        tax
        sec
        rts

.segment "BSS"
heap_base:       .res 2
heap_end:        .res 2
heap_free_head:  .res 2
heap_free_total: .res 2
heap_error:      .res 1
heap_arg:        .res 3
heap_req:        .res 2
heap_need:       .res 2
heap_bsize:      .res 2
heap_rem:        .res 2
heap_rem_block:  .res 2
heap_cur:        .res 2
heap_prev:       .res 2
heap_next:       .res 2
heap_block:      .res 2
heap_tmp:        .res 2
heap_zero_len:   .res 2