heap_size:
  ; 1 step here means 1 byte heap alloc table + 16 byte heap data
  ; 0x0400 = 1kB table, 16kb data
  dw 0x0400

heap_init:
  push ax ; byte value (0)
  push cx ; loop counter
  push di ; target ptr
  push es
  mov ax, cs
  mov es, ax
  xor ax, ax
  mov cx, [heap_size]
  mov di, heap
  rep stosb ; rep makes this loop cx times, incrementing di, writing al
  pop es
  pop di
  pop cx
  pop ax
  ret

; in	ax	number of bytes to alloc
; out	ds:di
malloc:
  push cx ; number of blocks to allocate
  push bp ; base offset for block table
  push es

  ; transfer 12 MSB from ax to cx
  add ax, 0x000F
  mov cl, 0x04
  shr ax, cl
  mov cx, ax

  mov bp, heap
  mov di, 0x0000
  mov ax, ds
  mov es, ax
  xor ah, ah ; we use ax for reading from al once
  jmp .firstblock

.testlength:
  push cx
  push di
  add di, bp ; scasb doesn't add bp on its own
  repe scasb
  pop di
  pop cx
  je .found
  ; find next free item
  mov al, 0x01
.nextblock:
  add di, ax
.firstblock:
  mov al, [bp+di]
  cmp al, 0x00
  jne .nextblock
  cmp cx, 0x0001 ; fast case: 1 block allocation, skip test if length fits
  jne .testlength
.found:
  mov [bp+di], cl
  mov cl, 0x04
  shl di, cl
  add di, bp
  add di, [heap_size]
  pop es
  pop bp
  pop cx
  ret

; in	ds:di	begin of data block to free (di trashed)
free:
  push cx
  push bp
  mov bp, heap
  sub di, [heap_size]
  sub di, bp
  mov cl, 0x04
  shr di, cl
  mov byte [bp+di], 0x00
  pop bp
  pop cx
  ret