/*++

Copyright (c) 2024, Quinn Stephens.
Provided under the BSD 3-Clause license.

Module Name:

   mmpa.c

Abstract:

   Provides memory manager page routines.

--*/

#include "bootlib.h"
#include "mm.h"

ULONG PapMinimumAllocationCount;
ULONGLONG PapMinimumPhysicalPage;
ULONGLONG PapMaximumPhysicalPage;

MEMORY_DESCRIPTOR_LIST MmMdlFwAllocationTracker;
MEMORY_DESCRIPTOR_LIST MmMdlBadMemory;
MEMORY_DESCRIPTOR_LIST MmMdlTruncatedMemory;
MEMORY_DESCRIPTOR_LIST MmMdlPersistentMemory;
MEMORY_DESCRIPTOR_LIST MmMdlReservedAllocated;
MEMORY_DESCRIPTOR_LIST MmMdlMappedAllocated;
MEMORY_DESCRIPTOR_LIST MmMdlMappedUnallocated;
MEMORY_DESCRIPTOR_LIST MmMdlUnmappedAllocated;
MEMORY_DESCRIPTOR_LIST MmMdlUnmappedUnallocated;

FORCEINLINE
VOID
InitializeMdl (
    IN PMEMORY_DESCRIPTOR_LIST Mdl
    )

/*++

Routine Description:

    Initializes a MDL.

Arguments:

    Mdl - Pointer to the MDL.

Return Value:

    None.

--*/

{
    Mdl->Head = NULL;
    Mdl->Current = NULL;
    Mdl->Type = MDL_TYPE_PHYSICAL;
}

NTSTATUS
MmPaInitialize (
    IN PBOOT_MEMORY_INFO MemoryInfo,
    IN ULONG MinimumAllocation
    )

/*++

Routine Description:

    Initializes the page allocator.

Arguments:

    MemoryInfo - Pointer to the memory info structure.

    MinimumAllocation - Minimum amount of pages to grow the pool by at a time.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if regions in MemoryInfo could not be removed.
    STATUS_NO_MEMORY if a new descriptor cannot be allocated.

--*/

{
    NTSTATUS Status;
    PMEMORY_DESCRIPTOR Descriptor, NewDescriptor;

    (VOID)MemoryInfo;

    DebugPrint(L"Initializing page allocator...\r\n");

    PapMinimumAllocationCount = MinimumAllocation;
    PapMinimumPhysicalPage = 1;
    PapMaximumPhysicalPage = MAXULONGLONG >> PAGE_SHIFT;

    InitializeMdl(&MmMdlFwAllocationTracker);
    InitializeMdl(&MmMdlBadMemory);
    InitializeMdl(&MmMdlTruncatedMemory);
    InitializeMdl(&MmMdlPersistentMemory);
    InitializeMdl(&MmMdlReservedAllocated);
    InitializeMdl(&MmMdlMappedAllocated);
    InitializeMdl(&MmMdlMappedUnallocated);
    InitializeMdl(&MmMdlUnmappedAllocated);
    InitializeMdl(&MmMdlUnmappedUnallocated);

    Status = MmFwGetMemoryMap(&MmMdlUnmappedUnallocated, 0x03);
    if (!NT_SUCCESS(Status)) {
        return Status;
    }

    //
    // Remove regions described in MemoryInfo from the
    // MDL the memory manager will use for allocation.
    //
    Descriptor = (PMEMORY_DESCRIPTOR)((PUCHAR)MemoryInfo + MemoryInfo->MdlOffset);
    for (ULONG DescriptorCount = MemoryInfo->DescriptorCount; DescriptorCount > 0; DescriptorCount--) {
        //
        // Remove from the usable MDL.
        //
        Status = MmMdRemoveRegionFromMdlEx(&MmMdlUnmappedUnallocated, Descriptor->FirstPage, Descriptor->PageCount, MDL_OPERATION_FLAGS_PHYSICAL, NULL);
        if (!NT_SUCCESS(Status)) {
            return STATUS_INVALID_PARAMETER;
        }

        //
        // ... and add to the reserved MDL.
        //
        NewDescriptor = MmMdInitDescriptor(
            Descriptor->FirstPage,
            Descriptor->MappedFirstPage,
            Descriptor->PageCount,
            Descriptor->Attributes,
            Descriptor->Type
        );
        if (NewDescriptor == NULL) {
            return STATUS_NO_MEMORY;
        }

        Status = MmMdAddDescriptorToList(&MmMdlReservedAllocated, NewDescriptor, 0x00);
        if (!NT_SUCCESS(Status)) {
            return Status;
        }

        Descriptor = (PMEMORY_DESCRIPTOR)((PUCHAR)Descriptor + MemoryInfo->DescriptorSize);
    }

    Status = BlpMmInitializeConstraints();
    if (!NT_SUCCESS(Status)) {
        return Status;
    }

    return STATUS_SUCCESS;
}