[PATCH 1/6] ntdll: Implement RtlRbInsertNodeEx().

[Paul Gofman]

From: Paul Gofman

Based on Wine RB tree implementation in include/wine/rb_tree.h.
---
 dlls/ntdll/ntdll.spec |  1 +
 dlls/ntdll/rtl.c      | 91 +++++++++++++++++++++++++++++++++++++++++++
 include/ntdef.h       |  6 +++
 include/winternl.h    |  1 +
 4 files changed, 99 insertions(+)

diff --git a/dlls/ntdll/ntdll.spec b/dlls/ntdll/ntdll.spec
index 85aa48eacb5..5281d2dfbf5 100644
--- a/dlls/ntdll/ntdll.spec
+++ b/dlls/ntdll/ntdll.spec
@@ -951,6 +951,7 @@
 @ stdcall RtlRaiseStatus(long)
 @ stdcall RtlRandom(ptr)
 @ stdcall RtlRandomEx(ptr)
+@ stdcall RtlRbInsertNodeEx(ptr ptr long ptr)
 @ stdcall RtlReAllocateHeap(long long ptr long)
 @ stub RtlReadMemoryStream
 @ stub RtlReadOutOfProcessMemoryStream
diff --git a/dlls/ntdll/rtl.c b/dlls/ntdll/rtl.c
index 5597daa061f..1890144f225 100644
--- a/dlls/ntdll/rtl.c
+++ b/dlls/ntdll/rtl.c
@@ -779,6 +779,97 @@ void WINAPI RtlGetCurrentProcessorNumberEx(PROCESSOR_NUMBER *processor)
     processor->Reserved = 0;
 }

+static RTL_BALANCED_NODE *rtl_node_parent( RTL_BALANCED_NODE *node )
+{
+    return (RTL_BALANCED_NODE *)(node->ParentValue & ~(ULONG_PTR)RTL_BALANCED_NODE_RESERVED_PARENT_MASK);
+}
+
+static void rtl_set_node_parent( RTL_BALANCED_NODE *node, RTL_BALANCED_NODE *parent )
+{
+    node->ParentValue = (ULONG_PTR)parent | (node->ParentValue & RTL_BALANCED_NODE_RESERVED_PARENT_MASK);
+}
+
+static void rtl_rotate( RTL_RB_TREE *tree, RTL_BALANCED_NODE *n, int right )
+{
+    RTL_BALANCED_NODE *child = n->Children[!right];
+    RTL_BALANCED_NODE *parent = rtl_node_parent( n );
+
+    if (!parent)                tree->root = child;
+    else if (parent->Left == n) parent->Left = child;
+    else                        parent->Right = child;
+
+    n->Children[!right] = child->Children[right];
+    if (n->Children[!right]) rtl_set_node_parent( n->Children[!right], n );
+    child->Children[right] = n;
+    rtl_set_node_parent( child, parent );
+    rtl_set_node_parent( n, child );
+}
+
+static void rtl_flip_color( RTL_BALANCED_NODE *node )
+{
+    node->Red = !node->Red;
+    node->Left->Red = !node->Left->Red;
+    node->Right->Red = !node->Right->Red;
+}
+
+/*********************************************************************
+ *           RtlRbInsertNodeEx [NTDLL.@]
+ */
+void WINAPI RtlRbInsertNodeEx( RTL_RB_TREE *tree, RTL_BALANCED_NODE *parent, BOOLEAN right, RTL_BALANCED_NODE *node )
+{
+    RTL_BALANCED_NODE *grandparent;
+
+    TRACE( "tree %p, parent %p, right %d, node %p.\n", tree, parent, right, node );
+
+    node->ParentValue = (ULONG_PTR)parent;
+    node->Left = NULL;
+    node->Right = NULL;
+
+    if (!parent)
+    {
+        tree->root = tree->min = node;
+        return;
+    }
+    if (right > 1)
+    {
+        ERR( "right %d.\n", right );
+        return;
+    }
+    if (parent->Children[right])
+    {
+        ERR( "parent %p, right %d, child %p.\n", parent, right, parent->Children[right] );
+        return;
+    }
+
+    node->Red = 1;
+    parent->Children[right] = node;
+    if (tree->min == parent && parent->Left == node) tree->min = node;
+    grandparent = rtl_node_parent( parent );
+    while (parent->Red)
+    {
+        right = (parent == grandparent->Right);
+        if (grandparent->Children[!right] && grandparent->Children[!right]->Red)
+        {
+            node = grandparent;
+            rtl_flip_color( node );
+            if (!(parent = rtl_node_parent( node ))) break;
+            grandparent = rtl_node_parent( parent );
+            continue;
+        }
+        if (node == parent->Children[!right])
+        {
+            node = parent;
+            rtl_rotate( tree, node, right );
+            parent = rtl_node_parent( node );
+            grandparent = rtl_node_parent( parent );
+        }
+        parent->Red = 0;
+        grandparent->Red = 1;
+        rtl_rotate( tree, grandparent, !right );
+    }
+    tree->root->Red = 0;
+}
+
 /***********************************************************************
  *           RtlInitializeGenericTableAvl  (NTDLL.@)
  */
diff --git a/include/ntdef.h b/include/ntdef.h
index 8d04f83da5d..e840bd3f9fe 100644
--- a/include/ntdef.h
+++ b/include/ntdef.h
@@ -84,6 +84,12 @@ typedef struct _RTL_BALANCED_NODE

#define RTL_BALANCED_NODE_RESERVED_PARENT_MASK 3

+typedef struct _RTL_RB_TREE
+{
+    RTL_BALANCED_NODE *root;
+    RTL_BALANCED_NODE *min;
+} RTL_RB_TREE, *PRTL_RB_TREE;
+
 #define RTL_CONSTANT_STRING(s) { sizeof(s) - sizeof(s[0]), sizeof(s), (void*)s }

#endif /* _NTDEF_ */
diff --git a/include/winternl.h b/include/winternl.h
index 828fbc6e915..8e172125e44 100644
--- a/include/winternl.h
+++ b/include/winternl.h
@@ -5005,6 +5005,7 @@ NTSYSAPI NTSTATUS  WINAPI RtlQueueWorkItem(PRTL_WORK_ITEM_ROUTINE,PVOID,ULONG);
 NTSYSAPI void      DECLSPEC_NORETURN WINAPI RtlRaiseStatus(NTSTATUS);
 NTSYSAPI ULONG     WINAPI RtlRandom(PULONG);
 NTSYSAPI ULONG     WINAPI RtlRandomEx(PULONG);
+NTSYSAPI void      WINAPI RtlRbInsertNodeEx(RTL_RB_TREE*,RTL_BALANCED_NODE*,BOOLEAN,RTL_BALANCED_NODE*);
 NTSYSAPI PVOID     WINAPI RtlReAllocateHeap(HANDLE,ULONG,PVOID,SIZE_T) __WINE_ALLOC_SIZE(4) __WINE_DEALLOC(RtlFreeHeap,3);
 NTSYSAPI NTSTATUS  WINAPI RtlRegisterWait(PHANDLE,HANDLE,RTL_WAITORTIMERCALLBACKFUNC,PVOID,ULONG,ULONG);
 NTSYSAPI void      WINAPI RtlReleaseActivationContext(HANDLE);