GitFOSS

module PMM
  PAGE_SIZE = 4096
  KERNEL_HEAP_START = 0x200000_u64
  KERNEL_HEAP_SIZE = 16 * 1024 * 1024

  TOTAL_PAGES = 4096 # KERNEL_HEAP_SIZE / PAGE_SIZE

  @@page_bitmap = StaticArray(UInt64, TOTAL_PAGES).new()
  @@heap_ptr = Pointer(UInt8).new(KERNEL_HEAP_START)
  @@heap_end = Pointer(UInt8).new(KERNEL_HEAP_START + KERNEL_HEAP_SIZE)

  def self.init : Nil
    TOTAL_PAGES.times do |i|
      clear_page_used(i.to_u64)
    end
  end

  private def self.set_page_used(page : UInt64) : Nil
    @@page_bitmap[page // 64] |= (1_u64 << (page % 64))
  end

  private def self.clear_page_used(page : UInt64) : Nil
    @@page_bitmap[page // 64] &= ~(1_u64 << (page % 64))
  end

  private def self.is_page_used(page : UInt64) : Bool
    ((@@page_bitmap[page // 64] >> (page % 64)) & 1_u64) == 1_u64
  end

  def self.alloc_page : Pointer(Void)?
    TOTAL_PAGES.times do |i|
      unless is_page_used(i)
        set_page_used(i)
        return Pointer(Void).new(KERNEL_HEAP_START + i * PAGE_SIZE)
      end
    end
    nil
  end

  def self.free_page(page : Pointer(Void)) : Nil
    addr = page.address
    return if addr < KERNEL_HEAP_START || addr >= KERNEL_HEAP_START + KERNEL_HEAP_SIZE
    page_idx = (addr - KERNEL_HEAP_START) // PAGE_SIZE
    clear_page_used(page_idx)
  end

  def self.kalloc(bytes : UInt64) : Pointer(Void)?
    size = (bytes + 7) & ~7
    ret = @@heap_ptr
    if @@heap_ptr + size > @@heap_end
      return nil
    end
    @@heap_ptr += size
    ret
  end
end