From 04cc96d6d9bc258948fc468e3d2c16c8cba44ecc Mon Sep 17 00:00:00 2001
From: Hongze Cheng <hzcheng@taosdata.com>
Date: Wed, 24 Aug 2022 19:01:27 +0800
Subject: [PATCH] more code

---
 include/util/trbtree.h    | 10 +++--
 source/util/src/trbtree.c | 80 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 83 insertions(+), 7 deletions(-)

diff --git a/include/util/trbtree.h b/include/util/trbtree.h
index 17739b7977..f02b746c34 100644
--- a/include/util/trbtree.h
+++ b/include/util/trbtree.h
@@ -21,7 +21,7 @@ typedef struct SRBTreeIter SRBTreeIter;
 
 typedef int32_t (*tRBTreeCmprFn)(void *, void *);
 
-// SRBTree
+// SRBTree =============================================
 #define tRBTreeCreate(compare) \
   (SRBTree) { .cmprFn = (compare), .root = NULL, .minNode = NULL, .maxNode = NULL }
 
@@ -30,12 +30,13 @@ void         tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode);
 SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey);
 SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey);
 
-// SRBTreeIter
+// SRBTreeIter =============================================
 #define tRBTreeIterCreate(tree) \
-  (SRBTreeIter) { .pTree = (tree) }
+  (SRBTreeIter) { .pTree = (tree), .pNode = (tree)->minNode }
 
 SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter);
 
+// STRUCT =============================================
 struct SRBTreeNode {
   enum { RED, BLACK } color;
   SRBTreeNode *parent;
@@ -52,5 +53,6 @@ struct SRBTree {
 };
 
 struct SRBTreeIter {
-  SRBTree *pTree;
+  SRBTree     *pTree;
+  SRBTreeNode *pNode;
 };
diff --git a/source/util/src/trbtree.c b/source/util/src/trbtree.c
index e21dcacc37..5b8f6c9d96 100644
--- a/source/util/src/trbtree.c
+++ b/source/util/src/trbtree.c
@@ -17,7 +17,7 @@
 
 #define RBTREE_NODE_COLOR(N) ((N) ? (N)->color : BLACK)
 
-// APIs ================================================
+// SRBTree ================================================
 static void tRBTreeRotateLeft(SRBTree *pTree, SRBTreeNode *pNode) {
   SRBTreeNode *right = pNode->right;
 
@@ -138,13 +138,73 @@ SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew) {
       }
     }
   }
-
   pTree->rootNode->color = BLACK;
+
+  // update min/max node
+  if (pTree->minNode == NULL || pTree->cmprFn(pTree->minNode->payload, pNew->payload) > 0) {
+    pTree->minNode = pNew;
+  }
+  if (pTree->maxNode == NULL || pTree->cmprFn(pTree->maxNode->payload, pNew->payload) < 0) {
+    pTree->maxNode = pNew;
+  }
+
   return pNew;
 }
 
 void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode) {
-  // TODO
+  // update min/max node
+  if (pTree->minNode == pNode) pTree->minNode = pNode->parent;
+  if (pTree->maxNode == pNode) pTree->maxNode = pNode->parent;
+
+  // drop impl
+  if (pNode->left == NULL) {
+    if (pNode->parent) {
+      if (pNode == pNode->parent->left) {
+        pNode->parent->left = pNode->right;
+      } else {
+        pNode->parent->right = pNode->right;
+      }
+    } else {
+      pTree->rootNode = pNode->right;
+    }
+
+    if (pNode->right) {
+      pNode->right->parent = pNode->parent;
+    }
+  } else if (pNode->right == NULL) {
+    if (pNode->parent) {
+      if (pNode == pNode->parent->left) {
+        pNode->parent->left = pNode->left;
+      } else {
+        pNode->parent->right = pNode->left;
+      }
+    } else {
+      pTree->rootNode = pNode->left;
+    }
+
+    if (pNode->left) {
+      pNode->left->parent = pNode->parent;
+    }
+  } else {
+    // TODO
+    SRBTreeNode *pSuccessorNode = pNode->right;
+    while (pSuccessorNode->left) {
+      pSuccessorNode = pSuccessorNode->left;
+    }
+
+    pSuccessorNode->parent->left = NULL;  // todo: not correct here
+
+    pSuccessorNode->parent = pNode->parent;
+    pSuccessorNode->left = pNode->left;
+    pSuccessorNode->right = pNode->right;
+    pNode->left->parent = pSuccessorNode;
+    pNode->right->parent = pSuccessorNode;
+  }
+
+  // fix
+  if (pNode->color == BLACK) {
+    // TODO
+  }
 }
 
 SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey) {
@@ -186,3 +246,17 @@ SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey) {
 
   return pNode;
 }
+
+// SRBTreeIter ================================================
+SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter) {
+  SRBTreeNode *pNode = pIter->pNode;
+  SRBTree     *pTree = pIter->pTree;
+
+  if (pIter->pNode) {
+    ASSERT(0);
+    // TODO
+  }
+
+_exit:
+  return pNode;
+}
\ No newline at end of file
-- 
GitLab