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

---
 include/util/trbtree.h           |   6 +-
 source/util/src/trbtree.c        | 120 ++++++++++++++++---------------
 source/util/test/trbtreeTest.cpp |   3 +-
 3 files changed, 66 insertions(+), 63 deletions(-)

diff --git a/include/util/trbtree.h b/include/util/trbtree.h
index d666cc11d1..1e34839cbe 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, descend) \
-  (SRBTreeIter) { .des = (descend), .pTree = (tree), .pNode = (descend) ? (tree)->maxNode : (tree)->minNode }
+#define tRBTreeIterCreate(tree, ascend) \
+  (SRBTreeIter) { .asc = (ascend), .pTree = (tree), .pNode = (ascend) ? (tree)->minNode : (tree)->maxNode }
 
 SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter);
 
@@ -60,7 +60,7 @@ struct SRBTree {
 };
 
 struct SRBTreeIter {
-  int8_t       des;
+  int8_t       asc;
   SRBTree     *pTree;
   SRBTreeNode *pNode;
 };
diff --git a/source/util/src/trbtree.c b/source/util/src/trbtree.c
index 14b841dc50..7364c129bb 100644
--- a/source/util/src/trbtree.c
+++ b/source/util/src/trbtree.c
@@ -60,6 +60,55 @@ static void tRBTreeRotateRight(SRBTree *pTree, SRBTreeNode *pNode) {
   pNode->parent = left;
 }
 
+static SRBTreeNode *tRBTreeSuccessor(SRBTreeNode *pNode) {
+  if (pNode->right) {
+    pNode = pNode->right;
+    while (pNode->left) {
+      pNode = pNode->left;
+    }
+  } else {
+    while (true) {
+      if (pNode->parent) {
+        if (pNode == pNode->parent->left) {
+          pNode = pNode->parent;
+          break;
+        } else {
+          pNode = pNode->parent;
+        }
+      } else {
+        pNode = NULL;
+        break;
+      }
+    }
+  }
+
+  return pNode;
+}
+
+static SRBTreeNode *tRBTreePredecessor(SRBTreeNode *pNode) {
+  if (pNode->left) {
+    pNode = pNode->left;
+    while (pNode->right) {
+      pNode = pNode->right;
+    }
+  } else {
+    while (true) {
+      if (pNode->parent) {
+        if (pNode == pNode->parent->right) {
+          pNode = pNode->parent;
+          break;
+        } else {
+          pNode = pNode->parent;
+        }
+      } else {
+        pNode = NULL;
+        break;
+      }
+    }
+  }
+  return NULL;
+}
+
 SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew) {
   pNew->left = NULL;
   pNew->right = NULL;
@@ -153,8 +202,12 @@ SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew) {
 
 void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode) {
   // update min/max node
-  if (pTree->minNode == pNode) pTree->minNode = pNode->parent;
-  if (pTree->maxNode == pNode) pTree->maxNode = pNode->parent;
+  if (pTree->minNode == pNode) {
+    pTree->minNode = tRBTreeSuccessor(pTree->minNode);
+  }
+  if (pTree->maxNode == pNode) {
+    pTree->maxNode = tRBTreePredecessor(pTree->maxNode);
+  }
 
   // drop impl
   if (pNode->left == NULL) {
@@ -208,19 +261,7 @@ void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode) {
 }
 
 SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey) {
-  SRBTreeNode *pNode = pTree->rootNode;
-
-  while (pNode) {
-    int32_t c = pTree->cmprFn(pKey, pNode->payload);
-
-    if (c < 0) {
-      pNode = pNode->left;
-    } else if (c > 0) {
-      pNode = pNode->right;
-    } else {
-      break;
-    }
-  }
+  SRBTreeNode *pNode = tRBTreeGet(pTree, pKey);
 
   if (pNode) {
     tRBTreeDrop(pTree, pNode);
@@ -250,53 +291,14 @@ SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey) {
 // SRBTreeIter ================================================
 SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter) {
   SRBTreeNode *pNode = pIter->pNode;
-  SRBTree     *pTree = pIter->pTree;
 
   if (pIter->pNode) {
-    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 {
+    if (pIter->asc) {
       // 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 = NULL;
-            break;
-          }
-        }
-      }
+      pIter->pNode = tRBTreeSuccessor(pIter->pNode);
+    } else {
+      // descend
+      pIter->pNode = tRBTreePredecessor(pIter->pNode);
     }
   }
 
diff --git a/source/util/test/trbtreeTest.cpp b/source/util/test/trbtreeTest.cpp
index be62740949..1e0768c32f 100644
--- a/source/util/test/trbtreeTest.cpp
+++ b/source/util/test/trbtreeTest.cpp
@@ -25,12 +25,13 @@ TEST(trbtreeTest, rbtree_test1) {
     tRBTreePut(&rt, pNode);
   }
 
-  SRBTreeIter  rti = tRBTreeIterCreate(&rt, 0);
+  SRBTreeIter  rti = tRBTreeIterCreate(&rt, 1);
   SRBTreeNode *pNode = tRBTreeIterNext(&rti);
   int          la = 0;
   while (pNode) {
     GTEST_ASSERT_GT(*(int *)pNode->payload, la);
     la = *(int *)pNode->payload;
+    // printf("%d\n", la);
     pNode = tRBTreeIterNext(&rti);
   }
 }
\ No newline at end of file
-- 
GitLab