[PATCH 3/6] ntdll: Implement RtlEnumerateGenericTable().

[Zhiyi Zhang]

From: Zhiyi Zhang

---
 dlls/ntdll/ntdll.spec               |  2 +-
 dlls/ntdll/rtl.c                    | 36 +++++++++++++++++++++++++++++
 dlls/ntoskrnl.exe/ntoskrnl.exe.spec |  2 +-
 include/ddk/ntddk.h                 |  1 +
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/dlls/ntdll/ntdll.spec b/dlls/ntdll/ntdll.spec
index cb6682d7e5c..c0d4c5eede7 100644
--- a/dlls/ntdll/ntdll.spec
+++ b/dlls/ntdll/ntdll.spec
@@ -659,7 +659,7 @@
 @ stdcall -arch=win32 -ret64 RtlEnlargedUnsignedMultiply(long long)
 @ stdcall RtlEnterCriticalSection(ptr)
 @ stub RtlEnumProcessHeaps
-@ stub RtlEnumerateGenericTable
+@ stdcall RtlEnumerateGenericTable(ptr long)
 # @ stub RtlEnumerateGenericTableAvl
 # @ stub RtlEnumerateGenericTableLikeADirectory
 @ stdcall RtlEnumerateGenericTableWithoutSplaying(ptr ptr)
diff --git a/dlls/ntdll/rtl.c b/dlls/ntdll/rtl.c
index a1647241283..a0588db6957 100644
--- a/dlls/ntdll/rtl.c
+++ b/dlls/ntdll/rtl.c
@@ -700,6 +700,42 @@ void * WINAPI RtlEnumerateGenericTableWithoutSplaying(RTL_GENERIC_TABLE *table,
     return get_data_from_splay_links(child);
 }

+/******************************************************************************
+ *  RtlEnumerateGenericTable           [NTDLL.@]
+ *
+ * RtlEnumerateGenericTable() uses TableRoot to keep track of enumeration status according to tests.
+ * This also means that other functions that change TableRoot should not be used during enumerations.
+ * Otherwise, RtlEnumerateGenericTable() won't be able to find the correct next element when restart
+ * is FALSE. This is also the case on Windows.
+ */
+void * WINAPI RtlEnumerateGenericTable(RTL_GENERIC_TABLE *table, BOOLEAN restart)
+{
+    RTL_SPLAY_LINKS *child;
+
+    TRACE("(%p, %d)\n", table, restart);
+
+    if (RtlIsGenericTableEmpty(table))
+        return NULL;
+
+    if (restart)
+    {
+        /* Find the smallest element */
+        child = table->TableRoot;
+        while (RtlLeftChild(child))
+            child = RtlLeftChild(child);
+    }
+    else
+    {
+        /* Find the successor of the root */
+        child = RtlRealSuccessor(table->TableRoot);
+        if (!child)
+            return NULL;
+    }
+
+    table->TableRoot = RtlSplay(child);
+    return get_data_from_splay_links(child);
+}
+
 /******************************************************************************
  *  RtlNumberGenericTableElements           [NTDLL.@]
  */
diff --git a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
index 49a29d06644..fea13ace291 100644
--- a/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
+++ b/dlls/ntoskrnl.exe/ntoskrnl.exe.spec
@@ -1070,7 +1070,7 @@
 @ stdcall -arch=win32 -ret64 RtlEnlargedIntegerMultiply(long long)
 @ stdcall -arch=win32 RtlEnlargedUnsignedDivide(int64 long ptr)
 @ stdcall -arch=win32 -ret64 RtlEnlargedUnsignedMultiply(long long)
-@ stub RtlEnumerateGenericTable
+@ stdcall RtlEnumerateGenericTable(ptr long)
 @ stub RtlEnumerateGenericTableAvl
 @ stub RtlEnumerateGenericTableLikeADirectory
 @ stdcall RtlEnumerateGenericTableWithoutSplaying(ptr ptr)
diff --git a/include/ddk/ntddk.h b/include/ddk/ntddk.h
index 36c52097c16..887d1749662 100644
--- a/include/ddk/ntddk.h
+++ b/include/ddk/ntddk.h
@@ -303,6 +303,7 @@ void      WINAPI RtlCopyString(STRING*,const STRING*);
 PRTL_SPLAY_LINKS WINAPI RtlDelete(PRTL_SPLAY_LINKS);
 BOOLEAN   WINAPI RtlDeleteElementGenericTable(PRTL_GENERIC_TABLE,PVOID);
 void      WINAPI RtlDeleteNoSplay(PRTL_SPLAY_LINKS,PRTL_SPLAY_LINKS *);
+void *    WINAPI RtlEnumerateGenericTable(PRTL_GENERIC_TABLE,BOOLEAN);
 void *    WINAPI RtlEnumerateGenericTableWithoutSplaying(PRTL_GENERIC_TABLE,PVOID*);
 void *    WINAPI RtlEnumerateGenericTableWithoutSplayingAvl(PRTL_AVL_TABLE,PVOID*);
 BOOLEAN   WINAPI RtlEqualString(const STRING*,const STRING*,BOOLEAN);