diff --git a/include/util/trbtree.h b/include/util/trbtree.h index 9d5cc606665e7bf9131eb7868e1ab63746fe8ba0..060d5a791840ee4e74049ce18fe5a2f8ae7459f8 100644 --- a/include/util/trbtree.h +++ b/include/util/trbtree.h @@ -29,11 +29,12 @@ typedef struct SRBTreeIter SRBTreeIter; typedef int32_t (*tRBTreeCmprFn)(const void *, const void *); // SRBTree ============================================= -#define tRBTreeCreate(compare) \ - (SRBTree) { .cmprFn = (compare), .root = NULL, .min = NULL, .max = NULL } +#define tRBTreeMin(T) ((T)->min == ((T)->NIL) ? NULL : (T)->min) +#define tRBTreeMax(T) ((T)->max == ((T)->NIL) ? NULL : (T)->max) -SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *pNew); -void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *pNode); +void tRBTreeCreate(SRBTree *pTree, tRBTreeCmprFn cmprFn); +SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *z); +void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *z); SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey); SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey); @@ -50,7 +51,7 @@ struct SRBTreeNode { SRBTreeNode *parent; SRBTreeNode *left; SRBTreeNode *right; - uint8_t payload[]; + uint8_t payload[0]; }; struct SRBTree { @@ -58,6 +59,8 @@ struct SRBTree { SRBTreeNode *root; SRBTreeNode *min; SRBTreeNode *max; + SRBTreeNode *NIL; + SRBTreeNode NILNODE; }; struct SRBTreeIter { diff --git a/source/dnode/vnode/src/tsdb/tsdbMerge.c b/source/dnode/vnode/src/tsdb/tsdbMerge.c index 3a078131a1c9d1e61831cc5fe0a1655a13dceddb..186db38461c3e5cacef9d5c89045574d09848c3f 100644 --- a/source/dnode/vnode/src/tsdb/tsdbMerge.c +++ b/source/dnode/vnode/src/tsdb/tsdbMerge.c @@ -55,7 +55,7 @@ static int32_t tRowInfoCmprFn(const void *p1, const void *p2) { static void tDataMergerInit(SDataMerger *pMerger, SArray *aNodeP) { pMerger->pNode = NULL; - pMerger->rbt = tRBTreeCreate(tRowInfoCmprFn); + tRBTreeCreate(&pMerger->rbt, tRowInfoCmprFn); for (int32_t iNode = 0; iNode < taosArrayGetSize(aNodeP); iNode++) { SRBTreeNode *pNode = (SRBTreeNode *)taosArrayGetP(aNodeP, iNode); @@ -90,8 +90,9 @@ static int32_t tDataMergeNext(SDataMerger *pMerger, SRowInfo **ppInfo) { } } - if (pMerger->pNode && pMerger->rbt.min) { - int32_t c = tRowInfoCmprFn(pMerger->pNode->payload, pMerger->rbt.min->payload); + SRBTreeNode *pMinNode = tRBTreeMin(&pMerger->rbt); + if (pMerger->pNode && pMinNode) { + int32_t c = tRowInfoCmprFn(pMerger->pNode->payload, pMinNode->payload); if (c > 0) { pMerger->pNode = tRBTreePut(&pMerger->rbt, pMerger->pNode); ASSERT(pMerger->pNode); @@ -103,7 +104,7 @@ static int32_t tDataMergeNext(SDataMerger *pMerger, SRowInfo **ppInfo) { } if (pMerger->pNode == NULL) { - pMerger->pNode = pMerger->rbt.min; + pMerger->pNode = tRBTreeMin(&pMerger->rbt); if (pMerger->pNode) { tRBTreeDrop(&pMerger->rbt, pMerger->pNode); } @@ -155,7 +156,7 @@ static int32_t tsdbMergeFileDataStart(STsdbMerger *pMerger, SDFileSet *pSet) { if (code) goto _err; pMerger->dReader.merger.pNode = NULL; - pMerger->dReader.merger.rbt = tRBTreeCreate(tRowInfoCmprFn); + tRBTreeCreate(&pMerger->dReader.merger.rbt, tRowInfoCmprFn); for (int8_t iLast = 0; iLast < pSet->nLastF; iLast++) { SRBTreeNode *pNode = (SRBTreeNode *)taosMemoryCalloc(1, sizeof(*pNode) + sizeof(SLDataIter)); if (pNode == NULL) { diff --git a/source/util/src/trbtree.c b/source/util/src/trbtree.c index 02c5113904401910a42c962b48d89c5be01b64bf..ba011de0d7bfe020153c2a47b63e9c3fa9b8a496 100644 --- a/source/util/src/trbtree.c +++ b/source/util/src/trbtree.c @@ -15,17 +15,14 @@ #include "trbtree.h" -#define RBTREE_NODE_COLOR(N) ((N) ? (N)->color : BLACK) - -// SRBTree ================================================ static void tRBTreeRotateLeft(SRBTree *pTree, SRBTreeNode *x) { SRBTreeNode *y = x->right; x->right = y->left; - if (y->left) { + if (y->left != pTree->NIL) { y->left->parent = x; } y->parent = x->parent; - if (x->parent == NULL) { + if (x->parent == pTree->NIL) { pTree->root = y; } else if (x == x->parent->left) { x->parent->left = y; @@ -39,269 +36,263 @@ static void tRBTreeRotateLeft(SRBTree *pTree, SRBTreeNode *x) { static void tRBTreeRotateRight(SRBTree *pTree, SRBTreeNode *x) { SRBTreeNode *y = x->left; x->left = y->right; - if (y->right) { + if (y->right != pTree->NIL) { y->right->parent = x; } y->parent = x->parent; - if (x->parent == NULL) { + if (x->parent == pTree->NIL) { pTree->root = y; - } else if (x == x->parent->left) { - x->parent->left = y; - } else { + } else if (x == x->parent->right) { x->parent->right = y; + } else { + x->parent->left = y; } y->right = x; x->parent = y; } -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; - pNew->color = RED; - - // insert - if (pTree->root == NULL) { - pNew->parent = NULL; - pTree->root = pNew; - } else { - SRBTreeNode *pNode = pTree->root; - while (true) { - ASSERT(pNode); - - int32_t c = pTree->cmprFn(pNew->payload, pNode->payload); - if (c < 0) { - if (pNode->left) { - pNode = pNode->left; - } else { - pNew->parent = pNode; - pNode->left = pNew; - break; - } - } else if (c > 0) { - if (pNode->right) { - pNode = pNode->right; - } else { - pNew->parent = pNode; - pNode->right = pNew; - break; +static void tRBTreePutFix(SRBTree *pTree, SRBTreeNode *z) { + while (z->parent->color == RED) { + if (z->parent == z->parent->parent->left) { // z.parent is the left child + + SRBTreeNode *y = z->parent->parent->right; // uncle of z + + if (y->color == RED) { // case 1 + z->parent->color = BLACK; + y->color = BLACK; + z->parent->parent->color = RED; + z = z->parent->parent; + } else { // case2 or case3 + if (z == z->parent->right) { // case2 + z = z->parent; // marked z.parent as new z + tRBTreeRotateLeft(pTree, z); } - } else { - return NULL; + // case3 + z->parent->color = BLACK; // made parent black + z->parent->parent->color = RED; // made parent red + tRBTreeRotateRight(pTree, z->parent->parent); } - } - } - - // fix - SRBTreeNode *pNode = pNew; - while (pNode->parent && pNode->parent->color == RED) { - SRBTreeNode *p = pNode->parent; - SRBTreeNode *g = p->parent; - - if (p == g->left) { - SRBTreeNode *u = g->right; - - if (RBTREE_NODE_COLOR(u) == RED) { - p->color = BLACK; - u->color = BLACK; - g->color = RED; - pNode = g; - } else { - if (pNode == p->right) { - pNode = p; - tRBTreeRotateLeft(pTree, pNode); - } - pNode->parent->color = BLACK; - pNode->parent->parent->color = RED; - tRBTreeRotateRight(pTree, pNode->parent->parent); - } - } else { - SRBTreeNode *u = g->left; - - if (RBTREE_NODE_COLOR(u) == RED) { - p->color = BLACK; - u->color = BLACK; - g->color = RED; + } else { // z.parent is the right child + SRBTreeNode *y = z->parent->parent->left; // uncle of z + + if (y->color == RED) { + z->parent->color = BLACK; + y->color = BLACK; + z->parent->parent->color = RED; + z = z->parent->parent; } else { - if (pNode == p->left) { - pNode = p; - tRBTreeRotateRight(pTree, pNode); + if (z == z->parent->left) { + z = z->parent; // marked z.parent as new z + tRBTreeRotateRight(pTree, z); } - pNode->parent->color = BLACK; - pNode->parent->parent->color = RED; - tRBTreeRotateLeft(pTree, pNode->parent->parent); + z->parent->color = BLACK; // made parent black + z->parent->parent->color = RED; // made parent red + tRBTreeRotateLeft(pTree, z->parent->parent); } } } pTree->root->color = BLACK; - - // update min/max node - if (pTree->min == NULL || pTree->cmprFn(pTree->min->payload, pNew->payload) > 0) { - pTree->min = pNew; - } - if (pTree->max == NULL || pTree->cmprFn(pTree->max->payload, pNew->payload) < 0) { - pTree->max = pNew; - } - - return pNew; } static void tRBTreeTransplant(SRBTree *pTree, SRBTreeNode *u, SRBTreeNode *v) { - if (u->parent == NULL) { + if (u->parent == pTree->NIL) pTree->root = v; - } else if (u == u->parent->left) { + else if (u == u->parent->left) u->parent->left = v; - } else { + else u->parent->right = v; - } - if (v) { - v->parent = u->parent; - } + v->parent = u->parent; } -static void tRBTreeDropFixup(SRBTree *t, SRBTreeNode *x) { - while (x != t->root && x->color == BLACK) { +static void tRBTreeDropFix(SRBTree *pTree, SRBTreeNode *x) { + while (x != pTree->root && x->color == BLACK) { if (x == x->parent->left) { SRBTreeNode *w = x->parent->right; - if (RBTREE_NODE_COLOR(w) == RED) { + if (w->color == RED) { w->color = BLACK; x->parent->color = RED; - tRBTreeRotateLeft(t, x->parent); + tRBTreeRotateLeft(pTree, x->parent); w = x->parent->right; } - if (RBTREE_NODE_COLOR(w->left) == BLACK && RBTREE_NODE_COLOR(w->right) == BLACK) { + if (w->left->color == BLACK && w->right->color == BLACK) { w->color = RED; x = x->parent; } else { - if (RBTREE_NODE_COLOR(w->right) == BLACK) { + if (w->right->color == BLACK) { w->left->color = BLACK; w->color = RED; - tRBTreeRotateRight(t, w); + tRBTreeRotateRight(pTree, w); w = x->parent->right; } w->color = x->parent->color; x->parent->color = BLACK; w->right->color = BLACK; - tRBTreeRotateLeft(t, x->parent); - x = t->root; + tRBTreeRotateLeft(pTree, x->parent); + x = pTree->root; } } else { SRBTreeNode *w = x->parent->left; - if (RBTREE_NODE_COLOR(w) == RED) { + if (w->color == RED) { w->color = BLACK; x->parent->color = RED; - tRBTreeRotateRight(t, x->parent); + tRBTreeRotateRight(pTree, x->parent); w = x->parent->left; } - if (RBTREE_NODE_COLOR(w->right) == BLACK && RBTREE_NODE_COLOR(w->left) == BLACK) { + if (w->right->color == BLACK && w->left->color == BLACK) { w->color = RED; x = x->parent; } else { - if (RBTREE_NODE_COLOR(w->left) == BLACK) { + if (w->left->color == BLACK) { w->right->color = BLACK; w->color = RED; - tRBTreeRotateLeft(t, w); + tRBTreeRotateLeft(pTree, w); w = x->parent->left; } w->color = x->parent->color; x->parent->color = BLACK; w->left->color = BLACK; - tRBTreeRotateRight(t, x->parent); - x = t->root; + tRBTreeRotateRight(pTree, x->parent); + x = pTree->root; } } } x->color = BLACK; } -void tRBTreeDrop(SRBTree *t, SRBTreeNode *z) { +static SRBTreeNode *tRBTreeSuccessor(SRBTree *pTree, SRBTreeNode *pNode) { + if (pNode->right != pTree->NIL) { + pNode = pNode->right; + while (pNode->left != pTree->NIL) { + pNode = pNode->left; + } + } else { + while (true) { + if (pNode->parent == pTree->NIL || pNode == pNode->parent->left) { + pNode = pNode->parent; + break; + } else { + pNode = pNode->parent; + } + } + } + + return pNode; +} + +static SRBTreeNode *tRBTreePredecessor(SRBTree *pTree, SRBTreeNode *pNode) { + if (pNode->left != pTree->NIL) { + pNode = pNode->left; + while (pNode->right != pTree->NIL) { + pNode = pNode->right; + } + } else { + while (true) { + if (pNode->parent == pTree->NIL || pNode == pNode->parent->right) { + pNode = pNode->parent; + break; + } else { + pNode = pNode->parent; + } + } + } + + return pNode; +} + +void tRBTreeCreate(SRBTree *pTree, tRBTreeCmprFn cmprFn) { + pTree->cmprFn = cmprFn; + pTree->NIL = &pTree->NILNODE; + pTree->NIL->color = BLACK; + pTree->NIL->parent = NULL; + pTree->NIL->left = NULL; + pTree->NIL->right = NULL; + pTree->root = pTree->NIL; + pTree->min = pTree->NIL; + pTree->max = pTree->NIL; +} + +SRBTreeNode *tRBTreePut(SRBTree *pTree, SRBTreeNode *z) { + SRBTreeNode *y = pTree->NIL; // variable for the parent of the added node + SRBTreeNode *temp = pTree->root; + + while (temp != pTree->NIL) { + y = temp; + + int32_t c = pTree->cmprFn(z->payload, temp->payload); + if (c < 0) { + temp = temp->left; + } else if (c > 0) { + temp = temp->right; + } else { + return NULL; + } + } + z->parent = y; + + if (y == pTree->NIL) { + pTree->root = z; + } else if (pTree->cmprFn(z->payload, y->payload) < 0) { + y->left = z; + } else { + y->right = z; + } + + z->color = RED; + z->left = pTree->NIL; + z->right = pTree->NIL; + + tRBTreePutFix(pTree, z); + // update min/max node - if (t->min == z) { - t->min = tRBTreeSuccessor(t->min); + if (pTree->min == pTree->NIL || pTree->cmprFn(pTree->min->payload, z->payload) > 0) { + pTree->min = z; } - if (t->max == z) { - t->max = tRBTreePredecessor(t->max); + if (pTree->max == pTree->NIL || pTree->cmprFn(pTree->max->payload, z->payload) < 0) { + pTree->max = z; } + return z; +} - // drop impl +void tRBTreeDrop(SRBTree *pTree, SRBTreeNode *z) { SRBTreeNode *y = z; SRBTreeNode *x; - ECOLOR oColor = y->color; + ECOLOR y_orignal_color = y->color; - if (z->left == NULL) { + // update min/max node + if (pTree->min == z) { + pTree->min = tRBTreeSuccessor(pTree, pTree->min); + } + if (pTree->max == z) { + pTree->max = tRBTreePredecessor(pTree, pTree->max); + } + + // drop impl + if (z->left == pTree->NIL) { x = z->right; - tRBTreeTransplant(t, z, z->right); - } else if (z->right == NULL) { + tRBTreeTransplant(pTree, z, z->right); + } else if (z->right == pTree->NIL) { x = z->left; - tRBTreeTransplant(t, z, z->left); + tRBTreeTransplant(pTree, z, z->left); } else { - y = tRBTreeSuccessor(z); - oColor = y->color; + y = tRBTreeSuccessor(pTree, z); + y_orignal_color = y->color; x = y->right; if (y->parent == z) { x->parent = z; } else { - tRBTreeTransplant(t, y, y->right); + tRBTreeTransplant(pTree, y, y->right); y->right = z->right; y->right->parent = y; } - tRBTreeTransplant(t, z, y); + tRBTreeTransplant(pTree, z, y); y->left = z->left; y->left->parent = y; y->color = z->color; } // fix - if (oColor == BLACK) { - tRBTreeDropFixup(t, x); + if (y_orignal_color == BLACK) { + tRBTreeDropFix(pTree, x); } } @@ -318,7 +309,7 @@ SRBTreeNode *tRBTreeDropByKey(SRBTree *pTree, void *pKey) { SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey) { SRBTreeNode *pNode = pTree->root; - while (pNode) { + while (pNode != pTree->NIL) { int32_t c = pTree->cmprFn(pKey, pNode->payload); if (c < 0) { @@ -330,23 +321,23 @@ SRBTreeNode *tRBTreeGet(SRBTree *pTree, void *pKey) { } } - return pNode; + return (pNode == pTree->NIL) ? NULL : pNode; } // SRBTreeIter ================================================ SRBTreeNode *tRBTreeIterNext(SRBTreeIter *pIter) { SRBTreeNode *pNode = pIter->pNode; - if (pIter->pNode) { + if (pIter->pNode != pIter->pTree->NIL) { if (pIter->asc) { // ascend - pIter->pNode = tRBTreeSuccessor(pIter->pNode); + pIter->pNode = tRBTreeSuccessor(pIter->pTree, pIter->pNode); } else { // descend - pIter->pNode = tRBTreePredecessor(pIter->pNode); + pIter->pNode = tRBTreePredecessor(pIter->pTree, pIter->pNode); } } _exit: - return pNode; + return (pNode == pIter->pTree->NIL) ? NULL : pNode; } \ No newline at end of file diff --git a/source/util/test/trbtreeTest.cpp b/source/util/test/trbtreeTest.cpp index 1e0768c32f017cea055eac1bc1ba9dde9a222a4c..518ffed2cfc87bec40538b980cbb6ea8136da2e9 100644 --- a/source/util/test/trbtreeTest.cpp +++ b/source/util/test/trbtreeTest.cpp @@ -15,8 +15,9 @@ static int32_t tCmprInteger(const void *p1, const void *p2) { } TEST(trbtreeTest, rbtree_test1) { - SRBTree rt = tRBTreeCreate(tCmprInteger); - int a[] = {1, 3, 4, 2, 7, 5, 8}; + SRBTree rt; + tRBTreeCreate(&rt, tCmprInteger); + int a[] = {1, 3, 4, 2, 7, 5, 8}; for (int i = 0; i < sizeof(a) / sizeof(a[0]); i++) { SRBTreeNode *pNode = (SRBTreeNode *)taosMemoryMalloc(sizeof(*pNode) + sizeof(int));