ptos: mm: Add virtual address descriptor allocation

Signed-off-by: Chloe M. <chloe@mensia.org>

service/ptos/head/mm/vad.h

@@ -54,4 +54,15 @@ VOID MmVadTreeInit(MM_VAD_TREE *Tree);
  */
 PT_STATUS MmVadTreeInsert(MM_VAD_TREE *Tree, MM_VAD *Vad);
 
+/*
+ * Allocate a number of virtual address descriptors from a
+ * VAD tree.
+ *
+ * @Tree:       VAD tree to allocate from
+ * @PageCount:  Page count
+ *
+ * Returns NULL on failure / out of memory
+ */
+MM_VAD *MmVadTreeAllocate(MM_VAD_TREE *Tree, USIZE PageCount);
+
 #endif  /* !_MM_VAD_H_ */

service/ptos/mm/vad.c

@@ -7,6 +7,7 @@
  */
 
 #include <mm/vad.h>
+#include <hal/page.h>
 
 /* Prototypes */
 static LONG VadCmp(MM_VAD *V0, MM_VAD *V1);
@@ -49,3 +50,56 @@ MmVadTreeInsert(MM_VAD_TREE *Tree, MM_VAD *Vad)
     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;
+}