malloc.rv

export {
  malloc!, free!, retain!, release!, blockSize, blockCount
  setBlockUsed!, headerSize, pageSize, allocationCount
}

# Current layout:
#
# l l l l u u u u m m m ....
# |       |       |
# |       used    |
# 32-bit size     user memory

# Internally we work with a pointer to the header, and consider block sizes
# to include the header. Pointers are shifted to point to user memory at the
# interface.
# (User) pointers are guaranteed to be 8-byte aligned, as long as all
# allocations are a multiple of 8 bytes.

pageSize = Int32(64*1024) # 64 KiB
headerSize = Int32(8)

fn getBlockSize(p: Ptr) { i32load(p) }
fn setBlockSize!(p: Ptr, sz: Int32) { i32store!(p, sz) }

fn getBlockUsed(p: Ptr) { i32load(p + Int32(4)) }
fn setBlockUsed!(p: Ptr, u: Int32) { i32store!(p + Int32(4), u) }

# Initialise first block
memoryGrow!(Int32(1))
setBlockSize!(Ptr(0), pageSize)

fn memorySize() {
  memoryPages() * pageSize
}

fn npages(bytes: Int32) {
  div(bytes + pageSize - Int32(1), pageSize)
}

# Block struct

bundle Block(ptr: Ptr, size: Int32)

fn getBlockSize(Block(_, size)) { size }

fn setBlockSize!(&bl: Block(ptr, _), size) {
  setBlockSize!(ptr, size)
  bl = Block(ptr, size)
}

fn getBlockUsed(Block(ptr, _)) { getBlockUsed(ptr) }
fn setBlockUsed!(Block(ptr, _), u: Int32) { setBlockUsed!(ptr, u) }

fn firstBlock() {
  Block(Ptr(0), getBlockSize(Ptr(0)))
}

fn next(&bl: Block(ptr, size)) {
  ptr = ptr + size
  bl = Block(ptr, getBlockSize(ptr))
}

fn isLastBlock(Block(ptr, size)) {
  addr(ptr + size) == memorySize()
}

fn publicPtr(Block(ptr, _)) {
  ptr + headerSize
}

# Main logic

# Combine a block with free neighbours to the right
fn coalesce(&bl: Block) {
  isLastBlock(bl) && return
  cl = next(bl)
  getBlockUsed(cl) && return
  coalesce(&cl)
  setBlockSize!(&bl, getBlockSize(bl) + getBlockSize(cl))
  return
}

# Find a free block of size >= n, creating it if necessary.
fn findFreeBlock(n: Int32) {
  bl = firstBlock()
  while not(isLastBlock(bl)) {
    if not(getBlockUsed(bl)) {
      coalesce(&bl)
      (getBlockSize(bl) >= n) && return bl
    }
    next(&bl)
  }
  if getBlockUsed(bl) {
    Block(ptr, size) = bl
    bl = Block(ptr+size, Int32(0))
  }
  size = getBlockSize(bl)
  if size < n {
    needed = n - size
    alloc = npages(needed)
    memoryGrow!(alloc)
    size = size + alloc*pageSize
    setBlockSize!(&bl, size)
  }
  return bl
}

fn trim(&bl: Block(ptr, size), newsize) {
  (size <= (newsize + headerSize)) && return
  setBlockSize!(&bl, newsize)
  newptr = ptr + newsize
  setBlockUsed!(newptr, false)
  setBlockSize!(newptr, size-newsize)
  return
}

# Public interface

fn malloc!(n: Int32) {
  n = n + headerSize
  bl = findFreeBlock(n)
  trim(&bl, n)
  setBlockUsed!(bl, true)
  return publicPtr(bl)
}

fn free!(p: Ptr) {
  p = p - headerSize
  setBlockUsed!(p, false)
  return
}

fn blockSize(p: Ptr) {
  getBlockSize(p - headerSize) - headerSize
}

fn blockCount(p: Ptr) {
  Int64(getBlockUsed(p - headerSize))
}

fn blockUnique(p: Ptr) {
  blockCount(p) == 1
}

fn retain!(p: Ptr) {
  p = p - headerSize
  n = getBlockUsed(p)
  if n == Int32(0) {
    abort("Memory management fault: retain")
  }
  setBlockUsed!(p, n+Int32(1))
  return
}

fn release!(p: Ptr, f: Int32) {
  p = p - headerSize
  n = getBlockUsed(p)
  if n == Int32(0) {
    abort("Memory management fault: release")
  } else if (n == Int32(1)) && (f != Int32(-1)) {
    invoke(Function(f, [TPtr], TNil), (p + headerSize) + 8)
  }
  # Setting to 0 equates to freeing
  setBlockUsed!(p, n-Int32(1))
}

fn release!(p: Ptr) { release!(p, widen(Int32(-1))) }

fn allocationCount() {
  i = 0
  bl = firstBlock()
  while true {
    getBlockUsed(bl) && (i = i + 1)
    isLastBlock(bl) && break
    next(&bl)
  }
  return i
}

fn checkAllocations() {
  if allocationCount() != 0 {
    println("Warning: retained memory")
  }
}