mDNSMacOSX/BonjourTop/source/LLRBTree.h

   1 //
   2 //  LLRBTree.h
   3 //  TestTB
   4 //
   5 //  Created by Terrin Eager on 7/9/12.
   6 //
   7 //  based on rbtree.h but converted to C++
   8 //  ll from http://www.cs.princeton.edu/~rs/talks/LLRB/RedBlack.pdf
   9
  10 #ifndef __TestTB__LLRBTree__
  11 #define __TestTB__LLRBTree__
  12
  13 #include <iostream>
  14 #include "bjtypes.h"
  15 #include <sys/socket.h>
  16 #include "bjstring.h"
  17 #include "bjIPAddr.h"
  18
  19 template <class KeyType>
  20 class CRBNode
  21 {
  22 public:
  23     CRBNode() {m_bIsRed = true; m_rbLeft = m_rbRight = NULL;};
  24     virtual ~CRBNode();
  25
  26     inline virtual BJ_COMPARE Compare(KeyType* pKey) = 0; // Key values are equal
  27
  28     inline virtual void CopyNode(CRBNode* pSource) = 0;
  29     inline virtual void Init()=0;
  30     inline virtual void Clear()=0;
  31
  32     CRBNode* GetMinNode();
  33     CRBNode* GetMaxNode();
  34
  35
  36     inline CRBNode* RotateNodeLeft();
  37     inline CRBNode* RotateNodeRight();
  38
  39
  40     CRBNode* AddRecord(CRBNode* pNewRecord);
  41
  42     void FlipColor();
  43
  44     BJ_UINT64 GetCount();
  45
  46     CRBNode* Fixup();
  47
  48     CRBNode* MoveRedLeft();
  49     CRBNode* MoveRedRight();
  50
  51     void  CallBack(int(*callback)(const void*, const void*),void* pParam);
  52
  53
  54     bool  m_bIsRed;
  55
  56 //protected:
  57     KeyType  m_Key;
  58     CRBNode* m_rbLeft;
  59     CRBNode* m_rbRight;
  60
  61
  62 };
  63
  64 template <class KeyType, class NodeType>
  65 class CLLRBTree
  66 {
  67
  68 public:
  69     CLLRBTree();
  70     virtual ~CLLRBTree() { if (m_Root) delete m_Root;};
  71
  72     NodeType* Find(KeyType* pKey);
  73
  74     NodeType* FindwithAddRecord(KeyType* pKey);
  75     NodeType* AddRecord(KeyType* pKey);
  76     void RemoveRecord(KeyType* pKey);
  77
  78     NodeType* GetRoot() { return m_Root;};
  79     void ClearAll() { delete m_Root; m_Root = NULL;};
  80
  81     BJ_UINT64 GetCount();
  82
  83
  84     NodeType* GetMinNode();
  85     NodeType* GetMaxNode();
  86
  87     NodeType* deleteMin(NodeType* pRecord);
  88
  89
  90 private:
  91     NodeType* RemoveRecord(NodeType* pRecord,KeyType* pKey);
  92
  93     virtual NodeType* newNode(KeyType* pKey) { return new NodeType(pKey);}
  94     virtual void freeNode(NodeType * pNode){ delete pNode;};
  95
  96
  97     NodeType* m_Root;
  98
  99
 100 };
 101 /////////////////
 102
 103
 104 template<class KeyType>
 105 CRBNode<KeyType>::~CRBNode()
 106 {
 107     if (m_rbLeft)
 108         delete m_rbLeft;
 109     if (m_rbRight)
 110         delete m_rbRight;
 111
 112 }
 113
 114
 115
 116
 117 template<class KeyType>
 118 CRBNode<KeyType>* CRBNode<KeyType>::RotateNodeLeft()
 119 {
 120     CRBNode<KeyType>* pTemp = m_rbRight;
 121     m_rbRight = pTemp->m_rbLeft;
 122     pTemp->m_rbLeft = this;
 123     pTemp->m_bIsRed = pTemp->m_rbLeft->m_bIsRed;
 124     pTemp->m_rbLeft->m_bIsRed = 1;
 125
 126     return pTemp;
 127 }
 128
 129
 130 template<class KeyType>
 131 CRBNode<KeyType>* CRBNode<KeyType>::RotateNodeRight()
 132 {
 133     CRBNode<KeyType>* pTemp = m_rbLeft;
 134     m_rbLeft = pTemp->m_rbRight;
 135     pTemp->m_rbRight = this;
 136     pTemp->m_bIsRed = pTemp->m_rbRight->m_bIsRed;
 137     pTemp->m_rbRight->m_bIsRed = 1;
 138
 139     return pTemp;
 140 }
 141
 142 template<class KeyType>
 143 BJ_UINT64 CRBNode<KeyType>::GetCount()
 144 {
 145     BJ_UINT64 Num = 1;
 146     if (m_rbLeft)
 147         Num += m_rbLeft->GetCount();
 148     if (m_rbRight)
 149         Num += m_rbRight->GetCount();
 150
 151     return Num;
 152
 153 }
 154
 155 template<class KeyType>
 156 void CRBNode<KeyType>::FlipColor()
 157 {
 158     m_bIsRed = !m_bIsRed;
 159     if (m_rbLeft)
 160         m_rbLeft->m_bIsRed = !m_rbLeft->m_bIsRed;
 161     if (m_rbRight)
 162         m_rbRight->m_bIsRed = !m_rbRight->m_bIsRed;
 163
 164 }
 165
 166 template<class KeyType>
 167 CRBNode<KeyType>* CRBNode<KeyType>::Fixup()
 168 {
 169     // fix the tree balance
 170     CRBNode<KeyType>* pNode = this;
 171
 172     if (m_rbRight && m_rbRight->m_bIsRed) // fix right leaning reds on the way up
 173         pNode = RotateNodeLeft();
 174
 175     if (pNode && pNode->m_rbLeft && pNode->m_rbLeft->m_bIsRed && pNode->m_rbLeft->m_rbLeft && pNode->m_rbLeft->m_rbLeft->m_bIsRed) // fix two reds in a row on the way up
 176         pNode = RotateNodeRight();
 177
 178     if (pNode && pNode->m_rbRight && pNode->m_rbRight->m_bIsRed && pNode->m_rbLeft && pNode->m_rbLeft->m_bIsRed) //split 4-nodes on the way up
 179         pNode->FlipColor();
 180
 181     return pNode;
 182 }
 183
 184 template<class KeyType>
 185 CRBNode<KeyType>* CRBNode<KeyType>::MoveRedLeft()
 186 {
 187     CRBNode* pNode = this;
 188     FlipColor();
 189
 190     if (m_rbRight && m_rbRight->m_rbLeft && m_rbRight->m_rbLeft->m_bIsRed)
 191     {
 192         m_rbRight = m_rbRight->RotateNodeRight();
 193         pNode = RotateNodeLeft();
 194         if (pNode)
 195             pNode->FlipColor();
 196     }
 197     return pNode;
 198 }
 199
 200 template<class KeyType>
 201 CRBNode<KeyType>* CRBNode<KeyType>::MoveRedRight()
 202 {
 203     CRBNode* pNode = this;
 204     FlipColor();
 205
 206     if (m_rbLeft && m_rbLeft->m_rbLeft && m_rbLeft->m_rbLeft->m_bIsRed)
 207     {
 208         pNode = RotateNodeRight();
 209         if (pNode)
 210             pNode->FlipColor();
 211     }
 212
 213     return pNode;
 214 }
 215 template<class KeyType>
 216 CRBNode<KeyType>* CRBNode<KeyType>::AddRecord(CRBNode* pNewRecord)
 217 {
 218
 219     switch (Compare(&pNewRecord->m_Key))
 220     {
 221         case BJ_GT:
 222             if (m_rbRight)
 223                 m_rbRight = m_rbRight->AddRecord(pNewRecord);
 224             else
 225                 m_rbRight = pNewRecord;
 226
 227             break;
 228         case BJ_LT:
 229             if (m_rbLeft)
 230                 m_rbLeft = m_rbLeft->AddRecord(pNewRecord);
 231             else
 232                 m_rbLeft = pNewRecord;
 233
 234             break;
 235         default: // equal
 236             pNewRecord->m_bIsRed = false;
 237             pNewRecord->m_rbLeft = m_rbLeft;
 238             m_rbLeft = pNewRecord;
 239             return this;
 240     };
 241
 242     // fix the tree balance
 243     CRBNode* pRecord = this;
 244
 245     // fix the tree balance
 246
 247     if (pRecord && pRecord->m_rbRight && pRecord->m_rbRight->m_bIsRed) // fix right leaning reds on the way up
 248         pRecord = pRecord->RotateNodeLeft();
 249
 250     if (pRecord && pRecord->m_rbLeft && pRecord->m_rbLeft->m_bIsRed && pRecord->m_rbLeft->m_rbLeft && pRecord->m_rbLeft->m_rbLeft->m_bIsRed) // fix two reds in a row on the way up
 251         pRecord = pRecord->RotateNodeRight();
 252
 253     if (pRecord && pRecord->m_rbRight && pRecord->m_rbRight->m_bIsRed && pRecord->m_rbLeft && pRecord->m_rbLeft->m_bIsRed) //split 4-nodes on the way up
 254         pRecord->FlipColor();
 255
 256
 257     return pRecord;
 258 }
 259
 260
 261
 262 template<class KeyType>
 263 CRBNode<KeyType>* CRBNode<KeyType>::GetMinNode()
 264 {
 265     CRBNode* pRecord = this;
 266     while (pRecord && pRecord->m_rbLeft)
 267         pRecord = pRecord->m_rbLeft;
 268
 269     return pRecord;
 270 }
 271 template<class KeyType>
 272 CRBNode<KeyType>* CRBNode<KeyType>::GetMaxNode()
 273 {
 274     CRBNode* pRecord = this;
 275     while (pRecord && pRecord->m_rbRight)
 276         pRecord = pRecord->m_rbRight;
 277
 278     return pRecord;
 279 }
 280
 281 template<class KeyType>
 282 void  CRBNode<KeyType>::CallBack(int(*callback)(const void*, const void*),void* pParam)
 283 {
 284
 285     if (m_rbLeft)
 286         m_rbLeft->CallBack(callback,pParam);
 287
 288     callback(this,pParam);
 289
 290     if (m_rbRight)
 291         m_rbRight->CallBack(callback,pParam);
 292
 293
 294 }
 295
 296
 297
 298 ///////////
 299 template<class KeyType,class NodeType>
 300 CLLRBTree<KeyType,NodeType>::CLLRBTree()
 301 {
 302     m_Root = NULL;
 303 }
 304
 305 template<class KeyType,class NodeType>
 306 NodeType* CLLRBTree<KeyType,NodeType>::Find(KeyType* pKey)
 307 {
 308
 309     CRBNode<KeyType>* pNode = m_Root;
 310
 311     while (pNode)
 312     {
 313         switch (pNode->Compare(pKey))
 314         {
 315             case BJ_GT:
 316                 pNode = pNode->m_rbRight;
 317                 break;
 318             case BJ_LT:
 319                 pNode = pNode->m_rbLeft;
 320                 break;
 321             default:
 322                 return (NodeType*)pNode;
 323                 break;
 324         }
 325     }
 326
 327     return NULL;
 328
 329 }
 330
 331 template<class KeyType,class NodeType>
 332 NodeType* CLLRBTree<KeyType,NodeType>::AddRecord(KeyType* pKey)
 333 {
 334     NodeType* pRecord = newNode(pKey);
 335     if (m_Root)
 336         m_Root = (NodeType*) m_Root->AddRecord(pRecord);
 337     else
 338         m_Root = pRecord;
 339
 340     if (m_Root)
 341         m_Root->m_bIsRed = false;
 342
 343     return pRecord;
 344 }
 345
 346 template<class KeyType,class NodeType>
 347 NodeType* CLLRBTree<KeyType,NodeType>::FindwithAddRecord(KeyType* pKey)
 348 {
 349     NodeType* pRecord = NULL;
 350
 351     pRecord = Find(pKey);
 352
 353     if (pRecord == NULL)
 354         pRecord = AddRecord(pKey);
 355
 356     return pRecord;
 357 }
 358
 359 template<class KeyType,class NodeType>
 360 void CLLRBTree<KeyType,NodeType>::RemoveRecord(KeyType* pKey)
 361 {
 362     m_Root = RemoveRecord(m_Root,pKey);
 363 }
 364
 365 template<class KeyType,class NodeType>
 366 NodeType* CLLRBTree<KeyType,NodeType>::deleteMin(NodeType* pRecord)
 367 {
 368     if (pRecord->m_rbLeft == NULL)
 369     {
 370         freeNode(pRecord);
 371         return NULL;
 372     }
 373
 374     if (!(pRecord->m_rbLeft && pRecord->m_rbLeft->m_bIsRed)  && !(pRecord->m_rbLeft &&  pRecord->m_rbLeft->m_rbLeft &&pRecord->m_rbLeft->m_rbLeft->m_bIsRed))
 375         pRecord = (NodeType*)pRecord->MoveRedLeft();
 376
 377     pRecord->m_rbLeft = deleteMin((NodeType*)pRecord->m_rbLeft);
 378
 379     return (NodeType*)pRecord->Fixup();
 380 }
 381
 382 template<class KeyType,class NodeType>
 383 NodeType* CLLRBTree<KeyType,NodeType>::RemoveRecord(NodeType* pRecord,KeyType* pKey)
 384 {
 385     NodeType* pTempRecord = NULL;
 386
 387     if (pRecord == NULL)
 388         return NULL;
 389
 390     if (pRecord->Compare(pKey) == BJ_LT)
 391     {
 392         if (!(pRecord->m_rbLeft &&pRecord->m_rbLeft->m_bIsRed)  && !(pRecord->m_rbLeft &&  pRecord->m_rbLeft->m_rbLeft &&pRecord->m_rbLeft->m_rbLeft->m_bIsRed))
 393             pRecord->MoveRedLeft();
 394         pRecord = RemoveRecord((NodeType*)pRecord->m_rbLeft, pKey);
 395     }
 396     else
 397     {
 398         if (pRecord->m_rbLeft &&pRecord->m_rbLeft->m_bIsRed)
 399             pRecord->RotateNodeRight();
 400
 401         if(pRecord->Compare(pKey) == BJ_EQUAL && pRecord->m_rbRight == NULL)
 402         {
 403             freeNode(pRecord);
 404             return NULL;
 405         }
 406
 407         if (!(pRecord->m_rbRight && pRecord->m_rbRight->m_bIsRed) && !(pRecord->m_rbRight && pRecord->m_rbRight->m_rbLeft && pRecord->m_rbRight->m_rbLeft->m_bIsRed))
 408             pRecord = (NodeType*)pRecord->MoveRedRight();
 409
 410         if (pRecord->Compare(pKey) == BJ_EQUAL)
 411         {
 412             pTempRecord = (NodeType*)pRecord->GetMinNode();
 413             pRecord->CopyNode(pTempRecord);
 414             pRecord->m_rbRight = deleteMin((NodeType*)pRecord->m_rbRight);
 415         }
 416         else
 417         {
 418             pRecord->m_rbRight = RemoveRecord((NodeType*)pRecord->m_rbRight, pKey);
 419         }
 420     }
 421     return pRecord?(NodeType*)pRecord->Fixup():NULL;
 422 }
 423
 424
 425
 426
 427 template<class KeyType,class NodeType>
 428 BJ_UINT64 CLLRBTree<KeyType,NodeType>::GetCount()
 429 {
 430     if (m_Root)
 431         return m_Root->GetCount();
 432     else
 433         return 0;
 434 }
 435
 436 template<class KeyType,class NodeType>
 437 NodeType* CLLRBTree<KeyType,NodeType>::GetMinNode()
 438 {
 439     if (m_Root)
 440         return m_Root->GetMinNode();
 441     else
 442         return NULL;
 443
 444 }
 445
 446 template<class KeyType,class NodeType>
 447 NodeType* CLLRBTree<KeyType,NodeType>::GetMaxNode()
 448 {
 449     if (m_Root)
 450         return m_Root->GetMaxNode();
 451     else
 452         return NULL;
 453
 454 }
 455
 456
 457
 458
 459
 460 #endif /* defined(__TestTB__LLRBTree__) */