/*
 * Copyright 2016-2022 Haiku, Inc. All rights reserved.
 * Copyright 2014, Jessica Hamilton, jessica.l.hamilton@gmail.com.
 * Copyright 2014, Henry Harrington, henry.harrington@gmail.com.
 * Distributed under the terms of the MIT License.
 */


#include <algorithm>

#include <boot/addr_range.h>
#include <boot/platform.h>
#include <boot/stage2.h>
#include <kernel/kernel.h>

#include "efi_platform.h"
#include "mmu.h"


//#define TRACE_MMU
#ifdef TRACE_MMU
#   define TRACE(x...) dprintf("efi/mmu: " x)
#else
#   define TRACE(x...) ;
#endif


struct memory_region {
        memory_region *next;
        addr_t vaddr;
        phys_addr_t paddr;
        size_t size;

        void dprint(const char * msg) {
          dprintf("%s memory_region v: %#" B_PRIxADDR " p: %#" B_PRIxPHYSADDR " size: %lu\n", msg, vaddr,
                        paddr, size);
        }

        bool matches(phys_addr_t expected_paddr, size_t expected_size) {
                return paddr == expected_paddr && size == expected_size;
        }
};


static const size_t kMaxKernelSize = 0x2000000;         // 32 MB
static addr_t sNextVirtualAddress = KERNEL_LOAD_BASE + kMaxKernelSize;

static memory_region *allocated_regions = NULL;


extern "C" phys_addr_t
mmu_allocate_page()
{
        TRACE("%s: called\n", __func__);

        efi_physical_addr addr;
        efi_status s = kBootServices->AllocatePages(AllocateAnyPages,
                EfiLoaderData, 1, &addr);

        if (s != EFI_SUCCESS)
                panic("Unabled to allocate memory: %li", s);

        return addr;
}


extern "C" addr_t
get_next_virtual_address(size_t size)
{
        TRACE("%s: called. size: %" B_PRIuSIZE "\n", __func__, size);

        addr_t address = sNextVirtualAddress;
        sNextVirtualAddress += ROUNDUP(size, B_PAGE_SIZE);
        return address;
}


extern "C" addr_t
get_current_virtual_address()
{
        TRACE("%s: called\n", __func__);
        return sNextVirtualAddress;
}


// Platform allocator.
// The bootloader assumes that bootloader address space == kernel address space.
// This is not true until just before the kernel is booted, so an ugly hack is
// used to cover the difference. platform_allocate_region allocates addresses
// in bootloader space, but can convert them to kernel space. The ELF loader
// accesses kernel memory via Mao(), and much later in the boot process,
// addresses in the kernel argument struct are converted from bootloader
// addresses to kernel addresses.

extern "C" status_t
platform_allocate_region(void **_address, size_t size, uint8 protection)
{
        TRACE("%s: called\n", __func__);

        size_t pages = ROUNDUP(size, B_PAGE_SIZE) / B_PAGE_SIZE;
        efi_physical_addr addr = 0;
        efi_status status = kBootServices->AllocatePages(AllocateAnyPages,
                EfiLoaderData, pages, &addr);

        if (status != EFI_SUCCESS)
                return B_NO_MEMORY;

        // Addresses above 512GB not supported.
        // Memory map regions above 512GB can be ignored, but if EFI returns pages
        // above that there's nothing that can be done to fix it.
        if (addr + size > (512ull * 1024 * 1024 * 1024))
                panic("Can't currently support more than 512GB of RAM!");

        memory_region *region = new(std::nothrow) memory_region {
                next: allocated_regions,
                vaddr: 0,
                paddr: (phys_addr_t)addr,
                size: size,
        };

        if (region == NULL) {
                kBootServices->FreePages(addr, pages);
                return B_NO_MEMORY;
        }

        if (*_address != NULL) {
                // This is only useful for mapping the kernel itself.
                // Validate base and size, but don't check for duplicates.
                addr_t virtualAddress = (addr_t)*_address;
                if (virtualAddress < KERNEL_LOAD_BASE
                                || (virtualAddress + size) > (KERNEL_LOAD_BASE + kMaxKernelSize)) {
                        kBootServices->FreePages(addr, pages);
                        delete region;
                        return B_BAD_VALUE;
                }

                region->vaddr = virtualAddress;
        }

#ifdef TRACE_MMU
        //region->dprint("Allocated");
#endif
        allocated_regions = region;
        *_address = (void *)region->paddr;
        return B_OK;
}


extern "C" status_t
platform_allocate_region_below(void **_address, size_t size, phys_addr_t maxAddress)
{
        TRACE("%s: called\n", __func__);

        efi_physical_addr addr = maxAddress;
        size_t pages = ROUNDUP(size, B_PAGE_SIZE) / B_PAGE_SIZE;
        efi_status status = kBootServices->AllocatePages(AllocateMaxAddress,
                EfiLoaderData, pages, &addr);
        if (status != EFI_SUCCESS)
                return B_NO_MEMORY;

        memory_region *region = new(std::nothrow) memory_region {
                next: allocated_regions,
                vaddr: 0,
                paddr: (phys_addr_t)addr,
                size: size
        };

        if (region == NULL) {
                kBootServices->FreePages(addr, pages);
                return B_NO_MEMORY;
        }

        allocated_regions = region;
        *_address = (void *)region->paddr;
        return B_OK;
}


/*!
        Neither \a virtualAddress nor \a size need to be aligned, but the function
        will map all pages the range intersects with.
        If physicalAddress is not page-aligned, the returned virtual address will
        have the same "misalignment".
*/
extern "C" addr_t
mmu_map_physical_memory(addr_t physicalAddress, size_t size, uint32 flags)
{
        TRACE("%s: called\n", __func__);

        addr_t pageOffset = physicalAddress & (B_PAGE_SIZE - 1);

        physicalAddress -= pageOffset;
        size += pageOffset;

        if (insert_physical_allocated_range(physicalAddress,
                        ROUNDUP(size, B_PAGE_SIZE)) != B_OK)
                return B_NO_MEMORY;

        return physicalAddress + pageOffset;
}


static void
convert_physical_ranges()
{
        TRACE("%s: called\n", __func__);

        addr_range *range = gKernelArgs.physical_allocated_range;
        uint32 num_ranges = gKernelArgs.num_physical_allocated_ranges;

        for (uint32 i = 0; i < num_ranges; ++i) {
                // Addresses above 512GB not supported.
                // Memory map regions above 512GB can be ignored, but if EFI returns
                // pages above that there's nothing that can be done to fix it.
                if (range[i].start + range[i].size > (512ull * 1024 * 1024 * 1024))
                        panic("Can't currently support more than 512GB of RAM!");

                memory_region *region = new(std::nothrow) memory_region {
                        next: allocated_regions,
                        vaddr: 0,
                        paddr: (phys_addr_t)range[i].start,
                        size: (size_t)range[i].size
                };

                if (!region)
                        panic("Couldn't add allocated region");

                allocated_regions = region;

                // Clear out the allocated range
                range[i].start = 0;
                range[i].size = 0;
                gKernelArgs.num_physical_allocated_ranges--;
        }
}


extern "C" status_t
platform_assign_kernel_address_for_region(void *address, addr_t assign)
{
        // Double cast needed to avoid sign extension issues on 32-bit architecture
        phys_addr_t addr = (phys_addr_t)(addr_t)address;

        for (memory_region *region = allocated_regions; region;
                        region = region->next) {
                if (region->paddr <= addr && addr < region->paddr + region->size) {
                        if (region->paddr != addr)
                                return EINVAL;
                        if (region->vaddr != 0)
                                return EALREADY;

                        region->vaddr = assign;
                        return B_OK;
                }
        }

        return B_ERROR;
}


extern "C" status_t
platform_bootloader_address_to_kernel_address(void *address, addr_t *_result)
{
        TRACE("%s: called\n", __func__);

        // Convert any physical ranges prior to looking up address
        convert_physical_ranges();

        // Double cast needed to avoid sign extension issues on 32-bit architecture
        phys_addr_t addr = (phys_addr_t)(addr_t)address;

        for (memory_region *region = allocated_regions; region;
                        region = region->next) {
                if (region->paddr <= addr && addr < region->paddr + region->size) {
                        // Lazily allocate virtual memory.
                        if (region->vaddr == 0)
                                region->vaddr = get_next_virtual_address(region->size);

                        *_result = region->vaddr + (addr - region->paddr);
                        //dprintf("Converted bootloader address %p in region %#lx-%#lx to %#lx\n",
                        //      address, region->paddr, region->paddr + region->size, *_result);
                        return B_OK;
                }
        }

        return B_ERROR;
}


extern "C" status_t
platform_kernel_address_to_bootloader_address(addr_t address, void **_result)
{
        TRACE("%s: called\n", __func__);

        for (memory_region *region = allocated_regions; region;
                        region = region->next) {
                if (region->vaddr != 0 && region->vaddr <= address
                                && address < region->vaddr + region->size) {
                        *_result = (void *)(region->paddr + (address - region->vaddr));
                        //dprintf("Converted kernel address %#lx in region %#lx-%#lx to %p\n",
                        //      address, region->vaddr, region->vaddr + region->size, *_result);
                        return B_OK;
                }
        }

        return B_ERROR;
}


extern "C" status_t
platform_free_region(void *address, size_t size)
{
        TRACE("%s: called to release region %p (%" B_PRIuSIZE ")\n", __func__,
                address, size);

        for (memory_region **ref = &allocated_regions; *ref;
                        ref = &(*ref)->next) {
                // Double cast needed to avoid sign extension issues on 32-bit architecture
                if ((*ref)->matches((phys_addr_t)(addr_t)address, size)) {
                        efi_status status;
                        status = kBootServices->FreePages((efi_physical_addr)(addr_t)address,
                                ROUNDUP(size, B_PAGE_SIZE) / B_PAGE_SIZE);
                        ASSERT_ALWAYS(status == EFI_SUCCESS);
                        memory_region* old = *ref;
                        //pointer to current allocated_memory_region* now points to next
                        *ref = (*ref)->next;
#ifdef TRACE_MMU
                        old->dprint("Freeing");
#endif
                        delete old;
                        return B_OK;
                }
        }
        panic("platform_free_region: Unknown region to free??");
        return B_ERROR; // NOT Reached
}


bool
mmu_next_region(void** cookie, addr_t* vaddr, phys_addr_t* paddr, size_t* size)
{
        if (*cookie == NULL)
                *cookie = allocated_regions;
        else
                *cookie = ((memory_region*)*cookie)->next;

        memory_region* region = (memory_region*)*cookie;
        if (region == NULL)
                return false;

        if (region->vaddr == 0)
                region->vaddr = get_next_virtual_address(region->size);

        *vaddr = region->vaddr;
        *paddr = region->paddr;
        *size = region->size;
        return true;
}