k-alloc.cc

#include "kernel.hh"
#include "k-lock.hh"
#include "k-list.hh"

constexpr int min_order = 12;
constexpr int max_order = 22;

static spinlock page_lock;
size_t total_pages;
size_t allocated_pages;

struct PageInfo {
    uint64_t addr;
    bool is_free = false;
    int order = min_order;
    list_links link_;
};
PageInfo physical_pages[MEMSIZE_PHYSICAL/PAGESIZE];

list<PageInfo, &PageInfo::link_> free_lists[max_order - min_order + 1];

void kalloc_coalesce(PageInfo* block);

// init_kalloc
//    Initialize stuff needed by `kalloc`. Called from `init_hardware`,
//    after `physical_ranges` is initialized.
void init_kalloc() {
    total_pages = 0;
    auto irqs = page_lock.lock();
    auto range = physical_ranges.find(0);
    while (range != physical_ranges.end()) {
        if (range->type() == mem_available) {
            // Skip current range if less than a single page
            if (range->size() < PAGESIZE) {
                ++range;
                continue;
            }
            // Initialize the PageInfo structure for each available page
            for (uint64_t addr = range->first(); addr < range->last(); addr += PAGESIZE) {
                // Skip to the first full free page if necessary
                if (addr % PAGESIZE != 0) {
                    addr += PAGESIZE - (addr % PAGESIZE);
                }
                assert(addr % PAGESIZE == 0);
                assert(addr < MEMSIZE_PHYSICAL);

                physical_pages[addr/PAGESIZE].addr = addr;
                physical_pages[addr/PAGESIZE].is_free = true;
                physical_pages[addr/PAGESIZE].order = min_order;
                free_lists[min_order - PAGEOFFBITS].push_back(&physical_pages[addr/PAGESIZE]);
                ++total_pages;

                kalloc_coalesce(&physical_pages[addr/PAGESIZE]);
            }
        }
        ++range;
    }
    assert(total_pages <= MEMSIZE_PHYSICAL/PAGESIZE);

    allocated_pages = 0;
    page_lock.unlock(irqs);
}


// kalloc(sz)
//    Allocate and return a pointer to at least `sz` contiguous bytes of
//    memory. Returns `nullptr` if `sz == 0` or on failure.
//
//    The caller should initialize the returned memory before using it.
//    The handout allocator sets returned memory to 0xCC (this corresponds
//    to the x86 `int3` instruction and may help you debug).
//
//    If `sz` is a multiple of `PAGESIZE`, the returned pointer is guaranteed
//    to be page-aligned.
//
//    The handout code does not free memory and allocates memory in units
//    of pages.
void* kalloc(size_t sz) {
    // Calculate the lowest order needed to contain sz bytes
    int needed_order = msb(sz - 1);
    // log_printf("sz %d needs order %d\n", sz, needed_order);
    if (sz == 0 || needed_order > max_order) {
        return nullptr;
    }
    assert((size_t(1) << needed_order) >= sz);

    // Clamp needed_order to be at least min_order
    if (needed_order < min_order) {
        needed_order = min_order;
    }

    auto irqs = page_lock.lock();
    void* ptr = nullptr;

    // Find a free physical page of the correct order (current only supports single pages)
    PageInfo* found_base;
    int found_order = needed_order;
    for (; found_order <= max_order; ++found_order) {
        found_base = free_lists[found_order - min_order].pop_front();
        if (found_base) {
            // log_printf("found order %d looking for %d\n", found_order, needed_order);
            assert(found_base->is_free);
            assert(found_base->order == found_order);
            // log_printf("found addr %p of order %d\n", found_base->addr, found_order);
            assert(found_base->addr != 0 && found_base->addr % (1 << (found_order - 1)) == 0);
            
            // Split the block if necessary
            while (found_order > needed_order) {
                --found_order;
                // Create a new block for second buddy
                uint64_t buddy_addr = found_base->addr + (1 << found_order);
                physical_pages[buddy_addr/PAGESIZE].order = found_order;
                physical_pages[buddy_addr/PAGESIZE].is_free = true;
                free_lists[found_order - min_order].push_back(&physical_pages[buddy_addr/PAGESIZE]);
                // log_printf("Created new buddy of order %d at %p (from %p)\n", found_order, buddy_addr, found_base->addr);
                // Update our block's order
                found_base->order = found_order;
            }

            assert(found_order == needed_order);
            assert(found_base->order == needed_order);

            assert(!found_base->link_.is_linked());
            found_base->is_free = false;
            allocated_pages += (1 << (found_order - min_order));
            ptr = pa2kptr<void*>(found_base->addr);
            break;
        }
    }

    page_lock.unlock(irqs);

    if (ptr) {
        // tell sanitizers the allocated page is accessible
        asan_mark_memory(ka2pa(ptr), 1 << found_order, false);
        // initialize to `int3`
        memset(ptr, 0xCC, 1 << found_order);

        // log_printf("got ptr %p-%p\n", ptr, reinterpret_cast<uint64_t>(ptr) + (1 << found_order) - 1);
        // assert(reinterpret_cast<uint64_t>(ptr) < -HIGHMEM_BASE);
    }
    return ptr;
}


// kfree(ptr)
//    Free a pointer previously returned by `kalloc`. Does nothing if
//    `ptr == nullptr`.
void kfree(void* ptr) {
    if (ptr == nullptr) {
        return;
    }
    uint64_t addr = kptr2pa<void>(ptr);

    // Address must be a physical page and properly aligned
    assert(addr < MEMSIZE_PHYSICAL);
    assert(addr % PAGESIZE == 0);

    auto irqs = page_lock.lock();
    PageInfo* page = &physical_pages[addr/PAGESIZE];

    // Page should not be free already (double free)
    assert(!page->is_free);

    // Mark the page as free and tell sanitizers it is inaccessible
    page->is_free = true;
    asan_mark_memory(addr, 1 << page->order, true);

    // Return the page to its corresponding free list
    free_lists[page->order - min_order].push_front(page);
    allocated_pages -= 1 << (page->order - min_order);

    // Try to coalesce adjacent pages into the highest possible order
    kalloc_coalesce(page);

    page_lock.unlock(irqs);
}


// kalloc_coalesce(block)
//    Coalesces the given block with its neighbors to the largest
//    order possible. Should hold the page_lock before calling

void kalloc_coalesce(PageInfo* block) {
    PageInfo* base = block;
    assert(base->is_free);
    assert(base->addr != 0);
    base->link_.erase();
    while (base->order < max_order) {
        uint64_t buddy_addr = base->addr ^ (1 << base->order);
        assert(buddy_addr != base->addr);
        assert(buddy_addr < MEMSIZE_PHYSICAL);
        PageInfo* buddy = &physical_pages[buddy_addr/PAGESIZE];
        if (buddy_addr == 0 || buddy->addr != buddy_addr || !buddy->is_free || buddy->order != base->order) {
            // log_printf("buddy %p not suitable for %p (order %d) (free %d) (b order %d)\n", buddy_addr, base->addr, base->order, buddy->is_free, buddy->order);
            break;
        }
        // log_printf("found buddy %p for %p (order %d)\n", buddy_addr, base->addr, base->order);
        buddy->link_.erase();

        // Assign the lower block to be the base
        if (buddy_addr < base->addr) {
            // log_printf("swapping base\n");
            PageInfo* tmp = base;
            base = buddy;
            buddy = tmp;
        }

        base->is_free = true;
        buddy->is_free = false;
        ++base->order;
        ++buddy->order;
        assert(base->order == buddy->order);

        // log_printf("now have base %p order %d free %d\n", base->addr, base->order, base->is_free);
        // log_printf("    and buddy %p order %d free %d\n", buddy->addr, buddy->order, buddy->is_free);
    }
    free_lists[base->order - min_order].push_front(base);
}


// operator new, operator delete
//    Expressions like `new (std::nothrow) T(...)` and `delete x` work,
//    and call kalloc/kfree.
void* operator new(size_t sz, const std::nothrow_t&) noexcept {
    return kalloc(sz);
}
void* operator new(size_t sz, std::align_val_t, const std::nothrow_t&) noexcept {
    return kalloc(sz);
}
void* operator new[](size_t sz, const std::nothrow_t&) noexcept {
    return kalloc(sz);
}
void* operator new[](size_t sz, std::align_val_t, const std::nothrow_t&) noexcept {
    return kalloc(sz);
}
void operator delete(void* ptr) noexcept {
    kfree(ptr);
}
void operator delete(void* ptr, size_t) noexcept {
    kfree(ptr);
}
void operator delete(void* ptr, std::align_val_t) noexcept {
    kfree(ptr);
}
void operator delete(void* ptr, size_t, std::align_val_t) noexcept {
    kfree(ptr);
}
void operator delete[](void* ptr) noexcept {
    kfree(ptr);
}
void operator delete[](void* ptr, size_t) noexcept {
    kfree(ptr);
}
void operator delete[](void* ptr, std::align_val_t) noexcept {
    kfree(ptr);
}
void operator delete[](void* ptr, size_t, std::align_val_t) noexcept {
    kfree(ptr);
}