/*++

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, MmGlobalMemoryDescriptorsUsed;

PMEMORY_DESCRIPTOR MmDynamicMemoryDescriptors;
ULONG MmDynamicMemoryDescriptorCount, MmDynamicMemoryDescriptorsUsed;

BOOLEAN
MmMdpHasPrecedence (
    IN MEMORY_TYPE A,
    IN MEMORY_TYPE B
    )

/*++

Routine Description:

    Compares two memory types to determine which has precedence.

Arguments:

    A - memory type A.

    B - memory type B.

Return Value:

    TRUE if A has precedence over B, or if neither has precedence.
    FALSE if B has precedence over A.

--*/

{
    ULONG ClassA, ClassB;

    ClassA = A >> 28;
    ClassB = B >> 28;

    if (B == MEMORY_TYPE_FREE) {
        return TRUE;
    }

    if (A == MEMORY_TYPE_FREE) {
        return FALSE;
    }

    if (ClassA > 0x0F) {
        return TRUE;
    }

    if (ClassB > 0x0F) {
        return FALSE;
    }

    if (ClassA != 0x0F) {
        return (ClassB != 0x0F && (ClassA == 0x0D || ClassB != 0x0D)) ? TRUE:FALSE;
    }

    if (ClassB != 0x0F) {
        return TRUE;
    }

    //
    // TODO: Implement the rest of this routine.
    //
    ConsolePrint(L"MmMdpHasPrecedence() is incomplete\r\n");
    return TRUE;
}

BOOLEAN
MmMdpTruncateDescriptor (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN PMEMORY_DESCRIPTOR Descriptor,
    IN ULONG Flags
    )

/*++

Routine Description:

    Trunactes a memory descriptor that overlaps
    with adjacent descriptors.

Arguments:

    Mdl - The MDL containing Descriptor.

    Descriptor - The descriptor to truncate.

Return Value:

    TRUE if Descriptor was deleted from Mdl.

--*/

{
    PMEMORY_DESCRIPTOR PrevDescriptor, NextDescriptor, NewDescriptor;
    ULONGLONG DescriptorEnd, PrevDescriptorEnd, NextDescriptorEnd;
    ULONGLONG MappedFirstPage;
    ULONGLONG Change;

    PrevDescriptor = (PMEMORY_DESCRIPTOR)Descriptor->ListEntry.Blink;
    NextDescriptor = (PMEMORY_DESCRIPTOR)Descriptor->ListEntry.Flink;

    DescriptorEnd = Descriptor->FirstPage + Descriptor->PageCount;
    PrevDescriptorEnd = PrevDescriptor->FirstPage + PrevDescriptor->PageCount;
    NextDescriptorEnd = NextDescriptor->FirstPage + NextDescriptor->PageCount;

    //
    // Check if overlapping with previous descriptor.
    //
    if (Descriptor->ListEntry.Blink != Mdl->Head && Descriptor->FirstPage < PrevDescriptorEnd) {
        if (MmMdpHasPrecedence(Descriptor->Type, PrevDescriptor->Type)) {
            PrevDescriptor->PageCount = Descriptor->FirstPage - PrevDescriptor->FirstPage;

            if (DescriptorEnd < PrevDescriptorEnd) {
                if (PrevDescriptor->MappedFirstPage) {
                    MappedFirstPage = PrevDescriptor->MappedFirstPage + DescriptorEnd - PrevDescriptor->FirstPage;
                } else {
                    MappedFirstPage = 0;
                }

                NewDescriptor = MmMdInitDescriptor(
                    DescriptorEnd,
                    MappedFirstPage,
                    PrevDescriptorEnd - DescriptorEnd,
                    PrevDescriptor->Attributes,
                    PrevDescriptor->Type
                );
                if (NewDescriptor != NULL) {
                    MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
                }
            }

            //
            // Delete and free the previous
            // descriptor if it is now empty.
            //
            if (PrevDescriptor->PageCount == 0) {
                MmMdRemoveDescriptorFromList(Mdl, PrevDescriptor);
                MmMdFreeDescriptor(PrevDescriptor);
            }
        } else {
            //
            // Remove if completely overlapping.
            //
            if (DescriptorEnd <= PrevDescriptorEnd) {
                MmMdRemoveDescriptorFromList(Mdl, Descriptor);
                MmMdFreeDescriptor(Descriptor);
                return TRUE;
            }

            //
            // Otherwise, move the descriptor.
            //
            Change = PrevDescriptorEnd - Descriptor->FirstPage;
            Descriptor->FirstPage += Change;
            Descriptor->PageCount -= Change;
            if (Descriptor->MappedFirstPage) {
                Descriptor->MappedFirstPage += Change;
            }
        }
    }

    //
    // Check if overlapping with next descriptor.
    //
    if (Descriptor->ListEntry.Flink != Mdl->Head && NextDescriptor->FirstPage < DescriptorEnd) {
        if (MmMdpHasPrecedence(NextDescriptor->Type, Descriptor->Type)) {
            Descriptor->PageCount = NextDescriptor->FirstPage - Descriptor->FirstPage;

            if (NextDescriptorEnd < DescriptorEnd) {
                if (Descriptor->MappedFirstPage) {
                    MappedFirstPage = Descriptor->MappedFirstPage + NextDescriptorEnd - Descriptor->FirstPage;
                } else {
                    MappedFirstPage = 0;
                }

                NewDescriptor = MmMdInitDescriptor(
                    NextDescriptorEnd,
                    MappedFirstPage,
                    DescriptorEnd - NextDescriptorEnd,
                    Descriptor->Attributes,
                    Descriptor->Type
                );
                if (NewDescriptor != NULL) {
                    MmMdAddDescriptorToList(Mdl, NewDescriptor, Flags);
                }
            }

            //
            // Delete and free the previous
            // descriptor if it is now empty.
            //
            if (Descriptor->PageCount == 0) {
                MmMdRemoveDescriptorFromList(Mdl, Descriptor);
                MmMdFreeDescriptor(Descriptor);
            }
        } else {
            //
            // Remove if completely overlapping.
            //
            if (NextDescriptorEnd <= DescriptorEnd) {
                MmMdRemoveDescriptorFromList(Mdl, NextDescriptor);
                MmMdFreeDescriptor(NextDescriptor);
                return TRUE;
            }

            //
            // Otherwise, move the next descriptor.
            //
            Change = DescriptorEnd - NextDescriptor->FirstPage;
            NextDescriptor->FirstPage += Change;
            NextDescriptor->PageCount -= Change;
            if (NextDescriptor->MappedFirstPage) {
                NextDescriptor->MappedFirstPage += Change;
            }
        }
    }

    return FALSE;
}

NTSTATUS
MmMdAddDescriptorToList (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN PMEMORY_DESCRIPTOR Descriptor,
    IN ULONG Flags
    )

/*++

Routine Description:

    Adds a descriptor to a MDL.

Arguments:

    Mdl - the MDL to add Descriptor to.

    Descriptor - the descriptor to add to Mdl.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if Mdl or Descriptor are invalid.

--*/

{
    PLIST_ENTRY Entry;
    PMEMORY_DESCRIPTOR CurrentDescriptor;

    if (Mdl == NULL || Descriptor == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    if (Mdl->Current) {
        if (Descriptor->FirstPage < ((PMEMORY_DESCRIPTOR)Mdl->Current)->FirstPage) {
            Entry = Mdl->Head->Flink;
        } else {
            Entry = Mdl->Current;
        }
    } else {
        Entry = Mdl->Head->Flink;
    }

    //
    // Search for an existing descriptor for this region.
    //
    while (Entry != Mdl->Head) {
        CurrentDescriptor = (PMEMORY_DESCRIPTOR)Entry;

        if (
            Descriptor->FirstPage >= CurrentDescriptor->FirstPage && (Descriptor->FirstPage != CurrentDescriptor->FirstPage
            || !MmMdpHasPrecedence(Descriptor->Type, CurrentDescriptor->Type))
        ) {
            Entry = Entry->Flink;
            continue;
        }

        //
        // Insert descriptor into the list
        // right before the current entry.
        //
        Descriptor->ListEntry.Blink = Entry->Blink;
        Descriptor->ListEntry.Flink = Entry;
        Entry->Blink->Flink = &Descriptor->ListEntry;
        Entry->Blink = &Descriptor->ListEntry;

        //
        // Truncate overlapping descriptors
        // into one larger descriptor.
        //
        if (Flags & MDL_OPERATION_FLAGS_TRUNCATE) {
            MmMdpTruncateDescriptor(Mdl, Descriptor, Flags);
        }

        return STATUS_SUCCESS;
    }

    //
    // If there are no existing descriptors for this region
    // (or the list is empty), insert and truncate the descriptor.
    //
    InsertTailList(Mdl->Head, &Descriptor->ListEntry);
    if (Flags & MDL_OPERATION_FLAGS_TRUNCATE) {
        MmMdpTruncateDescriptor(Mdl, Descriptor, Flags);
    }

    return STATUS_SUCCESS;
}

VOID
MmMdRemoveDescriptorFromList (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN PMEMORY_DESCRIPTOR Descriptor
    )

/*++

Routine Description:

    Removes a descriptor from a MDL.

Arguments:

    Mdl - MDL to remove Descriptor from.

    Descriptor - Descriptor to remove from Mdl.

Return Value:

    None.

--*/

{
    PLIST_ENTRY Blink;

    Blink = Descriptor->ListEntry.Blink;
    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;
    }
}

PMEMORY_DESCRIPTOR
MmMdFindDescriptorFromMdl (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN ULONGLONG Page,
    IN ULONG Flags
    )

/*++

Routine Description:

    Searches an MDL for the descriptor containing a page.

Arguments:

    Mdl - The MDL to search.

    Page - The page to search for.

    Flags - MDL_OPERATION_FLAGS_*
        MDL_OPERATION_FLAGS_PHYSICAL if Page is physical.
        MDL_OPERATION_FLAGS_VIRTUAL if Page is virtual.

Return Value:

    Pointer to the descriptor if successful.
    NULL if an error occurs.

--*/


{
    //
    // TODO: Implement this routine.
    //

    return NULL;
}

NTSTATUS
MmMdRemoveRegionFromMdlEx (
    IN PMEMORY_DESCRIPTOR_LIST Mdl,
    IN ULONGLONG FirstPage,
    IN ULONGLONG PageCount,
    IN ULONG Flags,
    OUT PMEMORY_DESCRIPTOR_LIST Unused
    )

/*++

Routine Description:

    Removes a region from a MDL.

Arguments:

    Mdl - MDL to remove the region from.

    FirstPage - The first page in the region.

    PageCount - The number of pages in the region.

    Flags - MDL_OPERATION_FLAGS_*.

    Unused - Unused.

Return Value:

    None.

--*/

{
    ULONGLONG RemoveEnd, DescriptorEnd;
    PLIST_ENTRY Entry;
    PMEMORY_DESCRIPTOR Descriptor;
    MEMORY_DESCRIPTOR RemovedDescriptor;

    (VOID)Flags;
    (VOID)Unused;

    RemoveEnd = FirstPage + PageCount;
    Entry = Mdl->Head->Flink;
    while (Entry != Mdl->Head) {
        Descriptor = (PMEMORY_DESCRIPTOR)Entry;
        DescriptorEnd = Descriptor->FirstPage + Descriptor->PageCount;

        RtlCopyMemory(&RemovedDescriptor, Descriptor, sizeof(MEMORY_DESCRIPTOR));

        // if (FirstPage <= Descriptor->FirstPage && Descriptor->FirstPage < RemoveEnd) {
        // }
    }

    //
    // TODO: Implement the rest of this routine.
    //

    return STATUS_SUCCESS;
}

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"MmMdFreeDescriptor(): need BlMmFreeHeap() to free descriptor\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.

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;
    }
}

PMEMORY_DESCRIPTOR
MmMdInitDescriptor (
    IN ULONGLONG FirstPage,
    IN ULONGLONG MappedFirstPage,
    IN ULONGLONG PageCount,
    IN ULONG Attributes,
    IN MEMORY_TYPE Type
    )

/*++

Routine Description:

    Initializes a memory descriptor.

Arguments:

    FirstPage - The first page in the region.

    MappedFirstPage - The first page in the mapping.

    Attributes - Memory attributes of the region.

    Type - Memory type of the region.

Return Value:

    Pointer to the memory descriptor,
    NULL if an error occurs.

--*/

{
    PMEMORY_DESCRIPTOR Descriptor;

    if (MmGlobalMemoryDescriptorsUsed >= MmGlobalMemoryDescriptorCount) {
        DebugPrint(L"MmMdInitDescriptor(): No free descriptors available\r\n");
        return NULL;
    }

    Descriptor = &MmGlobalMemoryDescriptors[MmGlobalMemoryDescriptorsUsed++];
    Descriptor->FirstPage = FirstPage;
    Descriptor->MappedFirstPage = MappedFirstPage;
    Descriptor->PageCount = PageCount;
    Descriptor->Attributes = Attributes;
    Descriptor->Type = Type;
    InitializeListHead(&Descriptor->ListEntry);

    return Descriptor;
}

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;
    MmGlobalMemoryDescriptorsUsed = 0;
    RtlZeroMemory(MmGlobalMemoryDescriptors, MAX_STATIC_DESCRIPTOR_COUNT * sizeof(MEMORY_DESCRIPTOR));
    DebugPrintf(L"Global memory descriptor count: %x\r\n", MmGlobalMemoryDescriptorCount);
}