/*++

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

Module Name:

   mmmd.c

Abstract:

   Provides memory descriptor routines.

--*/

#include <ntrtl.h>
#include "bootmgr.h"
#include "mm.h"

#define MAX_STATIC_DESCRIPTOR_COUNT 1024

MEMORY_DESCRIPTOR MmStaticMemoryDescriptors[MAX_STATIC_DESCRIPTOR_COUNT];

PMEMORY_DESCRIPTOR MmGlobalMemoryDescriptors;
ULONG MmGlobalMemoryDescriptorCount;

PMEMORY_DESCRIPTOR MmDynamicMemoryDescriptors;
ULONG MmDynamicMemoryDescriptorCount;

VOID
MmMdRemoveDescriptorFromList (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN PMEMORY_DESCRIPTOR Descriptor
    )

/*++

Routine Description:

    Removes a descriptor from a MDL.

Arguments:

    Mdl - the MDL to remove Descriptor from.

    Descriptor - the descriptor to remove from Mdl.

Return Value:

    None.

--*/

{
    PLIST_ENTRY Blink;

    Blink = Descriptor->ListEntry.Blink;

    //
    // Remove the descriptor from the MDL.
    //
    RemoveEntryList(&Descriptor->ListEntry);

    //
    // Check if the removed descriptor was cached.
    //
    if (Mdl->Current != &Descriptor->ListEntry) {
        return;
    }

    //
    // Cache the previous descriptor if possible.
    //
    if (
        (
            (ULONG_PTR)Blink < (ULONG_PTR)MmGlobalMemoryDescriptors
            || (ULONG_PTR)Blink >= (ULONG_PTR)&MmGlobalMemoryDescriptors[MmGlobalMemoryDescriptorCount]
        )
        && Blink != Mdl->Head
    ) {
        Mdl->Current = Blink;
    } else {
        Mdl->Current = NULL;
    }
}

NTSTATUS
MmMdFreeDescriptor (
    IN PMEMORY_DESCRIPTOR Descriptor
    )

/*++

Routine Description:

    Frees a memory descriptor.

Arguments:

    Descriptor - the descriptor to free.

Return Value:

    STATUS_SUCCESS if successful,
    other NTSTATUS value if an error occurs.

--*/

{
    if (
        (
            MmDynamicMemoryDescriptors != NULL
            && (ULONG_PTR)Descriptor >= (ULONG_PTR)MmDynamicMemoryDescriptors
            && (ULONG_PTR)Descriptor <= (ULONG_PTR)&MmDynamicMemoryDescriptors[MmDynamicMemoryDescriptorCount]
        )
        || (
            (ULONG_PTR)Descriptor >= (ULONG_PTR)MmStaticMemoryDescriptors
            && (ULONG_PTR)Descriptor <= (ULONG_PTR)&MmStaticMemoryDescriptors[MAX_STATIC_DESCRIPTOR_COUNT]
        )
    ) {
        //
        // Clear the descriptor from static/dynamic MDL.
        //
        RtlZeroMemory(Descriptor, sizeof(MEMORY_DESCRIPTOR));
        return STATUS_SUCCESS;
    }

    //
    // Free the descriptor from the heap.
    // TODO: Use BlMmFreeHeap()
    //
    ConsolePrint(L"Cannot free memory descriptor: BlMmFreeHeap() is not implemented\r\n");
    return STATUS_NOT_IMPLEMENTED;
    // return BlMmFreeHeap(Descriptor);
}

VOID
MmMdFreeList (
    IN PMEMORY_DESCRIPTOR_LIST Mdl
    )

/*++

Routine Description:

    Frees a memory descriptor list (MDL).

Arguments:

    Mdl - the MDL to free.

Return Value:

    None.

--*/

{
    PLIST_ENTRY Entry;

    Entry = Mdl->Head->Flink;
    while (Entry != Mdl->Head) {
        MmMdRemoveDescriptorFromList(Mdl, (PMEMORY_DESCRIPTOR)Entry);
        MmMdFreeDescriptor((PMEMORY_DESCRIPTOR)Entry);
        Entry = Entry->Flink;
    }
}

VOID
MmMdInitialize (
    IN ULONG Unused,
    IN PBOOT_LIBRARY_PARAMETERS LibraryParameters
    )

/*++

Routine Description:

    Initializes the memory descriptor manager.

Arguments:

    Unused - Ignored.

    LibraryParameters - pointer to the library parameters structure.

Return Value:

    None.

--*/

{
    (VOID)Unused;
    (VOID)LibraryParameters;

    DebugPrint(L"Initializing memory descriptor manager...\r\n");

    //
    // Initialize global memory descriptor list.
    //
    MmGlobalMemoryDescriptors = &MmStaticMemoryDescriptors[0];
    MmGlobalMemoryDescriptorCount = MAX_STATIC_DESCRIPTOR_COUNT;
    RtlZeroMemory(MmGlobalMemoryDescriptors, MAX_STATIC_DESCRIPTOR_COUNT * sizeof(MEMORY_DESCRIPTOR));
    DebugPrintf(L"Global memory descriptor count: %x\r\n", MmGlobalMemoryDescriptorCount);
}