From: Paul Gofman pgofman@codeweavers.com
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);
From: Paul Gofman pgofman@codeweavers.com
Based on Wine RB tree implementation in include/wine/rb_tree.h. --- dlls/ntdll/ntdll.spec | 1 + dlls/ntdll/rtl.c | 83 +++++++++++++++++++++++++++++++++++++++++++ include/winternl.h | 1 + 3 files changed, 85 insertions(+)
diff --git a/dlls/ntdll/ntdll.spec b/dlls/ntdll/ntdll.spec index 5281d2dfbf5..41d0500214d 100644 --- a/dlls/ntdll/ntdll.spec +++ b/dlls/ntdll/ntdll.spec @@ -952,6 +952,7 @@ @ stdcall RtlRandom(ptr) @ stdcall RtlRandomEx(ptr) @ stdcall RtlRbInsertNodeEx(ptr ptr long ptr) +@ stdcall RtlRbRemoveNode(ptr ptr) @ stdcall RtlReAllocateHeap(long long ptr long) @ stub RtlReadMemoryStream @ stub RtlReadOutOfProcessMemoryStream diff --git a/dlls/ntdll/rtl.c b/dlls/ntdll/rtl.c index 1890144f225..55b0f5b8416 100644 --- a/dlls/ntdll/rtl.c +++ b/dlls/ntdll/rtl.c @@ -870,6 +870,89 @@ void WINAPI RtlRbInsertNodeEx( RTL_RB_TREE *tree, RTL_BALANCED_NODE *parent, BOO tree->root->Red = 0; }
+/********************************************************************* + * RtlRbRemoveNode [NTDLL.@] + */ +void WINAPI RtlRbRemoveNode( RTL_RB_TREE *tree, RTL_BALANCED_NODE *node ) +{ + RTL_BALANCED_NODE *iter = NULL, *child, *parent, *w; + BOOL need_fixup; + int right; + + TRACE( "tree %p, node %p.\n", tree, node ); + + if (node->Right && (node->Left || tree->min == node)) + { + for (iter = node->Right; iter->Left; iter = iter->Left) + ; + if (tree->min == node) tree->min = iter; + } + else if (tree->min == node) tree->min = rtl_node_parent( node ); + if (!iter || !node->Left) iter = node; + + child = iter->Left ? iter->Left : iter->Right; + + if (!(parent = rtl_node_parent( iter ))) tree->root = child; + else if (iter == parent->Left) parent->Left = child; + else parent->Right = child; + + if (child) rtl_set_node_parent( child, parent ); + + need_fixup = !iter->Red; + + if (node != iter) + { + *iter = *node; + if (!(w = rtl_node_parent( iter ))) tree->root = iter; + else if (node == w->Left) w->Left = iter; + else w->Right = iter; + + if (iter->Right) rtl_set_node_parent( iter->Right, iter ); + if (iter->Left) rtl_set_node_parent( iter->Left, iter ); + if (parent == node) parent = iter; + } + + if (!need_fixup) + { + if (tree->root) tree->root->Red = 0; + return; + } + + while (parent && !(child && child->Red)) + { + right = (child == parent->Right); + w = parent->Children[!right]; + if (w->Red) + { + w->Red = 0; + parent->Red = 1; + rtl_rotate( tree, parent, right ); + w = parent->Children[!right]; + } + if ((w->Left && w->Left->Red) || (w->Right && w->Right->Red)) + { + if (!(w->Children[!right] && w->Children[!right]->Red)) + { + w->Children[right]->Red = 0; + w->Red = 1; + rtl_rotate( tree, w, !right ); + w = parent->Children[!right]; + } + w->Red = parent->Red; + parent->Red = 0; + if (w->Children[!right]) w->Children[!right]->Red = 0; + rtl_rotate( tree, parent, right ); + child = NULL; + break; + } + w->Red = 1; + child = parent; + parent = rtl_node_parent( child ); + } + if (child) child->Red = 0; + if (tree->root) tree->root->Red = 0; +} + /*********************************************************************** * RtlInitializeGenericTableAvl (NTDLL.@) */ diff --git a/include/winternl.h b/include/winternl.h index 8e172125e44..4f24a921bb1 100644 --- a/include/winternl.h +++ b/include/winternl.h @@ -5006,6 +5006,7 @@ 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 void WINAPI RtlRbRemoveNode(RTL_RB_TREE*,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);
From: Paul Gofman pgofman@codeweavers.com
--- dlls/ntdll/tests/rtl.c | 148 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+)
diff --git a/dlls/ntdll/tests/rtl.c b/dlls/ntdll/tests/rtl.c index e678f7f43a0..d013848c07a 100644 --- a/dlls/ntdll/tests/rtl.c +++ b/dlls/ntdll/tests/rtl.c @@ -36,6 +36,7 @@ #include "ddk/ntifs.h" #include "wine/test.h" #include "wine/asm.h" +#include "wine/rbtree.h"
#ifndef __WINE_WINTERNL_H
@@ -107,6 +108,9 @@ static NTSTATUS (WINAPI *pLdrEnumerateLoadedModules)(void *, void *, void *); static NTSTATUS (WINAPI *pLdrRegisterDllNotification)(ULONG, PLDR_DLL_NOTIFICATION_FUNCTION, void *, void **); static NTSTATUS (WINAPI *pLdrUnregisterDllNotification)(void *); static VOID (WINAPI *pRtlGetDeviceFamilyInfoEnum)(ULONGLONG *,DWORD *,DWORD *); +static void (WINAPI *pRtlRbInsertNodeEx)(RTL_RB_TREE *, RTL_BALANCED_NODE *, BOOLEAN, RTL_BALANCED_NODE *); +static void (WINAPI *pRtlRbRemoveNode)(RTL_RB_TREE *, RTL_BALANCED_NODE *); +
static HMODULE hkernel32 = 0; static BOOL (WINAPI *pIsWow64Process)(HANDLE, PBOOL); @@ -151,6 +155,8 @@ static void InitFunctionPtrs(void) pLdrRegisterDllNotification = (void *)GetProcAddress(hntdll, "LdrRegisterDllNotification"); pLdrUnregisterDllNotification = (void *)GetProcAddress(hntdll, "LdrUnregisterDllNotification"); pRtlGetDeviceFamilyInfoEnum = (void *)GetProcAddress(hntdll, "RtlGetDeviceFamilyInfoEnum"); + pRtlRbInsertNodeEx = (void *)GetProcAddress(hntdll, "RtlRbInsertNodeEx"); + pRtlRbRemoveNode = (void *)GetProcAddress(hntdll, "RtlRbRemoveNode"); } hkernel32 = LoadLibraryA("kernel32.dll"); ok(hkernel32 != 0, "LoadLibrary failed\n"); @@ -3784,6 +3790,147 @@ static void test_RtlGetDeviceFamilyInfoEnum(void) trace( "UAP version is %#I64x, device family is %lu, form factor is %lu\n", version, family, form ); }
+struct test_rb_tree_entry +{ + int value; + struct rb_entry wine_rb_entry; + RTL_BALANCED_NODE rtl_entry; +}; + +static int test_rb_tree_entry_compare( const void *key, const struct wine_rb_entry *entry ) +{ + const struct test_rb_tree_entry *t = WINE_RB_ENTRY_VALUE(entry, struct test_rb_tree_entry, wine_rb_entry); + const int *value = key; + + return *value - t->value; +} + +static int test_rtl_rb_tree_entry_compare( const void *key, const RTL_BALANCED_NODE *entry ) +{ + const struct test_rb_tree_entry *t = CONTAINING_RECORD(entry, struct test_rb_tree_entry, rtl_entry); + const int *value = key; + + return *value - t->value; +} + +static int rtl_rb_tree_put( RTL_RB_TREE *tree, const void *key, RTL_BALANCED_NODE *entry, + int (*compare_func)( const void *key, const RTL_BALANCED_NODE *entry )) +{ + RTL_BALANCED_NODE *parent = tree->root; + BOOLEAN right = 0; + int c; + + while (parent) + { + if (!(c = compare_func( key, parent ))) return -1; + right = c > 0; + if (!parent->Children[right]) break; + parent = parent->Children[right]; + } + pRtlRbInsertNodeEx( tree, parent, right, entry ); + return 0; +} + +static struct test_rb_tree_entry *test_rb_tree_entry_from_wine_rb( struct rb_entry *entry ) +{ + if (!entry) return NULL; + return CONTAINING_RECORD(entry, struct test_rb_tree_entry, wine_rb_entry); +} + +static struct test_rb_tree_entry *test_rb_tree_entry_from_rtl_rb( RTL_BALANCED_NODE *entry ) +{ + if (!entry) return NULL; + return CONTAINING_RECORD(entry, struct test_rb_tree_entry, rtl_entry); +} + +static struct test_rb_tree_entry *test_rb_tree_entry_rtl_parent( struct test_rb_tree_entry *node ) +{ + return test_rb_tree_entry_from_rtl_rb( (void *)(node->rtl_entry.ParentValue + & ~(ULONG_PTR)RTL_BALANCED_NODE_RESERVED_PARENT_MASK) ); +} + +static void test_rb_tree(void) +{ + static int test_values[] = { 44, 51, 6, 66, 69, 20, 87, 80, 72, 86, 90, 16, 54, 61, 62, 14, 27, 39, 42, 41 }; + static const unsigned int count = ARRAY_SIZE(test_values); + + struct test_rb_tree_entry *nodes, *parent, *parent2; + RTL_BALANCED_NODE *prev_min_entry = NULL; + int ret, is_red, min_val; + struct rb_tree wine_tree; + RTL_RB_TREE rtl_tree; + unsigned int i; + + if (!pRtlRbInsertNodeEx) + { + win_skip( "RtlRbInsertNodeEx is not present.\n" ); + return; + } + + memset( &rtl_tree, 0, sizeof(rtl_tree) ); + nodes = malloc( count * sizeof(*nodes) ); + memset( nodes, 0xcc, count * sizeof(*nodes) ); + + min_val = test_values[0]; + rb_init( &wine_tree, test_rb_tree_entry_compare ); + for (i = 0; i < count; ++i) + { + winetest_push_context( "i %u", i ); + nodes[i].value = test_values[i]; + ret = rb_put( &wine_tree, &nodes[i].value, &nodes[i].wine_rb_entry ); + ok( !ret, "got %d.\n", ret ); + parent = test_rb_tree_entry_from_wine_rb( nodes[i].wine_rb_entry.parent ); + ret = rtl_rb_tree_put( &rtl_tree, &nodes[i].value, &nodes[i].rtl_entry, test_rtl_rb_tree_entry_compare ); + ok( !ret, "got %d.\n", ret ); + parent2 = test_rb_tree_entry_rtl_parent( &nodes[i] ); + ok( parent == parent2, "got %p, %p.\n", parent, parent2 ); + is_red = nodes[i].rtl_entry.ParentValue & RTL_BALANCED_NODE_RESERVED_PARENT_MASK; + ok( is_red == rb_is_red( &nodes[i].wine_rb_entry ), "got %d, expected %d.\n", is_red, + rb_is_red( &nodes[i].wine_rb_entry )); + + parent = test_rb_tree_entry_from_wine_rb( wine_tree.root ); + parent2 = test_rb_tree_entry_from_rtl_rb( rtl_tree.root ); + ok( parent == parent2, "got %p, %p.\n", parent, parent2 ); + if (nodes[i].value <= min_val) + { + min_val = nodes[i].value; + prev_min_entry = &nodes[i].rtl_entry; + } + ok( rtl_tree.min == prev_min_entry, "unexpected min tree entry.\n" ); + winetest_pop_context(); + } + + for (i = 0; i < count; ++i) + { + struct test_rb_tree_entry *node; + + winetest_push_context( "i %u", i ); + rb_remove( &wine_tree, &nodes[i].wine_rb_entry ); + pRtlRbRemoveNode( &rtl_tree, &nodes[i].rtl_entry ); + + parent = test_rb_tree_entry_from_wine_rb( wine_tree.root ); + parent2 = test_rb_tree_entry_from_rtl_rb( rtl_tree.root ); + ok( parent == parent2, "got %p, %p.\n", parent, parent2 ); + + parent = test_rb_tree_entry_from_wine_rb( rb_head( wine_tree.root )); + parent2 = test_rb_tree_entry_from_rtl_rb( rtl_tree.min ); + ok( parent == parent2, "got %p, %p.\n", parent, parent2 ); + + RB_FOR_EACH_ENTRY(node, &wine_tree, struct test_rb_tree_entry, wine_rb_entry) + { + is_red = node->rtl_entry.ParentValue & RTL_BALANCED_NODE_RESERVED_PARENT_MASK; + ok( is_red == rb_is_red( &node->wine_rb_entry ), "got %d, expected %d.\n", is_red, rb_is_red( &node->wine_rb_entry )); + parent = test_rb_tree_entry_from_wine_rb( node->wine_rb_entry.parent ); + parent2 = test_rb_tree_entry_rtl_parent( node ); + ok( parent == parent2, "got %p, %p.\n", parent, parent2 ); + } + winetest_pop_context(); + } + ok( !rtl_tree.root, "got %p.\n", rtl_tree.root ); + ok( !rtl_tree.min, "got %p.\n", rtl_tree.min ); + free(nodes); +} + START_TEST(rtl) { InitFunctionPtrs(); @@ -3833,4 +3980,5 @@ START_TEST(rtl) test_RtlValidSecurityDescriptor(); test_RtlFindExportedRoutineByName(); test_RtlGetDeviceFamilyInfoEnum(); + test_rb_tree(); }
From: Paul Gofman pgofman@codeweavers.com
--- dlls/kernel32/tests/module.c | 93 ++++++++++++++++++++++++++++++++++++ dlls/ntdll/loader.c | 35 ++++++++++++++ 2 files changed, 128 insertions(+)
diff --git a/dlls/kernel32/tests/module.c b/dlls/kernel32/tests/module.c index 6b576ea25fd..4b4269c821b 100644 --- a/dlls/kernel32/tests/module.c +++ b/dlls/kernel32/tests/module.c @@ -1674,6 +1674,98 @@ static void test_tls_links(void) CloseHandle(test_tls_links_done); }
+ +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 unsigned int check_address_index_tree( RTL_BALANCED_NODE *node ) +{ + LDR_DATA_TABLE_ENTRY *mod; + unsigned int count; + char *base; + + if (!node) return 0; + ok( (node->ParentValue & RTL_BALANCED_NODE_RESERVED_PARENT_MASK) <= 1, "got ParentValue %#Ix.\n", + node->ParentValue ); + + mod = CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + base = mod->DllBase; + if (node->Left) + { + mod = CONTAINING_RECORD(node->Left, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + ok( (char *)mod->DllBase < base, "wrong ordering.\n" ); + } + if (node->Right) + { + mod = CONTAINING_RECORD(node->Right, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + ok( (char *)mod->DllBase > base, "wrong ordering.\n" ); + } + + count = check_address_index_tree( node->Left ); + count += check_address_index_tree( node->Right ); + return count + 1; +} + +static void test_base_address_index_tree(void) +{ + LIST_ENTRY *first = &NtCurrentTeb()->Peb->LdrData->InLoadOrderModuleList; + unsigned int tree_count, list_count = 0; + LDR_DATA_TABLE_ENTRY *mod, *mod2; + RTL_BALANCED_NODE *root, *node; + LDR_DDAG_NODE *ddag_node; + NTSTATUS status; + HMODULE hexe; + char *base; + + /* Check for old LDR data strcuture. */ + hexe = GetModuleHandleW( NULL ); + ok( !!hexe, "Got NULL exe handle.\n" ); + status = LdrFindEntryForAddress( hexe, &mod ); + ok( !status, "got %#lx.\n", status ); + if (!(ddag_node = mod->DdagNode)) + { + win_skip( "DdagNode is NULL, skipping tests.\n" ); + return; + } + ok( !!ddag_node->Modules.Flink, "Got NULL module link.\n" ); + mod2 = CONTAINING_RECORD(ddag_node->Modules.Flink, LDR_DATA_TABLE_ENTRY, NodeModuleLink); + ok( mod2 == mod || broken( (void **)mod2 == (void **)mod - 1 ), "got %p, expected %p.\n", mod2, mod ); + if (mod2 != mod) + { + win_skip( "Old LDR_DATA_TABLE_ENTRY structure, skipping tests.\n" ); + return; + } + + mod = CONTAINING_RECORD(first->Flink, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks); + ok( mod->BaseAddressIndexNode.ParentValue || mod->BaseAddressIndexNode.Left || mod->BaseAddressIndexNode.Right, + "got zero BaseAddressIndexNode.\n" ); + root = &mod->BaseAddressIndexNode; + while (rtl_node_parent( root )) + root = rtl_node_parent( root ); + tree_count = check_address_index_tree( root ); + for (LIST_ENTRY *entry = first->Flink; entry != first; entry = entry->Flink) + { + ++list_count; + mod = CONTAINING_RECORD(entry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks); + base = mod->DllBase; + node = root; + mod2 = NULL; + while (1) + { + ok( !!node, "got NULL.\n" ); + if (!node) break; + mod2 = CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + if (base == (char *)mod2->DllBase) break; + if (base < (char *)mod2->DllBase) node = node->Left; + else node = node->Right; + } + ok( base == (char *)mod2->DllBase, "module %s not found.\n", debugstr_w(mod->BaseDllName.Buffer) ); + } + ok( tree_count == list_count, "count mismatch %u, %u.\n", tree_count, list_count ); +} + START_TEST(module) { WCHAR filenameW[MAX_PATH]; @@ -1711,4 +1803,5 @@ START_TEST(module) test_apisets(); test_ddag_node(); test_tls_links(); + test_base_address_index_tree(); } diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c index 24652d5a663..57b33c0864a 100644 --- a/dlls/ntdll/loader.c +++ b/dlls/ntdll/loader.c @@ -175,6 +175,8 @@ static PEB_LDR_DATA ldr = { &ldr.InInitializationOrderModuleList, &ldr.InInitializationOrderModuleList } };
+static RTL_RB_TREE base_address_index_tree; + static RTL_BITMAP tls_bitmap; static RTL_BITMAP tls_expansion_bitmap;
@@ -237,6 +239,24 @@ static void module_push_unload_trace( const WINE_MODREF *wm ) unload_trace_ptr = unload_traces; }
+static int rtl_rb_tree_put( RTL_RB_TREE *tree, const void *key, RTL_BALANCED_NODE *entry, + int (*compare_func)( const void *key, const RTL_BALANCED_NODE *entry )) +{ + RTL_BALANCED_NODE *parent = tree->root; + BOOLEAN right = 0; + int c; + + while (parent) + { + if (!(c = compare_func( key, parent ))) return -1; + right = c > 0; + if (!parent->Children[right]) break; + parent = parent->Children[right]; + } + RtlRbInsertNodeEx( tree, parent, right, entry ); + return 0; +} + #ifdef __arm64ec__
static void update_hybrid_pointer( void *module, const IMAGE_SECTION_HEADER *sec, UINT rva, void *ptr ) @@ -558,6 +578,17 @@ static void call_ldr_notifications( ULONG reason, LDR_DATA_TABLE_ENTRY *module ) } }
+/* compare base address */ +static int base_address_compare( const void *key, const RTL_BALANCED_NODE *entry ) +{ + const LDR_DATA_TABLE_ENTRY *mod = CONTAINING_RECORD(entry, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + const char *base = key; + + if (base < (char *)mod->DllBase) return -1; + if (base > (char *)mod->DllBase) return 1; + return 0; +} + /************************************************************************* * get_modref * @@ -1559,6 +1590,8 @@ static WINE_MODREF *alloc_module( HMODULE hModule, const UNICODE_STRING *nt_name &wm->ldr.InLoadOrderLinks); InsertTailList(&NtCurrentTeb()->Peb->LdrData->InMemoryOrderModuleList, &wm->ldr.InMemoryOrderLinks); + if (rtl_rb_tree_put( &base_address_index_tree, wm->ldr.DllBase, &wm->ldr.BaseAddressIndexNode, base_address_compare )) + ERR( "rtl_rb_tree_put failed.\n" ); /* wait until init is called for inserting into InInitializationOrderModuleList */
if (!(nt->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_NX_COMPAT)) @@ -2255,6 +2288,7 @@ static NTSTATUS build_module( LPCWSTR load_path, const UNICODE_STRING *nt_name, /* the module has only be inserted in the load & memory order lists */ RemoveEntryList(&wm->ldr.InLoadOrderLinks); RemoveEntryList(&wm->ldr.InMemoryOrderLinks); + RtlRbRemoveNode( &base_address_index_tree, &wm->ldr.BaseAddressIndexNode );
/* FIXME: there are several more dangling references * left. Including dlls loaded by this dll before the @@ -3918,6 +3952,7 @@ static void free_modref( WINE_MODREF *wm )
RemoveEntryList(&wm->ldr.InLoadOrderLinks); RemoveEntryList(&wm->ldr.InMemoryOrderLinks); + RtlRbRemoveNode( &base_address_index_tree, &wm->ldr.BaseAddressIndexNode ); if (wm->ldr.InInitializationOrderLinks.Flink) RemoveEntryList(&wm->ldr.InInitializationOrderLinks);
From: Paul Gofman pgofman@codeweavers.com
--- dlls/ntdll/loader.c | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-)
diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c index 57b33c0864a..9ed863c6b0b 100644 --- a/dlls/ntdll/loader.c +++ b/dlls/ntdll/loader.c @@ -257,6 +257,20 @@ static int rtl_rb_tree_put( RTL_RB_TREE *tree, const void *key, RTL_BALANCED_NOD return 0; }
+static RTL_BALANCED_NODE *rtl_rb_tree_get( RTL_RB_TREE *tree, const void *key, + int (*compare_func)( const void *key, const RTL_BALANCED_NODE *entry )) +{ + RTL_BALANCED_NODE *parent = tree->root; + int c; + + while (parent) + { + if (!(c = compare_func( key, parent ))) return parent; + parent = parent->Children[c > 0]; + } + return NULL; +} + #ifdef __arm64ec__
static void update_hybrid_pointer( void *module, const IMAGE_SECTION_HEADER *sec, UINT rva, void *ptr ) @@ -597,19 +611,14 @@ static int base_address_compare( const void *key, const RTL_BALANCED_NODE *entry */ static WINE_MODREF *get_modref( HMODULE hmod ) { - PLIST_ENTRY mark, entry; PLDR_DATA_TABLE_ENTRY mod; + RTL_BALANCED_NODE *node;
if (cached_modref && cached_modref->ldr.DllBase == hmod) return cached_modref;
- mark = &NtCurrentTeb()->Peb->LdrData->InMemoryOrderModuleList; - for (entry = mark->Flink; entry != mark; entry = entry->Flink) - { - mod = CONTAINING_RECORD(entry, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks); - if (mod->DllBase == hmod) - return cached_modref = CONTAINING_RECORD(mod, WINE_MODREF, ldr); - } - return NULL; + if (!(node = rtl_rb_tree_get( &base_address_index_tree, hmod, base_address_compare ))) return NULL; + mod = CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + return cached_modref = CONTAINING_RECORD(mod, WINE_MODREF, ldr); }
From: Paul Gofman pgofman@codeweavers.com
--- dlls/ntdll/loader.c | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-)
diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c index 9ed863c6b0b..0ed37ad390c 100644 --- a/dlls/ntdll/loader.c +++ b/dlls/ntdll/loader.c @@ -1922,6 +1922,17 @@ NTSTATUS WINAPI LdrDisableThreadCalloutsForDll(HMODULE hModule) return ret; }
+/* compare base address */ +static int module_address_search_compare( const void *key, const RTL_BALANCED_NODE *entry ) +{ + const LDR_DATA_TABLE_ENTRY *mod = CONTAINING_RECORD(entry, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + const char *addr = key; + + if (addr < (char *)mod->DllBase) return -1; + if (addr >= (char *)mod->DllBase + mod->SizeOfImage) return 1; + return 0; +} + /****************************************************************** * LdrFindEntryForAddress (NTDLL.@) * @@ -1929,21 +1940,12 @@ NTSTATUS WINAPI LdrDisableThreadCalloutsForDll(HMODULE hModule) */ NTSTATUS WINAPI LdrFindEntryForAddress( const void *addr, PLDR_DATA_TABLE_ENTRY *pmod ) { - PLIST_ENTRY mark, entry; - PLDR_DATA_TABLE_ENTRY mod; + RTL_BALANCED_NODE *node;
- mark = &NtCurrentTeb()->Peb->LdrData->InMemoryOrderModuleList; - for (entry = mark->Flink; entry != mark; entry = entry->Flink) - { - mod = CONTAINING_RECORD(entry, LDR_DATA_TABLE_ENTRY, InMemoryOrderLinks); - if (mod->DllBase <= addr && - (const char *)addr < (char*)mod->DllBase + mod->SizeOfImage) - { - *pmod = mod; - return STATUS_SUCCESS; - } - } - return STATUS_NO_MORE_ENTRIES; + if (!(node = rtl_rb_tree_get( &base_address_index_tree, addr, module_address_search_compare ))) + return STATUS_NO_MORE_ENTRIES; + *pmod = CONTAINING_RECORD(node, LDR_DATA_TABLE_ENTRY, BaseAddressIndexNode); + return STATUS_SUCCESS; }
/******************************************************************
Hi,
It looks like your patch introduced the new failures shown below. Please investigate and fix them before resubmitting your patch. If they are not new, fixing them anyway would help a lot. Otherwise please ask for the known failures list to be updated.
The tests also ran into some preexisting test failures. If you know how to fix them that would be helpful. See the TestBot job for the details:
The full results can be found at: https://testbot.winehq.org/JobDetails.pl?Key=148549
Your paranoid android.
=== w1064v1507 (64 bit report) ===
Report validation errors: kernel32:module crashed (c0000005)
Test Drive Unlimited Solar Crown depends on that. Besides, this serves as a small optimization for module search by address.
-
Marvin -
Paul Gofman -
Paul Gofman (＠gofman)