From b8230d5a3e16ca0512579ad68ff3a668f80daaa5 Mon Sep 17 00:00:00 2001
From: Hongze Cheng <hzcheng@taosdata.com>
Date: Thu, 25 Aug 2022 10:04:28 +0800
Subject: [PATCH] more code

---
 include/util/trbtree.h           |  5 ++--
 source/util/src/trbtree.c        | 49 ++++++++++++++++++++++++--------
 source/util/test/trbtreeTest.cpp |  2 +-
 3 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/include/util/trbtree.h b/include/util/trbtree.h
index d6a237fb77..d666cc11d1 100644
--- a/include/util/trbtree.h
+++ b/include/util/trbtree.h
@@ -38,8 +38,8 @@ SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey);
 SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey);
 
 // SRBTreeIter =============================================
-#define tRBTreeIterCreate(tree) \
-  (SRBTreeIter) { .pTree = (tree), .pNode = (tree)->minNode }
+#define tRBTreeIterCreate(tree, descend) \
+  (SRBTreeIter) { .des = (descend), .pTree = (tree), .pNode = (descend) ? (tree)->maxNode : (tree)->minNode }
 
 SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter);
 
@@ -60,6 +60,7 @@ struct SRBTree {
 };
 
 struct SRBTreeIter {
+  int8_t       des;
   SRBTree     *pTree;
   SRBTreeNode *pNode;
 };
diff --git a/source/util/src/trbtree.c b/source/util/src/trbtree.c
index 3c3bad7340..14b841dc50 100644
--- a/source/util/src/trbtree.c
+++ b/source/util/src/trbtree.c
@@ -253,23 +253,48 @@ SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter) {
   SRBTree     *pTree = pIter->pTree;
 
   if (pIter->pNode) {
-    if (pIter->pNode->right) {
-      pIter->pNode = pIter->pNode->right;
-      while (pIter->pNode->left) {
+    if (pIter->des) {
+      // descend
+      if (pIter->pNode->left) {
         pIter->pNode = pIter->pNode->left;
+        while (pIter->pNode->right) {
+          pIter->pNode = pIter->pNode->right;
+        }
+      } else {
+        while (true) {
+          if (pIter->pNode->parent) {
+            if (pIter->pNode == pIter->pNode->parent->right) {
+              pIter->pNode = pIter->pNode->parent;
+              break;
+            } else {
+              pIter->pNode = pIter->pNode->parent;
+            }
+          } else {
+            pIter->pNode = NULL;
+            break;
+          }
+        }
       }
     } else {
-      while (true) {
-        if (pIter->pNode->parent) {
-          if (pIter->pNode == pIter->pNode->parent->left) {
-            pIter->pNode = pIter->pNode->parent;
-            break;
+      // ascend
+      if (pIter->pNode->right) {
+        pIter->pNode = pIter->pNode->right;
+        while (pIter->pNode->left) {
+          pIter->pNode = pIter->pNode->left;
+        }
+      } else {
+        while (true) {
+          if (pIter->pNode->parent) {
+            if (pIter->pNode == pIter->pNode->parent->left) {
+              pIter->pNode = pIter->pNode->parent;
+              break;
+            } else {
+              pIter->pNode = pIter->pNode->parent;
+            }
           } else {
-            pIter->pNode = pIter->pNode->parent;
+            pIter->pNode = NULL;
+            break;
           }
-        } else {
-          pIter->pNode = NULL;
-          break;
         }
       }
     }
diff --git a/source/util/test/trbtreeTest.cpp b/source/util/test/trbtreeTest.cpp
index 4ff990a46d..be62740949 100644
--- a/source/util/test/trbtreeTest.cpp
+++ b/source/util/test/trbtreeTest.cpp
@@ -25,7 +25,7 @@ TEST(trbtreeTest, rbtree_test1) {
     tRBTreePut(&rt, pNode);
   }
 
-  SRBTreeIter  rti = tRBTreeIterCreate(&rt);
+  SRBTreeIter  rti = tRBTreeIterCreate(&rt, 0);
   SRBTreeNode *pNode = tRBTreeIterNext(&rti);
   int          la = 0;
   while (pNode) {
-- 
GitLab