/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Virtual address descriptor management
 * Author:      Chloe M.
 */

#include <mm/vad.h>
#include <hal/page.h>

/* Prototypes */
static LONG VadCmp(MM_VAD *V0, MM_VAD *V1);

/* RBT generation */
RB_PROTOTYPE(VadTree, _MM_VAD, TreeLink, VadCmp);
RB_GENERATE(VadTree, _MM_VAD, TreeLink, VadCmp);

static LONG
VadCmp(MM_VAD *V0, MM_VAD *V1)
{
    MM_RANGE *Range0, *Range1;

    Range0 = &V0->Range;
    Range1 = &V1->Range;

    return (Range0->VirtualBase < Range1->VirtualBase)
        ? -1
        : Range0->VirtualBase > Range1->VirtualBase;
}

VOID
MmVadTreeInit(MM_VAD_TREE *Tree)
{
    if (Tree == NULL) {
        return;
    }

    Tree->Elements = 0;
    RB_INIT(&Tree->Head);
}

PT_STATUS
MmVadTreeInsert(MM_VAD_TREE *Tree, MM_VAD *Vad)
{
    if (Tree == NULL || Vad == NULL) {
        return STATUS_INVALID_PARAM;
    }

    RB_INSERT(VadTree, &Tree->Head, Vad);
    return STATUS_SUCCESS;
}

MM_VAD *
MmVadTreeAllocate(MM_VAD_TREE *Tree, USIZE PageCount)
{
    MM_VAD *VadBase, *VadIter;
    MM_VAD *VadTmp, *VadNext;
    MM_RANGE *RangeIter;
    UPTR RangeStart, RangeEnd;
    USIZE PagesFound = 0;
    USIZE RangeDelta;

    if (Tree == NULL || PageCount == 0) {
        return NULL;
    }

    VadBase = NULL;
    RB_FOREACH(VadIter, VadTree, &Tree->Head) {
        RangeIter = &VadIter->Range;
        RangeStart = RangeIter->VirtualBase;
        RangeEnd = RangeIter->VirtualBase + RangeIter->Length;
        RangeDelta = RangeEnd - RangeStart;

        /*
         * Do we have exactly a page here? Track that if so, otherwise
         * reset what we've already tracked.
         */
        if (RangeDelta == PAGESIZE) {
            if (VadBase == NULL)
                VadBase = (MM_VAD *)RangeIter->VirtualBase;
            ++PagesFound;
        } else {
            VadBase = NULL;
            --PagesFound;
        }

        /*
         * If we have found what we requested, remove the entire
         * range and return the base.
         */
        if (PagesFound >= PageCount) {
            VadTmp = VadBase;
            while (VadTmp != VadIter) {
                VadNext = RB_NEXT(VadTree, &Tree->Head, VadTmp);
                RB_REMOVE(VadTree, &Tree->Head, VadTmp);
                VadTmp = VadNext;
            }
            RB_REMOVE(VadTree, &Tree->Head, VadTmp);
            return VadBase;
        }
    }

    return NULL;
}