From 13341c8a94689f249e06c7cc05f4952f2fd381cf Mon Sep 17 00:00:00 2001
From: Richard Brooksby <rb@ravenbrook.com>
Date: Sun, 26 Jan 2014 14:01:37 +0000
Subject: [PATCH] Allowing splayfindfirst to continue to search even if the
 treetest is not precise, allowing the arena allocator to search for free
 areas using a zone preference.

Copied from Perforce
 Change: 184129
 ServerID: perforce.ravenbrook.com
---
 mps/code/arenavm.c |  2 +-
 mps/code/splay.c   | 71 +++++++++++++++++++++++++++++++++++++---------
 2 files changed, 58 insertions(+), 15 deletions(-)

diff --git a/mps/code/arenavm.c b/mps/code/arenavm.c
index ad8cb8299f6..c1261d9fb1d 100644
--- a/mps/code/arenavm.c
+++ b/mps/code/arenavm.c
@@ -1746,7 +1746,6 @@ static void VMCompact(Arena arena, Trace trace)
   vmem1 = VMArenaReserved(arena);
 
   /* Destroy any empty chunks (except the primary). */
-  sparePagesPurge(vmArena);
   RING_FOR(node, &arena->chunkRing, next) {
     Chunk chunk = RING_ELT(Chunk, chunkRing, node);
     if(chunk != arena->primary
@@ -1754,6 +1753,7 @@ static void VMCompact(Arena arena, Trace trace)
       Addr base = chunk->base;
       Size size = AddrOffset(chunk->base, chunk->limit);
 
+      sparePagesPurge(vmArena);
       vmChunkDestroy(chunk);
 
       vmArena->contracted(arena, base, size);
diff --git a/mps/code/splay.c b/mps/code/splay.c
index 4042f2c0e07..e0b4ca73c5d 100644
--- a/mps/code/splay.c
+++ b/mps/code/splay.c
@@ -803,6 +803,7 @@ typedef struct {
   void *p;
   Size s;
   SplayTree tree;
+  Bool found;
 } SplayFindClosureStruct, *SplayFindClosure;
 
 static Compare SplayFindFirstCompare(void *key, SplayNode node)
@@ -823,15 +824,23 @@ static Compare SplayFindFirstCompare(void *key, SplayNode node)
   testNode = closure->testNode;
   testTree = closure->testTree;
   tree = closure->tree;
-
+  
   if (SplayNodeLeftChild(node) != NULL &&
      (*testTree)(tree, SplayNodeLeftChild(node), closureP, closureS)) {
     return CompareLESS;
   } else if ((*testNode)(tree, node, closureP, closureS)) {
+    closure->found = TRUE;
     return CompareEQUAL;
   } else {
-    AVER(SplayNodeRightChild(node) != NULL);
-    AVER((*testTree)(tree, SplayNodeRightChild(node), closureP, closureS));
+    /* If there's a right subtree but it doesn't satisfy the tree test
+       then we want to terminate the splay right now.  SplaySplay will
+       return TRUE, so the caller must check closure->found to find out
+       whether the result node actually satisfies testNode. */
+    if (SplayNodeRightChild(node) != NULL &&
+        !(*testTree)(tree, SplayNodeRightChild(node), closureP, closureS)) {
+      closure->found = FALSE;
+      return CompareEQUAL;
+    }
     return CompareGREATER;
   }
 }
@@ -859,10 +868,15 @@ static Compare SplayFindLastCompare(void *key, SplayNode node)
      (*testTree)(tree, SplayNodeRightChild(node), closureP, closureS)) {
      return CompareGREATER;
   } else if ((*testNode)(tree, node, closureP, closureS)) {
+    closure->found = TRUE;
     return CompareEQUAL;
   } else {
-    AVER(SplayNodeLeftChild(node) != NULL);
-    AVER((*testTree)(tree, SplayNodeLeftChild(node), closureP, closureS));
+    /* See SplayFindFirstCompare. */
+    if (SplayNodeLeftChild(node) != NULL &&
+        !(*testTree)(tree, SplayNodeLeftChild(node), closureP, closureS)) {
+      closure->found = FALSE;
+      return CompareEQUAL;
+    }
     return CompareLESS;
   }
 }
@@ -881,12 +895,13 @@ static Compare SplayFindLastCompare(void *key, SplayNode node)
  */
 
 Bool SplayFindFirst(SplayNode *nodeReturn, SplayTree tree,
-                           SplayTestNodeMethod testNode,
-                           SplayTestTreeMethod testTree,
-                           void *closureP, Size closureS)
+                    SplayTestNodeMethod testNode,
+                    SplayTestTreeMethod testTree,
+                    void *closureP, Size closureS)
 {
   SplayNode node;
   SplayFindClosureStruct closureStruct;
+  Bool found;
 
   AVER(nodeReturn != NULL);
   AVERT(SplayTree, tree);
@@ -903,14 +918,42 @@ Bool SplayFindFirst(SplayNode *nodeReturn, SplayTree tree,
   closureStruct.testNode = testNode;
   closureStruct.testTree = testTree;
   closureStruct.tree = tree;
+  closureStruct.found = FALSE;
 
-  if (SplaySplay(&node, tree, (void *)&closureStruct,
-                 &SplayFindFirstCompare)) {
-    *nodeReturn = node;
-    return TRUE;
-  } else {
-    return FALSE;
+  found = SplaySplay(&node, tree, (void *)&closureStruct,
+                     &SplayFindFirstCompare) && closureStruct.found;
+
+  while (!found) {
+    SplayNode oldRoot, newRoot;
+    
+    oldRoot = SplayTreeRoot(tree);
+    newRoot = SplayNodeRightChild(oldRoot);
+
+    if (newRoot == NULL || !(*testTree)(tree, newRoot, closureP, closureS))
+      return FALSE; /* no suitable nodes in the rest of the tree */
+  
+    /* Temporarily chop off the left half-tree, inclusive of root,
+       so that the search excludes any nodes we've seen already. */
+    SplayTreeSetRoot(tree, newRoot);
+    SplayNodeSetRightChild(oldRoot, NULL);
+
+    found = SplaySplay(&node, tree, (void *)&closureStruct,
+                       &SplayFindFirstCompare) && closureStruct.found;
+
+    /* Restore the left tree, then rotate left so that the node we
+       just splayed is at the root.  Update both. */
+    newRoot = SplayTreeRoot(tree);
+    SplayNodeSetRightChild(oldRoot, newRoot);
+    SplayTreeSetRoot(tree, oldRoot);
+    SplayRotateLeft(&tree->root, tree);
+    if (tree->updateNode != NULL) {
+      SplayNodeUpdate(tree, oldRoot);
+      SplayNodeUpdate(tree, newRoot);
+    }
   }
+
+  *nodeReturn = node;
+  return TRUE;
 }