From 597628a6446452d2a165ef697cade67b6e9d3f8c Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken91@gmail.com>
Date: Mon, 23 Mar 2026 12:38:31 +0100
Subject: [PATCH] Refactor big allocation tracker to use Tag

---
 sdk/xtdk/mmtypes.h   |  7 ++++---
 xtoskrnl/mm/alloc.cc | 10 +++++-----
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/sdk/xtdk/mmtypes.h b/sdk/xtdk/mmtypes.h
index 93a1142..ceddf22 100644
--- a/sdk/xtdk/mmtypes.h
+++ b/sdk/xtdk/mmtypes.h
@@ -33,8 +33,9 @@
 #define MM_POOL_INVALID_ALLOC_RUNLEVEL             8
 #define MM_POOL_INVALID_FREE_RUNLEVEL              9
 
-/* Big allocations entry flags */
+/* Pool flags */
 #define MM_POOL_BIG_ALLOCATIONS_ENTRY_FREE         0x1
+#define MM_POOL_RAISE_EXCEPTION                    0x10
 
 /* Number of reserved zeroed PTEs */
 #define MM_RESERVED_ZERO_PTES                      32
@@ -225,10 +226,10 @@ typedef struct _POOL_HEADER
 /* Pool descriptor structure definition */
 typedef struct _POOL_TRACKER_BIG_ALLOCATIONS
 {
-    PVOID VirtualAddress;
-    ULONG Key;
     ULONG NumberOfPages;
     PVOID QuotaObject;
+    ULONG Tag;
+    PVOID VirtualAddress;
 } POOL_TRACKER_BIG_ALLOCATIONS, *PPOOL_TRACKER_BIG_ALLOCATIONS;
 
 #endif /* __XTDK_MMTYPES_H */
diff --git a/xtoskrnl/mm/alloc.cc b/xtoskrnl/mm/alloc.cc
index a6fe08e..c2cefe1 100644
--- a/xtoskrnl/mm/alloc.cc
+++ b/xtoskrnl/mm/alloc.cc
@@ -1282,8 +1282,8 @@ MM::Allocator::InitializeBigAllocationsTable(VOID)
  * @param VirtualAddress
  *        Supplies the virtual address of the big allocation.
  *
- * @param Key
- *        Supplies the key used to identify the allocation.
+ * @param Tag
+ *        Supplies the tag used to identify the allocation.
  *
  * @param NumberOfPages
  *        Supplies the number of physical pages backing the allocation.
@@ -1298,7 +1298,7 @@ MM::Allocator::InitializeBigAllocationsTable(VOID)
 BOOLEAN
 XTAPI
 MM::Allocator::RegisterBigAllocationTag(IN PVOID VirtualAddress,
-                                        IN ULONG Key,
+                                        IN ULONG Tag,
                                         IN ULONG NumberOfPages,
                                         IN MMPOOL_TYPE PoolType)
 {
@@ -1336,8 +1336,8 @@ MM::Allocator::RegisterBigAllocationTag(IN PVOID VirtualAddress,
                 if((ULONG_PTR)Entry->VirtualAddress & MM_POOL_BIG_ALLOCATIONS_ENTRY_FREE)
                 {
                     /* Populate the available bucket with the allocation metadata */
-                    Entry->Key = Key;
                     Entry->NumberOfPages = NumberOfPages;
+                    Entry->Tag = Tag;
                     Entry->VirtualAddress = VirtualAddress;
 
                     /* Increment the global usage counter */
@@ -1455,7 +1455,7 @@ MM::Allocator::UnregisterBigAllocationTag(IN PVOID VirtualAddress,
             {
                 /* Capture the allocation metadata */
                 *NumberOfPages = Entry->NumberOfPages;
-                PoolTag = Entry->Key;
+                PoolTag = Entry->Tag;
 
                 /* Invalidate the entry */
                 Entry->VirtualAddress = (PVOID)MM_POOL_BIG_ALLOCATIONS_ENTRY_FREE;