genusOS/src/alloc.c

#include "alloc.h"
#include "../include/limine.h"
#include "common.h"
#include "print.h"
#include "lib/stdio.h"

#ifdef DBG
  #define dprint(x) print(x)
  #define dprint8(x) print8(x)
  #define dprint16(x) print16(x)
  #define dprint32(x) print32(x)
  #define dprint64(x) print64(x)
#else
  #define dprint(x)
  #define dprint8(x)
  #define dprint16(x)
  #define dprint32(x)
  #define dprint64(x)
#endif
__attribute__((
    used,
    section(".limine_requests"))) static volatile struct limine_memmap_request
    mem_req = {
        .id = LIMINE_MEMMAP_REQUEST_ID,
        .revision = 0,
};

extern u32 fg, bg;

static void print_memmap_entry(struct limine_memmap_entry *entry) {
  u32 savefg = fg;
  u32 savebg = bg;
  static const char *type_strs[] = {
      "Usable",
      "Reserved",
      "ACPI Reclaimable",
      "ACPI NVS",
      "BAD!!!",
      "Bootloader Reclaimable",
      "Executable and Modules",
      "Framebuffer",
      "ACPI Tables",
  };

  static u32 type_colors[] = {0x00AF00, 0xAF1F00, 0x7f7f00, 0x7f7f00, 0xFF0000,
                              0x007f7f, 0x408000, 0x0000FF, 0x00FF00};

  bg = type_colors[entry->type];
  fg = ~bg;

  u64 addr = entry->base;
  u64 end = addr + entry->length;
  print("[");
  print(type_strs[entry->type]);
  print("]");
  print64(addr);
  print(" - ");
  print64(end);
  print("\n");
  fg = savefg;
  bg = savebg;
}

static size_t total_pages;
static u8 *alloc_start;
static u8 *alloc_end;
static u8 *alloc_next;
static u64 *alloc_bitmap;

void alloc_init() {
  print("Initializing Page Allocator\n");

  if (!mem_req.response) {
    print("Failed to get memory map :(\n)");
    hang();
  }

  u64 largest_base = 0;
  size_t largest_size = 0;

  for (size_t i = 0; i < mem_req.response->entry_count; i++) {
    struct limine_memmap_entry *entry = mem_req.response->entries[i];
    //print_memmap_entry(entry);

    if (entry->type == LIMINE_MEMMAP_USABLE) {
      if (entry->length > largest_size) {
        largest_base = entry->base;
        largest_size = entry->length;
        print("\t\\LARGEST USABLE: ");
        print64(entry->length >> 20);
        print("MiB\n");
      }
    }
  }

  if (!largest_size) {
    print("Couldn't find a free page!!!\n");
    hang();
  }

  total_pages = largest_size / PAGE_SIZE;
  size_t usage_bitmap_size = (total_pages + 7) / 8;
  total_pages -= (usage_bitmap_size + PAGE_SIZE - 1) / PAGE_SIZE;
  alloc_bitmap = (u64 *)(phys_to_virt(largest_base));
  alloc_start = (u8 *)(phys_to_virt(largest_base + usage_bitmap_size));
  alloc_next = alloc_start;
  alloc_end = (u8 *)(phys_to_virt(largest_base + largest_size));
}

static size_t page(void *addr) {
  return ((u8 *)addr - alloc_start) / PAGE_SIZE;
}

static void *addr(size_t page) { return alloc_start + page * PAGE_SIZE; }

static void usepage(size_t page) {
  size_t idx = page / 64;
  size_t bit = page & 63;
  alloc_bitmap[idx] |= (1ULL << bit);
}

static void usepagerange(size_t page, size_t len) {
  for (size_t i = 0; i < len; i++) {
    usepage(i + page);
  }
}

static void unusepage(size_t page) {
  size_t idx = page / 64;
  size_t bit = page & 63;
  alloc_bitmap[idx] &= ~(1ULL << bit);
}

static void unusepagerange(size_t page, size_t len) {
  for (size_t i = 0; i < len; i++) {
    unusepage(i + page);
  }
}

static u8 getpageusage(size_t page) {
  return (alloc_bitmap[page / 8] >> (page & 7)) & 1;
}

void dealloc(void *addr, size_t size) {
  // adjust to page size
  size_t pages = (size + PAGE_SIZE - 1) / PAGE_SIZE;
  unusepagerange(page(addr), pages);
}

void *alloc(size_t size) {
  size_t pages = (size + PAGE_SIZE - 1) / PAGE_SIZE;
  size_t page = 0;

  if (pages == 0 || pages >= total_pages) {
    print("Too many or too few pages: ");
    print8(pages);
    print("\n");
    return NULL;
  }


  dprint("Attempting to find ");
  dprint64(pages);
  dprint(" consecutive pages.....");
  size_t consecutive = 0;

  for (size_t i = 0; i < total_pages / 64; i++) {
    u64 word = alloc_bitmap[i];
    if (word == 0) {
      if (!consecutive)
        page = i * 64;
      consecutive += 64;
      if (consecutive >= pages) {
        usepagerange(page, pages);
        dprint("succeeded!\n");
        return addr(page);
      }
      continue;
    }

    if (word == ~0ULL) {
      consecutive = 0;
      page = -1;
    }

    for (size_t bit = 0; bit < 64; bit++) {
      if ((word & (1ULL << bit)) == 0) {
        if (!consecutive)
          page = i * 64 + bit;

        if (++consecutive >= pages) {
          usepagerange(page, pages);  
          dprint("succeeded!\n");
          
          return addr(page);
        }
      } else {
        consecutive = 0;
        page = -1;
      }
    }
  }

  dprint("failed :(!\n");

  return NULL;
}