; ***************************************************************************
; copyright   : (C) 2025 by Martin Preuss
; email       : martin@libchipcard.de
;
; ***************************************************************************
; * This file is part of the project "AqHome".                              *
; * Please see toplevel file COPYING of that project for license details.   *
; ***************************************************************************


; Needed vars:
; - HEAP_START
; - HEAP_SIZE


; ***************************************************************************
; defines



.equ HEAP_HEADER_BIT_USED = 0 



; ***************************************************************************
; data

.dseg


heapPtr:     .byte 2
heapUsed:    .byte 2
heapFree:    .byte 2
heapDblFree: .byte 1



; ***************************************************************************
; code

.cseg



; ---------------------------------------------------------------------------
; @routine Heap_Init

Heap_Init:
  ldi xl, LOW(HEAP_START)
  ldi xh, HIGH(HEAP_START)
  ldi r24, LOW((HEAP_SIZE-6) & 0xfffc)  ; size minus chunk header, chunk footer, start, end
  ldi r25, HIGH((HEAP_SIZE-6) & 0xfffc) ; we allocate multiples of 4 bytes
  clr r16
  sts heapUsed, r16
  sts heapUsed+1, r16
  st X+, r16                 ; write start header
  st X+, r16
  sts heapPtr, xl            ; store global vars
  sts heapPtr+1, xh
  sts heapFree, r24
  sts heapFree+1, r25
  st X+, r24                 ; store header of first chunk
  st X+, r25
  add xl, r24
  adc xh, r25
  st X+, r24                 ; store footer of first chunk
  st X+, r25
  st X+, r16                 ; write end footer
  sec
  ret
; @end



; ---------------------------------------------------------------------------
; @routine Heap_Alloc
;
; @param r25:r24 number of bytes to alloc
; @return CFLAG set of okay, cleared otherwise
; @return X start of allocated memory
; @clobbers r16, r17, r18, r19, r24, r25, X

Heap_Alloc:
  rjmp heapAllocFirstFit
; @end



; ---------------------------------------------------------------------------
; @routine Heap_AllocAndZero
;
; Allocates data on the heap and presets all bytes with zero.
;
; @param r25:r24 number of bytes to alloc
; @return CFLAG set of okay, cleared otherwise
; @return X start of allocated memory
; @clobbers r16, r17, r18, r19, r24, r25, X

Heap_AllocAndZero:
  push r24
    push r25
      rcall heapAllocFirstFit
    pop r25
  pop r24
  brcc Heap_AllocAndZero_ret
  push xl
    push xh
      clr r18
Heap_AllocAndZero_loop:
        st X+, r18
        sbiw r25:r24, 1
      brne Heap_AllocAndZero_loop
    pop xh
  pop xl
  sec
Heap_AllocAndZero_ret:
  ret
; @end




; ---------------------------------------------------------------------------
; @routine heapAllocFirstFit
;
; @param r25:r24 number of bytes to alloc
; @return CFLAG set of okay, cleared otherwise
; @return X start of allocated memory
; @clobbers r16, r17, r18, r19, r24, r25, X

heapAllocFirstFit:
  adiw r25:r24, 3 ; align size to next multiple of 4
  andi r24, 0xfc  ; mask lower two bits
  
  lds xl, heapPtr
  lds xh, heapPtr+1

heapAllocFirstFit_loop:
    ld r18, X+                    ; read chunk header
    ld r19, X+
    mov r16, r18                  ; heap end reached?
    or r16, r19
    breq heapAllocFirstFit_memFull
    sbrc r18, HEAP_HEADER_BIT_USED
      rjmp heapAllocFirstFit_next ; jump if used
    ; current chunk free, check size
    andi r18, 0xfc
    mov r16, r18
    mov r17, r19
    sub r16, r24
    sbc r17, r25                  ; r17:r16=remainder
    brcs heapAllocFirstFit_next   ; too small
    rcall heapUseChunk
    sec
    ret
heapAllocFirstFit_next:
    andi r18, 0xfc
    add xl, r18
    adc xh, r19
    adiw xh:xl, 2 ; skip footer
  rjmp heapAllocFirstFit_loop  

heapAllocFirstFit_memFull:
  clc
  ret
; @end



; ---------------------------------------------------------------------------
; @routine heapAdjustStatsAlloc
;
; @param r25:r24 size of allocated chunk
; @clobbers r16, r17

heapAdjustStatsAlloc:
  ; adjust heapUsed
  lds r16, heapUsed
  lds r17, heapUsed+1
  add r16, r24
  adc r17, r25
  sts heapUsed, r16
  sts heapUsed+1, r17
  
  ; adjust heapFree
  lds r16, heapFree
  lds r17, heapFree+1
  sub r16, r24
  sbc r17, r25
  sts heapFree, r16
  sts heapFree+1, r17
  ret
; @end



; ---------------------------------------------------------------------------
; @routine heapAdjustStatsFree
;
; @param r25:r24 size of released chunk
; @clobbers r16, r17

heapAdjustStatsFree:
  ; adjust heapUsed
  lds r16, heapUsed
  lds r17, heapUsed+1
  sub r16, r24
  sbc r17, r25
  sts heapUsed, r16
  sts heapUsed+1, r17
  
  ; adjust heapFree
  lds r16, heapFree
  lds r17, heapFree+1
  add r16, r24
  adc r17, r25
  sts heapFree, r16
  sts heapFree+1, r17
  ret
; @end



; ---------------------------------------------------------------------------
; @routine Heap_Free
;
; @param X pointer to previously allocated data
; @clobbers r16, r17, r24, r25, X

Heap_Free:
  sbiw xh:xl, 2 ; go back to chunk header
  ld r24, X+    ; read allocated size (and USED bit)
  ld r25, X+
  sbrs r24, HEAP_HEADER_BIT_USED
    rjmp Heap_Free_dblFree              ; not in use (double free!!)
  cbr r24, (1<<HEAP_HEADER_BIT_USED)    ; clear USED bit
  st -X, r25                            ; write back chunk header
  st -X, r24                            ; X now points to begin of chunk header
  push xl
    push xh
      adiw xh:xl, 2                     ; go skip chunk header
      add xl, r24                       ; skip data
      adc xh, r25
      st X+, r24                        ; write chunk footer
      st X+, r25                        ; X points to the next chunk header now
      rcall heapAdjustStatsFree         ; update stats (r16, r17)
      rcall heapCoalecseUp              ; try to coalecse following chunk with this (r16, r17, r24, r25, X)
    pop xh
  pop xl
  rcall heapCoalecseUp                  ; try to coalesce this chunk with predecessor (r16, r17, r24, r25, X)
  rjmp Heap_Free_ret
Heap_Free_dblFree:
  lds r24, heapDblFree
  lds r25, heapDblFree+1
  adiw r25:r24, 1
  breq Heap_Free_ret
  sts heapDblFree, r24
  sts heapDblFree+1, r25
Heap_Free_ret:
  ret
; @end




; ---------------------------------------------------------------------------
; @routine heapCoalecseUp
;
; @param X pointer to chunk header
; @clobbers r16, r17, r24, r25, X

heapCoalecseUp:
  ld r16, X+                    ; end of heap reached(header==0x0000)?
  ld r17, X
  sbiw xh:xl, 1
  or r16, r17
  breq heapCoalecseUp_ret        ; yes, jump
  ; read footer of preceeding chunk 
  ld r25, -X
  ld r24, -X
  ; R25:R24=footer of preceeding chunk (i.e. size of previous chunk), X points to previous footer
  mov r16, r24                   ; check: beginning of heap reached (header==0x0000)?
  or r16, r25
  breq heapCoalecseUp_ret        ; yes, jump
  sbrc r24, HEAP_HEADER_BIT_USED ; block used?
    rjmp heapCoalecseUp_ret      ; jump if used
  ; previous chunk also free, coalesce, X points to current chunk header
  adiw xh:xl, 2
; skip footer
  ld r16, X+
  ld r17, X+
 sbrc r16, HEAP_HEADER_BIT_USED ; current block used?
    rjmp heapCoalecseUp_ret      ; jump if used
 
  ; X now points to beginning of current chunk's data, R17:r16=size of current chunk, R25:R24=size of previous chunk
  ; to go to start of previous chunk header: sub this chunk header, previous chunk footer, previous chunk header and previous data
  sbiw xh:xl, 6
  sub xl, r24
  sbc xh, r25
  ; X points now to the beginning of the previous chunk header
  ; calculate size of new coalecsed chunk: add size of this chunk, size of previous chunk, 1 chunk header and 1 footer which are
  ; no longer needed now (since the one header and footer of the previous chunk now cover both chunks)
  add r24, r16
  adc r25, r17
  adiw r25:r24, 4
  ; write new chunk header
  st X+, r24
  st X+, r25
  ; skip full data size of the coalesced chunk
  add xl, r24
  adc xh, r25
  ; write new chunk footer
  st X+, r24
  st X+, r25
  ; X now points to header of following chunk
heapCoalecseUp_ret:
  ret
; @end



; ---------------------------------------------------------------------------
; @routine heapUseChunk
;
; @param r25:r24 wanted size
; @param r17:r16 size of current chunk minus wanted size (remainder)
; @param X points to data of current chunk
; @clobbers 

heapUseChunk:
  push xl
    push xh
      tst r17
      brne heapUseChunk_split
      cpi r16, 8                     ; at least 8 bytes left?
      brcs heapUseChunk_directAlloc  ; nope, use full chunk
heapUseChunk_split:
      sbiw xh:xl, 2             ; go back to chunk header
      ; X points to start of chunk header
      mov r18, r24
      sbr r18, (1<<HEAP_HEADER_BIT_USED)
      st X+, r18                ; set used bit
      st X+, r25
      add xl, r24
      adc xh, r25
      ; X now points to the start of chunk footer
      st X+, r18
      st X+, r25
      ; create chunk header for new free chunk
      subi r16, 4               ; sub chunk header and footer from size
      sbci r17, 0
      andi r16, 0xfc            ; clear USED bit
      st X+, r16                ; write new chunk header
      st X+, r17
      add xl, r16
      adc xh, r17
      st X+, r16                ; write new chunk footer
      st X+, r17
      ; update stats, pop X and return
      rjmp heapUseChunk_statsPopRet
heapUseChunk_directAlloc:
      ld r17, -X                          ; load allocated size from chunk header
      ld r16, -X
      andi r16, 0xfc                      ; mask off lower two bits
      mov r18, r16
      sbr r18, (1<<HEAP_HEADER_BIT_USED)  
      st X, r18                           ; set USED bit in chunk header
      adiw xh:xl, 2                       ; skip chunk header
      add xl, r16                         ; skip chunk data
      adc xh, r17
      st X, r18                           ; set USED bit in chunk footer
      ; update stats, pop X and return
heapUseChunk_statsPopRet:
      rcall heapAdjustStatsAlloc ; (r16, r17)
    pop xh
  pop xl
  sec
  ret
; @end







