[tkreuzer] 63336: [NTOSKRNL] Modify the VAD node code to return a TABLE_SEARCH_RESULT instead of an NTSTATUS, this allows us to avoid a 2nd tree iteration when inserting VADs. Fix several bugs in M...

Author: tkreuzer
Date: Sat May 17 20:34:11 2014
New Revision: 63336

URL: http://svn.reactos.org/svn/reactos?rev=63336&view=rev
Log:
[NTOSKRNL]
Modify the VAD node code to return a TABLE_SEARCH_RESULT instead of an NTSTATUS, this
allows us to avoid a 2nd tree iteration when inserting VADs. Fix several bugs in
MiFindEmptyAddressRangeDownBasedTree. This code now contains a compatibility hack, that
emulates a bug in the Windows 2003 kernel. Note that this bug is also present in
MiFindEmptyAddressRangeDownTree on Windows 2003, but will not be exposed to the user,
since it only affects the region above the top-most VAD, which will always be occupied by
the PEB or TEB. Implement MEM_TOPDOWN in NtAllocateVirtualMemory. See CORE-6392

Modified:
    trunk/reactos/ntoskrnl/mm/ARM3/miarm.h
    trunk/reactos/ntoskrnl/mm/ARM3/procsup.c
    trunk/reactos/ntoskrnl/mm/ARM3/section.c
    trunk/reactos/ntoskrnl/mm/ARM3/vadnode.c
    trunk/reactos/ntoskrnl/mm/ARM3/virtual.c

Modified: trunk/reactos/ntoskrnl/mm/ARM3/miarm.h
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/miarm.h?r…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/miarm.h	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/miarm.h	[iso-8859-1] Sat May 17 20:34:11 2014
@@ -2093,12 +2093,13 @@
     IN PVOID VirtualAddress
 );
 
-PMMADDRESS_NODE
+TABLE_SEARCH_RESULT
 NTAPI
 MiCheckForConflictingNode(
     IN ULONG_PTR StartVpn,
     IN ULONG_PTR EndVpn,
-    IN PMM_AVL_TABLE Table
+    IN PMM_AVL_TABLE Table,
+    OUT PMMADDRESS_NODE *NodeOrParent
 );
 
 TABLE_SEARCH_RESULT
@@ -2122,7 +2123,7 @@
     OUT PULONG_PTR Base
 );
 
-NTSTATUS
+TABLE_SEARCH_RESULT
 NTAPI
 MiFindEmptyAddressRangeInTree(
     IN SIZE_T Length,

Modified: trunk/reactos/ntoskrnl/mm/ARM3/procsup.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/procsup.c…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/procsup.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/procsup.c	[iso-8859-1] Sat May 17 20:34:11 2014
@@ -110,7 +110,7 @@
     {
         /* For TEBs, or if a PEB location couldn't be found, scan the VAD root */
         Result = MiFindEmptyAddressRangeDownTree(ROUND_TO_PAGES(Size),
-                                                 (ULONG_PTR)MM_HIGHEST_VAD_ADDRESS + 1,
+                                                 (ULONG_PTR)MM_HIGHEST_VAD_ADDRESS,
                                                  PAGE_SIZE,
                                                  &Process->VadRoot,
                                                  Base,

Modified: trunk/reactos/ntoskrnl/mm/ARM3/section.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/section.c…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/section.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/section.c	[iso-8859-1] Sat May 17 20:34:11 2014
@@ -1214,6 +1214,8 @@
     ULONG QuotaCharge = 0, QuotaExcess = 0;
     PMMPTE PointerPte, LastPte;
     MMPTE TempPte;
+    PMMADDRESS_NODE Parent;
+    TABLE_SEARCH_RESULT Result;
 
     /* Get the segment for this section */
     Segment = ControlArea->Segment;
@@ -1302,23 +1304,29 @@
         if (AllocationType & MEM_TOP_DOWN)
         {
             /* No, find an address top-down */
-            Status = MiFindEmptyAddressRangeDownTree(*ViewSize,
+            Result = MiFindEmptyAddressRangeDownTree(*ViewSize,
                                                      (ULONG_PTR)MM_HIGHEST_VAD_ADDRESS,
                                                      _64K,
                                                      &Process->VadRoot,
                                                      &StartAddress,
-                                                    
(PMMADDRESS_NODE*)&Process->VadFreeHint);
-            ASSERT(NT_SUCCESS(Status));
+                                                     &Parent);
         }
         else
         {
             /* No, find an address bottom-up */
-            Status = MiFindEmptyAddressRangeInTree(*ViewSize,
+            Result = MiFindEmptyAddressRangeInTree(*ViewSize,
                                                    _64K,
                                                    &Process->VadRoot,
-                                                  
(PMMADDRESS_NODE*)&Process->VadFreeHint,
+                                                   &Parent,
                                                    &StartAddress);
-            ASSERT(NT_SUCCESS(Status));
+        }
+
+        /* Check if we found a suitable location */
+        if (Result == TableFoundNode)
+        {
+            DPRINT1("Not enough free space to insert this section!\n");
+            MiDereferenceControlArea(ControlArea);
+            return STATUS_CONFLICTING_ADDRESSES;
         }
 
         /* Get the ending address, which is the last piece we need for the VAD */
@@ -1343,9 +1351,11 @@
         EndingAddress = (StartAddress + *ViewSize - 1) | (PAGE_SIZE - 1);
 
         /* Make sure it doesn't conflict with an existing allocation */
-        if (MiCheckForConflictingNode(StartAddress >> PAGE_SHIFT,
-                                      EndingAddress >> PAGE_SHIFT,
-                                      &Process->VadRoot))
+        Result = MiCheckForConflictingNode(StartAddress >> PAGE_SHIFT,
+                                           EndingAddress >> PAGE_SHIFT,
+                                           &Process->VadRoot,
+                                           &Parent);
+        if (Result == TableFoundNode)
         {
             DPRINT1("Conflict with SEC_BASED or manually based section!\n");
             MiDereferenceControlArea(ControlArea);
@@ -1395,7 +1405,8 @@
     MiLockProcessWorkingSetUnsafe(Process, Thread);
 
     /* Insert the VAD */
-    MiInsertVad((PMMVAD)Vad, Process);
+    Process->VadRoot.NodeHint = Vad;
+    MiInsertNode(&Process->VadRoot, (PVOID)Vad, Parent, Result);
 
     /* Release the working set */
     MiUnlockProcessWorkingSetUnsafe(Process, Thread);

Modified: trunk/reactos/ntoskrnl/mm/ARM3/vadnode.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/vadnode.c…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/vadnode.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/vadnode.c	[iso-8859-1] Sat May 17 20:34:11 2014
@@ -73,22 +73,25 @@
     return FoundVad;
 }
 
-PMMADDRESS_NODE
+TABLE_SEARCH_RESULT
 NTAPI
 MiCheckForConflictingNode(IN ULONG_PTR StartVpn,
                           IN ULONG_PTR EndVpn,
-                          IN PMM_AVL_TABLE Table)
-{
-    PMMADDRESS_NODE CurrentNode;
+                          IN PMM_AVL_TABLE Table,
+                          OUT PMMADDRESS_NODE *NodeOrParent)
+{
+    PMMADDRESS_NODE ParentNode, CurrentNode;
 
     /* If the tree is empty, there is no conflict */
-    if (!Table->NumberGenericTableElements) return NULL;
-
-    /* Start looping from the right */
+    if (Table->NumberGenericTableElements == 0) return TableEmptyTree;
+
+    /* Start looping from the root node */
     CurrentNode = RtlRightChildAvl(&Table->BalancedRoot);
     ASSERT(CurrentNode != NULL);
     while (CurrentNode)
     {
+        ParentNode = CurrentNode;
+
         /* This address comes after */
         if (StartVpn > CurrentNode->EndingVpn)
         {
@@ -103,12 +106,21 @@
         else
         {
             /* This address is part of this node, return it */
-            break;
-        }
-    }
-
-    /* Return either the conflicting node, or no node at all */
-    return CurrentNode;
+            *NodeOrParent = ParentNode;
+            return TableFoundNode;
+        }
+    }
+
+    /* There is no more child, save the current node as parent */
+    *NodeOrParent = ParentNode;
+    if (StartVpn > ParentNode->EndingVpn)
+    {
+        return TableInsertAsRight;
+    }
+    else
+    {
+        return TableInsertAsLeft;
+    }
 }
 
 VOID
@@ -338,7 +350,7 @@
     return NULL;
 }
 
-NTSTATUS
+TABLE_SEARCH_RESULT
 NTAPI
 MiFindEmptyAddressRangeInTree(IN SIZE_T Length,
                               IN ULONG_PTR Alignment,
@@ -346,67 +358,78 @@
                               OUT PMMADDRESS_NODE *PreviousVad,
                               OUT PULONG_PTR Base)
 {
-    PMMADDRESS_NODE Node;
-    PMMADDRESS_NODE NextNode;
-    ULONG_PTR StartingVpn, HighestVpn, AlignmentVpn, LengthVpn, LowVpn;
+    PMMADDRESS_NODE Node, PreviousNode;
+    ULONG_PTR PageCount, AlignmentVpn, LowVpn, HighVpn;
     ASSERT(Length != 0);
 
-    /* Precompute page numbers for the length, alignment, and starting address */
-    LengthVpn = (Length + (PAGE_SIZE - 1)) >> PAGE_SHIFT;
+    /* Calculate page numbers for the length, alignment, and starting address */
+    PageCount = BYTES_TO_PAGES(Length);
     AlignmentVpn = Alignment >> PAGE_SHIFT;
-    StartingVpn = ROUND_UP((ULONG_PTR)MM_LOWEST_USER_ADDRESS >> PAGE_SHIFT,
-                           AlignmentVpn);
-
-    /* Check if the table is free, so the lowest possible address is available */
-    if (!Table->NumberGenericTableElements) goto FoundAtBottom;
+    LowVpn = ALIGN_UP_BY((ULONG_PTR)MM_LOWEST_USER_ADDRESS >> PAGE_SHIFT,
AlignmentVpn);
+
+    /* Check if the table is empty */
+    if (Table->NumberGenericTableElements == 0)
+    {
+        /* Tree is empty, the candidate address is already the best one */
+        *Base = LowVpn << PAGE_SHIFT;
+        return TableEmptyTree;
+    }
 
     /* Otherwise, follow the leftmost child of the right root node's child */
     Node = RtlRightChildAvl(&Table->BalancedRoot);
     while (RtlLeftChildAvl(Node)) Node = RtlLeftChildAvl(Node);
 
-    /* This is the node for the remaining gap at the bottom, can it be used? */
-    if ((Node->StartingVpn > StartingVpn) &&
-        (LengthVpn < Node->StartingVpn - StartingVpn))
-    {
-FoundAtBottom:
-        /* Use this VAD to store the allocation */
-        *PreviousVad = NULL;
-        *Base = StartingVpn << PAGE_SHIFT;
-        return STATUS_SUCCESS;
-    }
-
-    /* Otherwise, we start a search to find a gap */
-    while (TRUE)
-    {
-        /* The last aligned page number in this entry */
-        LowVpn = ROUND_UP(Node->EndingVpn + 1, AlignmentVpn);
-
-        /* Keep going as long as there's still a next node */
-        NextNode = MiGetNextNode(Node);
-        if (!NextNode) break;
-
-        /* Can this allocation fit in this node? */
-        if ((LengthVpn <= (NextNode->StartingVpn - LowVpn)) &&
-            (NextNode->StartingVpn > LowVpn))
-        {
-Found:
-            /* Yes! Use this VAD to store the allocation */
-            *PreviousVad = Node;
-            *Base = ROUND_UP((Node->EndingVpn << PAGE_SHIFT) | (PAGE_SIZE - 1),
-                             Alignment);
-            return STATUS_SUCCESS;
-        }
-
-        /* Try the next node */
-        Node = NextNode;
-    }
-
-    /* We're down to the last (top) VAD, will this allocation fit inside it? */
-    HighestVpn = ((ULONG_PTR)MM_HIGHEST_VAD_ADDRESS + 1) >> PAGE_SHIFT;
-    if ((HighestVpn > LowVpn) && (LengthVpn <= HighestVpn - LowVpn)) goto
Found;
+    /* Start a search to find a gap */
+    PreviousNode = NULL;
+    while (Node != NULL)
+    {
+        /* Check if the gap below the current node is suitable */
+        if (Node->StartingVpn >= LowVpn + PageCount)
+        {
+            /* There is enough space to add our node */
+            *Base = LowVpn << PAGE_SHIFT;
+
+            /* Can we use the current node as parent? */
+            if (RtlLeftChildAvl(Node) == NULL)
+            {
+                /* Node has no left child, so use it as parent */
+                *PreviousVad = Node;
+                return TableInsertAsLeft;
+            }
+            else
+            {
+                /* Node has a left child, this means that the previous node is
+                   the right-most child of it's left child and can be used as
+                   the parent. In case we use the space before the left-most
+                   node, it's left child must be NULL. */
+                ASSERT(PreviousNode != NULL);
+                ASSERT(RtlRightChildAvl(PreviousNode) == NULL);
+                *PreviousVad = PreviousNode;
+                return TableInsertAsRight;
+            }
+        }
+
+        /* The next candidate is above the current node */
+        if (Node->EndingVpn >= LowVpn)
+            LowVpn = ALIGN_UP_BY(Node->EndingVpn + 1, AlignmentVpn);
+
+        /* Remember the current node and go to the next node */
+        PreviousNode = Node;
+        Node = MiGetNextNode(Node);
+    }
+
+    /* We're up to the highest VAD, will this allocation fit above it? */
+    HighVpn = ((ULONG_PTR)MM_HIGHEST_VAD_ADDRESS + 1) / PAGE_SIZE;
+    if (HighVpn >= LowVpn + PageCount)
+    {
+        /* Yes! Use this VAD to store the allocation */
+        *PreviousVad = PreviousNode;
+        *Base = LowVpn << PAGE_SHIFT;
+        return TableInsertAsRight;
+    }
 
     /* Nyet, there's no free address space for this allocation, so we'll fail */
-    return STATUS_NO_MEMORY;
+    return TableFoundNode;
 }
 
 TABLE_SEARCH_RESULT
@@ -418,55 +441,59 @@
                                 OUT PULONG_PTR Base,
                                 OUT PMMADDRESS_NODE *Parent)
 {
-    PMMADDRESS_NODE Node, LowestNode, Child;
-    ULONG_PTR LowVpn, HighVpn;
+    PMMADDRESS_NODE Node, OldNode, Child;
+    ULONG_PTR LowVpn, HighVpn, AlignmentVpn;
     PFN_NUMBER PageCount;
 
     /* Sanity checks */
     ASSERT(BoundaryAddress);
-    ASSERT(BoundaryAddress <= ((ULONG_PTR)MM_HIGHEST_VAD_ADDRESS + 1));
-
-    /* Compute page length, make sure the boundary address is valid */
+    ASSERT(BoundaryAddress <= ((ULONG_PTR)MM_HIGHEST_VAD_ADDRESS));
+    ASSERT((Alignment & (PAGE_SIZE - 1)) == 0);
+
+    /* Calculate page numbers for the length and alignment */
     Length = ROUND_TO_PAGES(Length);
     PageCount = Length >> PAGE_SHIFT;
-    if ((BoundaryAddress + 1) < Length) return TableFoundNode;
+    AlignmentVpn = Alignment / PAGE_SIZE;
+
+    /* Check if there is enough space below the boundary */
+    if ((ALIGN_UP_BY((ULONG_PTR)MM_LOWEST_USER_ADDRESS, Alignment) + Length) >
+        (BoundaryAddress + 1))
+    {
+        return TableFoundNode;
+    }
 
     /* Check if the table is empty */
     if (Table->NumberGenericTableElements == 0)
     {
         /* Tree is empty, the candidate address is already the best one */
-        *Base = ROUND_UP(BoundaryAddress + 1 - Length, Alignment);
+        *Base = ALIGN_DOWN_BY(BoundaryAddress + 1 - Length, Alignment);
         return TableEmptyTree;
     }
 
     /* Calculate the initial upper margin */
-    HighVpn = BoundaryAddress >> PAGE_SHIFT;
-
-    /* Starting from the root, go down until the right-most child
-     * which is just behind the boundary*/
-    LowestNode = Node = RtlRightChildAvl(&Table->BalancedRoot);
-    while (((Child = RtlRightChildAvl(Node)) != 0 )
-            && (Node->EndingVpn < HighVpn )) Node = Child;
+    HighVpn = (BoundaryAddress + 1) >> PAGE_SHIFT;
+
+    /* Starting from the root, follow the right children until we found a node
+       that ends above the boundary */
+    Node = RtlRightChildAvl(&Table->BalancedRoot);
+    while ((Node->EndingVpn < HighVpn) &&
+           ((Child = RtlRightChildAvl(Node)) != NULL)) Node = Child;
 
     /* Now loop the Vad nodes */
     while (Node)
     {
-        /* Keep track of the lowest node */
-        LowestNode = Node;
-
         /* Calculate the lower margin */
-        LowVpn = ROUND_UP(Node->EndingVpn + 1, Alignment >> PAGE_SHIFT);
+        LowVpn = ALIGN_UP_BY(Node->EndingVpn + 1, AlignmentVpn);
 
         /* Check if the current bounds are suitable */
         if ((HighVpn > LowVpn) && ((HighVpn - LowVpn) >= PageCount))
         {
             /* There is enough space to add our node */
-            LowVpn = HighVpn - PageCount;
+            LowVpn = ALIGN_DOWN_BY(HighVpn - PageCount, AlignmentVpn);
             *Base = LowVpn << PAGE_SHIFT;
 
             /* Can we use the current node as parent? */
-            Child = RtlRightChildAvl(Node);
-            if (!Child)
+            if (!RtlRightChildAvl(Node))
             {
                 /* Node has no right child, so use it as parent */
                 *Parent = Node;
@@ -474,29 +501,29 @@
             }
             else
             {
-                /* Node has a right child, find most left grand child */
-                Node = Child;
-                while ((Child = RtlLeftChildAvl(Node))) Node = Child;
-                *Parent = Node;
+                /* Node has a right child, the node we had before is the most
+                   left grandchild of that right child, use it as parent. */
+                *Parent = OldNode;
                 return TableInsertAsLeft;
             }
         }
 
-        /* Update the upper margin if neccessary */
+        /* Update the upper margin if necessary */
         if (Node->StartingVpn < HighVpn) HighVpn = Node->StartingVpn;
 
-        /* Go to the next lower node */
+        /* Remember the current node and go to the previous node */
+        OldNode = Node;
         Node = MiGetPreviousNode(Node);
     }
 
     /* Check if there's enough space before the lowest Vad */
-    LowVpn = ROUND_UP((ULONG_PTR)MI_LOWEST_VAD_ADDRESS, Alignment) >> PAGE_SHIFT;
+    LowVpn = ALIGN_UP_BY((ULONG_PTR)MI_LOWEST_VAD_ADDRESS, Alignment) / PAGE_SIZE;
     if ((HighVpn > LowVpn) && ((HighVpn - LowVpn) >= PageCount))
     {
         /* There is enough space to add our address */
-        LowVpn = HighVpn - PageCount;
+        LowVpn = ALIGN_DOWN_BY(HighVpn - PageCount, Alignment >> PAGE_SHIFT);
         *Base = LowVpn << PAGE_SHIFT;
-        *Parent = LowestNode;
+        *Parent = OldNode;
         return TableInsertAsLeft;
     }
 
@@ -527,7 +554,7 @@
     if ((BoundaryAddress + 1) < Length) return STATUS_NO_MEMORY;
 
     /* Check if the table is empty */
-    BestVpn = ROUND_UP(BoundaryAddress + 1 - Length, Alignment);
+    BestVpn = ROUND_DOWN(BoundaryAddress + 1 - Length, Alignment);
     if (Table->NumberGenericTableElements == 0)
     {
         /* Tree is empty, the candidate address is already the best one */
@@ -540,11 +567,18 @@
     while (RtlRightChildAvl(Node)) Node = RtlRightChildAvl(Node);
 
     /* Check if we can fit in here */
-    LowVpn = ROUND_UP(Node->EndingVpn, Alignment);
-    if ((LowVpn < BoundaryAddress) && (Length < (BoundaryAddress -
LowVpn)))
-    {
-        /* Return the address */
-        *Base = ROUND_UP(BoundaryAddress - Length, Alignment);
+    LowVpn = ROUND_UP(Node->EndingVpn + 1, Alignment);
+    if ((LowVpn < BoundaryAddress) && (Length <= (BoundaryAddress -
LowVpn)))
+    {
+#if (NTDDI_VERSION >= NTDDI_VISTA)
+        /* Return the address. */
+        *Base = BestVpn;
+#else
+        /* Note: this is a compatibility hack that mimics a bug in the 2k3
+           kernel. It will can waste up to Alignment bytes of memory above
+           the allocation. This bug was fixed in Windows Vista */
+        *Base = ROUND_DOWN(BoundaryAddress - Length, Alignment);
+#endif
         return STATUS_SUCCESS;
     }
 
@@ -558,22 +592,20 @@
         /* Check if this node could contain the requested address */
         LowVpn = ROUND_UP(LowestNode->EndingVpn + 1, Alignment);
         if ((LowestNode->EndingVpn < BestVpn) &&
+            (LowVpn < Node->StartingVpn) &&
             (Length <= (Node->StartingVpn - LowVpn)))
         {
-            /* Check if it fits in perfectly */
-            if ((BestVpn > LowestNode->EndingVpn) &&
-                (BoundaryAddress < Node->StartingVpn))
+            /* Check if we need to take BoundaryAddress into account */
+            if (BoundaryAddress < Node->StartingVpn)
             {
                 /* Return the optimal VPN address */
                 *Base = BestVpn;
                 return STATUS_SUCCESS;
             }
-
-            /* It doesn't, check if it can partly fit */
-            if (Node->StartingVpn > LowVpn)
-            {
-                /* Return an aligned base address within this node */
-                *Base = ROUND_UP(Node->StartingVpn - Length, Alignment);
+            else
+            {
+                /* The upper margin is given by the Node's starting address */
+                *Base = ROUND_DOWN(Node->StartingVpn - Length, Alignment);
                 return STATUS_SUCCESS;
             }
         }
@@ -584,7 +616,7 @@
 
     /* Check if there's enough space before the lowest Vad */
     if ((Node->StartingVpn > (ULONG_PTR)MI_LOWEST_VAD_ADDRESS) &&
-        ((Node->StartingVpn - (ULONG_PTR)MI_LOWEST_VAD_ADDRESS) > Length))
+        ((Node->StartingVpn - (ULONG_PTR)MI_LOWEST_VAD_ADDRESS) >= Length))
     {
         /* Check if it fits in perfectly */
         if (BoundaryAddress < Node->StartingVpn)
@@ -595,7 +627,7 @@
         }
 
         /* Return an aligned base address within this node */
-        *Base = ROUND_UP(Node->StartingVpn - Length, Alignment);
+        *Base = ROUND_DOWN(Node->StartingVpn - Length, Alignment);
         return STATUS_SUCCESS;
     }
 
@@ -669,3 +701,4 @@
 }
 
 /* EOF */
+

Modified: trunk/reactos/ntoskrnl/mm/ARM3/virtual.c
URL:
http://svn.reactos.org/svn/reactos/trunk/reactos/ntoskrnl/mm/ARM3/virtual.c…
==============================================================================
--- trunk/reactos/ntoskrnl/mm/ARM3/virtual.c	[iso-8859-1] (original)
+++ trunk/reactos/ntoskrnl/mm/ARM3/virtual.c	[iso-8859-1] Sat May 17 20:34:11 2014
@@ -1903,6 +1903,7 @@
     BOOLEAN Committed;
     NTSTATUS Status = STATUS_SUCCESS;
     PETHREAD Thread = PsGetCurrentThread();
+    TABLE_SEARCH_RESULT Result;
 
     /* Calculate base address for the VAD */
     StartingAddress = (ULONG_PTR)PAGE_ALIGN((*BaseAddress));
@@ -1939,10 +1940,11 @@
     }
 
     /* Get the VAD for this address range, and make sure it exists */
-    Vad = (PMMVAD)MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
-                                            EndingAddress >> PAGE_SHIFT,
-                                            &Process->VadRoot);
-    if (!Vad)
+    Result = MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
+                                       EndingAddress >> PAGE_SHIFT,
+                                       &Process->VadRoot,
+                                       (PMMADDRESS_NODE*)&Vad);
+    if (Result != TableFoundNode)
     {
         DPRINT("Could not find a VAD for this allocation\n");
         Status = STATUS_CONFLICTING_ADDRESSES;
@@ -4099,6 +4101,8 @@
     BOOLEAN Attached = FALSE, ChangeProtection = FALSE;
     MMPTE TempPte;
     PMMPTE PointerPte, PointerPde, LastPte;
+    TABLE_SEARCH_RESULT Result;
+    PMMADDRESS_NODE Parent;
     PAGED_CODE();
 
     /* Check for valid Zero bits */
@@ -4384,12 +4388,28 @@
         //
         if (!PBaseAddress)
         {
-            Status = MiFindEmptyAddressRangeInTree(PRegionSize,
-                                                   _64K,
-                                                   &Process->VadRoot,
-                                                  
(PMMADDRESS_NODE*)&Process->VadFreeHint,
-                                                   &StartingAddress);
-            if (!NT_SUCCESS(Status)) goto FailPath;
+            /* Which way should we search? */
+            if (AllocationType & MEM_TOP_DOWN)
+            {
+                /* Find an address top-down */
+                Result = MiFindEmptyAddressRangeDownTree(PRegionSize,
+                                                        
(ULONG_PTR)MM_HIGHEST_VAD_ADDRESS,
+                                                         _64K,
+                                                         &Process->VadRoot,
+                                                         &StartingAddress,
+                                                         &Parent);
+            }
+            else
+            {
+                /* Find an address bottom-up */
+                Result = MiFindEmptyAddressRangeInTree(PRegionSize,
+                                                       _64K,
+                                                       &Process->VadRoot,
+                                                       &Parent,
+                                                       &StartingAddress);
+            }
+
+            if (Result == TableFoundNode) goto FailPath;
 
             //
             // Now we know where the allocation ends. Make sure it doesn't end up
@@ -4402,15 +4422,21 @@
                 goto FailPath;
             }
         }
-        else if (MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
-                                           EndingAddress >> PAGE_SHIFT,
-                                           &Process->VadRoot))
-        {
-            //
-            // The address specified is in conflict!
-            //
-            Status = STATUS_CONFLICTING_ADDRESSES;
-            goto FailPath;
+        else
+        {
+            /* Make sure it doesn't conflict with an existing allocation */
+            Result = MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
+                                               EndingAddress >> PAGE_SHIFT,
+                                               &Process->VadRoot,
+                                               &Parent);
+            if (Result == TableFoundNode)
+            {
+                //
+                // The address specified is in conflict!
+                //
+                Status = STATUS_CONFLICTING_ADDRESSES;
+                goto FailPath;
+            }
         }
 
         //
@@ -4429,7 +4455,8 @@
         //
         MiLockProcessWorkingSetUnsafe(Process, CurrentThread);
         Vad->ControlArea = NULL; // For Memory-Area hack
-        MiInsertVad(Vad, Process);
+        Process->VadRoot.NodeHint = Vad;
+        MiInsertNode(&Process->VadRoot, (PVOID)Vad, Parent, Result);
         MiUnlockProcessWorkingSetUnsafe(Process, CurrentThread);
 
         //
@@ -4495,10 +4522,11 @@
     //
     // Get the VAD for this address range, and make sure it exists
     //
-    FoundVad = (PMMVAD)MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
-                                                 EndingAddress >> PAGE_SHIFT,
-                                                 &Process->VadRoot);
-    if (!FoundVad)
+    Result = MiCheckForConflictingNode(StartingAddress >> PAGE_SHIFT,
+                                       EndingAddress >> PAGE_SHIFT,
+                                       &Process->VadRoot,
+                                       (PMMADDRESS_NODE*)&FoundVad);
+    if (Result != TableFoundNode)
     {
         DPRINT1("Could not find a VAD for this allocation\n");
         Status = STATUS_CONFLICTING_ADDRESSES;