rbtree.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937
  1. package rbtree
  2. import (
  3. "math"
  4. "sync"
  5. )
  6. //
  7. // Public definitions
  8. //
  9. // Item is the object stored in each tree node.
  10. type Item struct {
  11. Key uint32
  12. Value uint32
  13. }
  14. // Allocator is the allocator for nodes in a RBTree.
  15. type Allocator struct {
  16. HibernationThreshold int
  17. storage []node
  18. gaps map[uint32]bool
  19. hibernatedStorage [6][]byte
  20. hibernatedLen int
  21. }
  22. // NewAllocator creates a new allocator for RBTree's nodes.
  23. func NewAllocator() *Allocator {
  24. return &Allocator{
  25. storage: []node{},
  26. gaps: map[uint32]bool{},
  27. }
  28. }
  29. // Size returns the currently allocated size.
  30. func (allocator Allocator) Size() int {
  31. return len(allocator.storage)
  32. }
  33. // Used returns the number of nodes contained in the allocator.
  34. func (allocator Allocator) Used() int {
  35. return len(allocator.storage) - len(allocator.gaps)
  36. }
  37. // Clone copies an existing RBTree allocator.
  38. func (allocator *Allocator) Clone() *Allocator {
  39. newAllocator := &Allocator{
  40. storage: make([]node, len(allocator.storage), cap(allocator.storage)),
  41. gaps: map[uint32]bool{},
  42. }
  43. copy(newAllocator.storage, allocator.storage)
  44. for key, val := range allocator.gaps {
  45. newAllocator.gaps[key] = val
  46. }
  47. return newAllocator
  48. }
  49. // Hibernate compresses the allocated memory.
  50. func (allocator *Allocator) Hibernate() {
  51. if allocator.HibernationThreshold == 0 || len(allocator.storage) < allocator.HibernationThreshold {
  52. return
  53. }
  54. allocator.hibernatedLen = len(allocator.storage)
  55. buffers := [6][]uint32{}
  56. for i := 0; i < len(buffers); i++ {
  57. buffers[i] = make([]uint32, len(allocator.storage))
  58. }
  59. // we deinterleave to achieve a better compression ratio
  60. for i, n := range allocator.storage {
  61. buffers[0][i] = n.item.Key
  62. buffers[1][i] = n.item.Value
  63. buffers[2][i] = n.left
  64. buffers[3][i] = n.parent
  65. buffers[4][i] = n.right
  66. if n.color {
  67. buffers[5][i] = 1
  68. }
  69. }
  70. allocator.storage = nil
  71. wg := &sync.WaitGroup{}
  72. wg.Add(len(buffers))
  73. for i, buffer := range buffers {
  74. go func(i int, buffer []uint32) {
  75. allocator.hibernatedStorage[i] = CompressUInt32Slice(buffer)
  76. buffers[i] = nil
  77. wg.Done()
  78. }(i, buffer)
  79. }
  80. wg.Wait()
  81. }
  82. // Boot performs the opposite of Hibernate() - decompresses and restores the allocated memory.
  83. func (allocator *Allocator) Boot() {
  84. if allocator.hibernatedLen == 0 {
  85. // not hibernated
  86. return
  87. }
  88. buffers := [6][]uint32{}
  89. wg := &sync.WaitGroup{}
  90. wg.Add(len(buffers))
  91. for i := 0; i < len(buffers); i++ {
  92. go func(i int) {
  93. buffers[i] = make([]uint32, allocator.hibernatedLen)
  94. DecompressUInt32Slice(allocator.hibernatedStorage[i], buffers[i])
  95. allocator.hibernatedStorage[i] = nil
  96. wg.Done()
  97. }(i)
  98. }
  99. wg.Wait()
  100. allocator.storage = make([]node, allocator.hibernatedLen, (allocator.hibernatedLen*3)/2)
  101. for i := range allocator.storage {
  102. n := &allocator.storage[i]
  103. n.item.Key = buffers[0][i]
  104. n.item.Value = buffers[1][i]
  105. n.left = buffers[2][i]
  106. n.parent = buffers[3][i]
  107. n.right = buffers[4][i]
  108. n.color = buffers[5][i] > 0
  109. }
  110. allocator.hibernatedLen = 0
  111. }
  112. func (allocator *Allocator) malloc() uint32 {
  113. if len(allocator.gaps) > 0 {
  114. var key uint32
  115. for key = range allocator.gaps {
  116. break
  117. }
  118. delete(allocator.gaps, key)
  119. return key
  120. }
  121. n := len(allocator.storage)
  122. if n == 0 {
  123. // zero is reserved
  124. allocator.storage = append(allocator.storage, node{})
  125. n = 1
  126. }
  127. if n == negativeLimitNode-1 {
  128. // math.MaxUint32 is reserved
  129. panic("the size of my RBTree allocator has reached the maximum value for uint32, sorry")
  130. }
  131. doAssert(n < negativeLimitNode)
  132. allocator.storage = append(allocator.storage, node{})
  133. return uint32(n)
  134. }
  135. func (allocator *Allocator) free(n uint32) {
  136. _, exists := allocator.gaps[n]
  137. doAssert(!exists)
  138. allocator.storage[n] = node{}
  139. allocator.gaps[n] = true
  140. }
  141. // RBTree is a red-black tree with an API similar to C++ STL's.
  142. //
  143. // The implementation is inspired (read: stolen) from:
  144. // http://en.literateprograms.org/Red-black_tree_(C)#chunk use:private function prototypes.
  145. //
  146. // The code was optimized for the simple integer types of Key and Value.
  147. // The code was further optimized for using allocators.
  148. // Credits: Yaz Saito.
  149. type RBTree struct {
  150. // Root of the tree
  151. root uint32
  152. // The minimum and maximum nodes under the tree.
  153. minNode, maxNode uint32
  154. // Number of nodes under root, including the root
  155. count int32
  156. // Nodes allocator
  157. allocator *Allocator
  158. }
  159. // NewRBTree creates a new red-black binary tree.
  160. func NewRBTree(allocator *Allocator) *RBTree {
  161. return &RBTree{allocator: allocator}
  162. }
  163. func (tree RBTree) storage() []node {
  164. return tree.allocator.storage
  165. }
  166. // Allocator returns the bound nodes allocator.
  167. func (tree RBTree) Allocator() *Allocator {
  168. return tree.allocator
  169. }
  170. // Len returns the number of elements in the tree.
  171. func (tree RBTree) Len() int {
  172. return int(tree.count)
  173. }
  174. // CloneShallow performs a shallow copy of the tree - the nodes are assumed to already exist in the allocator.
  175. func (tree RBTree) CloneShallow(allocator *Allocator) *RBTree {
  176. clone := tree
  177. clone.allocator = allocator
  178. return &clone
  179. }
  180. // CloneDeep performs a deep copy of the tree - the nodes are created from scratch.
  181. func (tree RBTree) CloneDeep(allocator *Allocator) *RBTree {
  182. clone := &RBTree{
  183. count: tree.count,
  184. allocator: allocator,
  185. }
  186. nodeMap := map[uint32]uint32{0: 0}
  187. originStorage := tree.storage()
  188. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  189. newNode := allocator.malloc()
  190. cloneNode := &allocator.storage[newNode]
  191. cloneNode.item = *iter.Item()
  192. cloneNode.color = originStorage[iter.node].color
  193. nodeMap[iter.node] = newNode
  194. }
  195. cloneStorage := allocator.storage
  196. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  197. cloneNode := &cloneStorage[nodeMap[iter.node]]
  198. originNode := originStorage[iter.node]
  199. cloneNode.left = nodeMap[originNode.left]
  200. cloneNode.right = nodeMap[originNode.right]
  201. cloneNode.parent = nodeMap[originNode.parent]
  202. }
  203. clone.root = nodeMap[tree.root]
  204. clone.minNode = nodeMap[tree.minNode]
  205. clone.maxNode = nodeMap[tree.maxNode]
  206. return clone
  207. }
  208. // Erase removes all the nodes from the tree.
  209. func (tree *RBTree) Erase() {
  210. nodes := make([]uint32, 0, tree.count)
  211. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  212. nodes = append(nodes, iter.node)
  213. }
  214. for _, node := range nodes {
  215. tree.allocator.free(node)
  216. }
  217. tree.root = 0
  218. tree.minNode = 0
  219. tree.maxNode = 0
  220. tree.count = 0
  221. }
  222. // Get is a convenience function for finding an element equal to Key. Returns
  223. // nil if not found.
  224. func (tree RBTree) Get(key uint32) *uint32 {
  225. n, exact := tree.findGE(key)
  226. if exact {
  227. return &tree.storage()[n].item.Value
  228. }
  229. return nil
  230. }
  231. // Min creates an iterator that points to the minimum item in the tree.
  232. // If the tree is empty, returns Limit()
  233. func (tree *RBTree) Min() Iterator {
  234. return Iterator{tree, tree.minNode}
  235. }
  236. // Max creates an iterator that points at the maximum item in the tree.
  237. //
  238. // If the tree is empty, returns NegativeLimit().
  239. func (tree *RBTree) Max() Iterator {
  240. if tree.maxNode == 0 {
  241. return Iterator{tree, negativeLimitNode}
  242. }
  243. return Iterator{tree, tree.maxNode}
  244. }
  245. // Limit creates an iterator that points beyond the maximum item in the tree.
  246. func (tree *RBTree) Limit() Iterator {
  247. return Iterator{tree, 0}
  248. }
  249. // NegativeLimit creates an iterator that points before the minimum item in the tree.
  250. func (tree *RBTree) NegativeLimit() Iterator {
  251. return Iterator{tree, negativeLimitNode}
  252. }
  253. // FindGE finds the smallest element N such that N >= Key, and returns the
  254. // iterator pointing to the element. If no such element is found,
  255. // returns tree.Limit().
  256. func (tree *RBTree) FindGE(key uint32) Iterator {
  257. n, _ := tree.findGE(key)
  258. return Iterator{tree, n}
  259. }
  260. // FindLE finds the largest element N such that N <= Key, and returns the
  261. // iterator pointing to the element. If no such element is found,
  262. // returns iter.NegativeLimit().
  263. func (tree *RBTree) FindLE(key uint32) Iterator {
  264. n, exact := tree.findGE(key)
  265. if exact {
  266. return Iterator{tree, n}
  267. }
  268. if n != 0 {
  269. return Iterator{tree, doPrev(n, tree.storage())}
  270. }
  271. if tree.maxNode == 0 {
  272. return Iterator{tree, negativeLimitNode}
  273. }
  274. return Iterator{tree, tree.maxNode}
  275. }
  276. // Insert an item. If the item is already in the tree, do nothing and
  277. // return false. Else return true.
  278. func (tree *RBTree) Insert(item Item) (bool, Iterator) {
  279. // TODO: delay creating n until it is found to be inserted
  280. n := tree.doInsert(item)
  281. if n == 0 {
  282. return false, Iterator{}
  283. }
  284. alloc := tree.storage()
  285. insN := n
  286. alloc[n].color = red
  287. for true {
  288. // Case 1: N is at the root
  289. if alloc[n].parent == 0 {
  290. alloc[n].color = black
  291. break
  292. }
  293. // Case 2: The parent is black, so the tree already
  294. // satisfies the RB properties
  295. if alloc[alloc[n].parent].color == black {
  296. break
  297. }
  298. // Case 3: parent and uncle are both red.
  299. // Then paint both black and make grandparent red.
  300. grandparent := alloc[alloc[n].parent].parent
  301. var uncle uint32
  302. if isLeftChild(alloc[n].parent, alloc) {
  303. uncle = alloc[grandparent].right
  304. } else {
  305. uncle = alloc[grandparent].left
  306. }
  307. if uncle != 0 && alloc[uncle].color == red {
  308. alloc[alloc[n].parent].color = black
  309. alloc[uncle].color = black
  310. alloc[grandparent].color = red
  311. n = grandparent
  312. continue
  313. }
  314. // Case 4: parent is red, uncle is black (1)
  315. if isRightChild(n, alloc) && isLeftChild(alloc[n].parent, alloc) {
  316. tree.rotateLeft(alloc[n].parent)
  317. n = alloc[n].left
  318. continue
  319. }
  320. if isLeftChild(n, alloc) && isRightChild(alloc[n].parent, alloc) {
  321. tree.rotateRight(alloc[n].parent)
  322. n = alloc[n].right
  323. continue
  324. }
  325. // Case 5: parent is read, uncle is black (2)
  326. alloc[alloc[n].parent].color = black
  327. alloc[grandparent].color = red
  328. if isLeftChild(n, alloc) {
  329. tree.rotateRight(grandparent)
  330. } else {
  331. tree.rotateLeft(grandparent)
  332. }
  333. break
  334. }
  335. return true, Iterator{tree, insN}
  336. }
  337. // DeleteWithKey deletes an item with the given Key. Returns true iff the item was
  338. // found.
  339. func (tree *RBTree) DeleteWithKey(key uint32) bool {
  340. n, exact := tree.findGE(key)
  341. if exact {
  342. tree.doDelete(n)
  343. return true
  344. }
  345. return false
  346. }
  347. // DeleteWithIterator deletes the current item.
  348. //
  349. // REQUIRES: !iter.Limit() && !iter.NegativeLimit()
  350. func (tree *RBTree) DeleteWithIterator(iter Iterator) {
  351. doAssert(!iter.Limit() && !iter.NegativeLimit())
  352. tree.doDelete(iter.node)
  353. }
  354. // Iterator allows scanning tree elements in sort order.
  355. //
  356. // Iterator invalidation rule is the same as C++ std::map<>'s. That
  357. // is, if you delete the element that an iterator points to, the
  358. // iterator becomes invalid. For other operation types, the iterator
  359. // remains valid.
  360. type Iterator struct {
  361. tree *RBTree
  362. node uint32
  363. }
  364. // Equal checks for the underlying nodes equality.
  365. func (iter Iterator) Equal(other Iterator) bool {
  366. return iter.node == other.node
  367. }
  368. // Limit checks if the iterator points beyond the max element in the tree.
  369. func (iter Iterator) Limit() bool {
  370. return iter.node == 0
  371. }
  372. // Min checks if the iterator points to the minimum element in the tree.
  373. func (iter Iterator) Min() bool {
  374. return iter.node == iter.tree.minNode
  375. }
  376. // Max checks if the iterator points to the maximum element in the tree.
  377. func (iter Iterator) Max() bool {
  378. return iter.node == iter.tree.maxNode
  379. }
  380. // NegativeLimit checks if the iterator points before the minimum element in the tree.
  381. func (iter Iterator) NegativeLimit() bool {
  382. return iter.node == negativeLimitNode
  383. }
  384. // Item returns the current element. Allows mutating the node
  385. // (key to be changed with care!).
  386. //
  387. // The result is nil if iter.Limit() || iter.NegativeLimit().
  388. func (iter Iterator) Item() *Item {
  389. if iter.Limit() || iter.NegativeLimit() {
  390. return nil
  391. }
  392. return &iter.tree.storage()[iter.node].item
  393. }
  394. // Next creates a new iterator that points to the successor of the current element.
  395. //
  396. // REQUIRES: !iter.Limit()
  397. func (iter Iterator) Next() Iterator {
  398. doAssert(!iter.Limit())
  399. if iter.NegativeLimit() {
  400. return Iterator{iter.tree, iter.tree.minNode}
  401. }
  402. return Iterator{iter.tree, doNext(iter.node, iter.tree.storage())}
  403. }
  404. // Prev creates a new iterator that points to the predecessor of the current
  405. // node.
  406. //
  407. // REQUIRES: !iter.NegativeLimit()
  408. func (iter Iterator) Prev() Iterator {
  409. doAssert(!iter.NegativeLimit())
  410. if !iter.Limit() {
  411. return Iterator{iter.tree, doPrev(iter.node, iter.tree.storage())}
  412. }
  413. if iter.tree.maxNode == 0 {
  414. return Iterator{iter.tree, negativeLimitNode}
  415. }
  416. return Iterator{iter.tree, iter.tree.maxNode}
  417. }
  418. func doAssert(b bool) {
  419. if !b {
  420. panic("rbtree internal assertion failed")
  421. }
  422. }
  423. const (
  424. red = false
  425. black = true
  426. negativeLimitNode = math.MaxUint32
  427. )
  428. type node struct {
  429. item Item
  430. parent, left, right uint32
  431. color bool // black or red
  432. }
  433. //
  434. // Internal node attribute accessors
  435. //
  436. func getColor(n uint32, allocator []node) bool {
  437. if n == 0 {
  438. return black
  439. }
  440. return allocator[n].color
  441. }
  442. func isLeftChild(n uint32, allocator []node) bool {
  443. return n == allocator[allocator[n].parent].left
  444. }
  445. func isRightChild(n uint32, allocator []node) bool {
  446. return n == allocator[allocator[n].parent].right
  447. }
  448. func sibling(n uint32, allocator []node) uint32 {
  449. doAssert(allocator[n].parent != 0)
  450. if isLeftChild(n, allocator) {
  451. return allocator[allocator[n].parent].right
  452. }
  453. return allocator[allocator[n].parent].left
  454. }
  455. // Return the minimum node that's larger than N. Return nil if no such
  456. // node is found.
  457. func doNext(n uint32, allocator []node) uint32 {
  458. if allocator[n].right != 0 {
  459. m := allocator[n].right
  460. for allocator[m].left != 0 {
  461. m = allocator[m].left
  462. }
  463. return m
  464. }
  465. for n != 0 {
  466. p := allocator[n].parent
  467. if p == 0 {
  468. return 0
  469. }
  470. if isLeftChild(n, allocator) {
  471. return p
  472. }
  473. n = p
  474. }
  475. return 0
  476. }
  477. // Return the maximum node that's smaller than N. Return nil if no
  478. // such node is found.
  479. func doPrev(n uint32, allocator []node) uint32 {
  480. if allocator[n].left != 0 {
  481. return maxPredecessor(n, allocator)
  482. }
  483. for n != 0 {
  484. p := allocator[n].parent
  485. if p == 0 {
  486. break
  487. }
  488. if isRightChild(n, allocator) {
  489. return p
  490. }
  491. n = p
  492. }
  493. return negativeLimitNode
  494. }
  495. // Return the predecessor of "n".
  496. func maxPredecessor(n uint32, allocator []node) uint32 {
  497. doAssert(allocator[n].left != 0)
  498. m := allocator[n].left
  499. for allocator[m].right != 0 {
  500. m = allocator[m].right
  501. }
  502. return m
  503. }
  504. //
  505. // Tree methods
  506. //
  507. //
  508. // Private methods
  509. //
  510. func (tree *RBTree) recomputeMinNode() {
  511. alloc := tree.storage()
  512. tree.minNode = tree.root
  513. if tree.minNode != 0 {
  514. for alloc[tree.minNode].left != 0 {
  515. tree.minNode = alloc[tree.minNode].left
  516. }
  517. }
  518. }
  519. func (tree *RBTree) recomputeMaxNode() {
  520. alloc := tree.storage()
  521. tree.maxNode = tree.root
  522. if tree.maxNode != 0 {
  523. for alloc[tree.maxNode].right != 0 {
  524. tree.maxNode = alloc[tree.maxNode].right
  525. }
  526. }
  527. }
  528. func (tree *RBTree) maybeSetMinNode(n uint32) {
  529. alloc := tree.storage()
  530. if tree.minNode == 0 {
  531. tree.minNode = n
  532. tree.maxNode = n
  533. } else if alloc[n].item.Key < alloc[tree.minNode].item.Key {
  534. tree.minNode = n
  535. }
  536. }
  537. func (tree *RBTree) maybeSetMaxNode(n uint32) {
  538. alloc := tree.storage()
  539. if tree.maxNode == 0 {
  540. tree.minNode = n
  541. tree.maxNode = n
  542. } else if alloc[n].item.Key > alloc[tree.maxNode].item.Key {
  543. tree.maxNode = n
  544. }
  545. }
  546. // Try inserting "item" into the tree. Return nil if the item is
  547. // already in the tree. Otherwise return a new (leaf) node.
  548. func (tree *RBTree) doInsert(item Item) uint32 {
  549. if tree.root == 0 {
  550. n := tree.allocator.malloc()
  551. tree.storage()[n].item = item
  552. tree.root = n
  553. tree.minNode = n
  554. tree.maxNode = n
  555. tree.count++
  556. return n
  557. }
  558. parent := tree.root
  559. storage := tree.storage()
  560. for true {
  561. parentNode := storage[parent]
  562. comp := int(item.Key) - int(parentNode.item.Key)
  563. if comp == 0 {
  564. return 0
  565. } else if comp < 0 {
  566. if parentNode.left == 0 {
  567. n := tree.allocator.malloc()
  568. storage = tree.storage()
  569. newNode := &storage[n]
  570. newNode.item = item
  571. newNode.parent = parent
  572. storage[parent].left = n
  573. tree.count++
  574. tree.maybeSetMinNode(n)
  575. return n
  576. }
  577. parent = parentNode.left
  578. } else {
  579. if parentNode.right == 0 {
  580. n := tree.allocator.malloc()
  581. storage = tree.storage()
  582. newNode := &storage[n]
  583. newNode.item = item
  584. newNode.parent = parent
  585. storage[parent].right = n
  586. tree.count++
  587. tree.maybeSetMaxNode(n)
  588. return n
  589. }
  590. parent = parentNode.right
  591. }
  592. }
  593. panic("should not reach here")
  594. }
  595. // Find a node whose item >= Key. The 2nd return Value is true iff the
  596. // node.item==Key. Returns (nil, false) if all nodes in the tree are <
  597. // Key.
  598. func (tree RBTree) findGE(key uint32) (uint32, bool) {
  599. alloc := tree.storage()
  600. n := tree.root
  601. for true {
  602. if n == 0 {
  603. return 0, false
  604. }
  605. comp := int(key) - int(alloc[n].item.Key)
  606. if comp == 0 {
  607. return n, true
  608. } else if comp < 0 {
  609. if alloc[n].left != 0 {
  610. n = alloc[n].left
  611. } else {
  612. return n, false
  613. }
  614. } else {
  615. if alloc[n].right != 0 {
  616. n = alloc[n].right
  617. } else {
  618. succ := doNext(n, alloc)
  619. if succ == 0 {
  620. return 0, false
  621. }
  622. return succ, key == alloc[succ].item.Key
  623. }
  624. }
  625. }
  626. panic("should not reach here")
  627. }
  628. // Delete N from the tree.
  629. func (tree *RBTree) doDelete(n uint32) {
  630. alloc := tree.storage()
  631. if alloc[n].left != 0 && alloc[n].right != 0 {
  632. pred := maxPredecessor(n, alloc)
  633. tree.swapNodes(n, pred)
  634. }
  635. doAssert(alloc[n].left == 0 || alloc[n].right == 0)
  636. child := alloc[n].right
  637. if child == 0 {
  638. child = alloc[n].left
  639. }
  640. if alloc[n].color == black {
  641. alloc[n].color = getColor(child, alloc)
  642. tree.deleteCase1(n)
  643. }
  644. tree.replaceNode(n, child)
  645. if alloc[n].parent == 0 && child != 0 {
  646. alloc[child].color = black
  647. }
  648. tree.allocator.free(n)
  649. tree.count--
  650. if tree.count == 0 {
  651. tree.minNode = 0
  652. tree.maxNode = 0
  653. } else {
  654. if tree.minNode == n {
  655. tree.recomputeMinNode()
  656. }
  657. if tree.maxNode == n {
  658. tree.recomputeMaxNode()
  659. }
  660. }
  661. }
  662. // Move n to the pred's place, and vice versa
  663. //
  664. func (tree *RBTree) swapNodes(n, pred uint32) {
  665. doAssert(pred != n)
  666. alloc := tree.storage()
  667. isLeft := isLeftChild(pred, alloc)
  668. tmp := alloc[pred]
  669. tree.replaceNode(n, pred)
  670. alloc[pred].color = alloc[n].color
  671. if tmp.parent == n {
  672. // swap the positions of n and pred
  673. if isLeft {
  674. alloc[pred].left = n
  675. alloc[pred].right = alloc[n].right
  676. if alloc[pred].right != 0 {
  677. alloc[alloc[pred].right].parent = pred
  678. }
  679. } else {
  680. alloc[pred].left = alloc[n].left
  681. if alloc[pred].left != 0 {
  682. alloc[alloc[pred].left].parent = pred
  683. }
  684. alloc[pred].right = n
  685. }
  686. alloc[n].item = tmp.item
  687. alloc[n].parent = pred
  688. alloc[n].left = tmp.left
  689. if alloc[n].left != 0 {
  690. alloc[alloc[n].left].parent = n
  691. }
  692. alloc[n].right = tmp.right
  693. if alloc[n].right != 0 {
  694. alloc[alloc[n].right].parent = n
  695. }
  696. } else {
  697. alloc[pred].left = alloc[n].left
  698. if alloc[pred].left != 0 {
  699. alloc[alloc[pred].left].parent = pred
  700. }
  701. alloc[pred].right = alloc[n].right
  702. if alloc[pred].right != 0 {
  703. alloc[alloc[pred].right].parent = pred
  704. }
  705. if isLeft {
  706. alloc[tmp.parent].left = n
  707. } else {
  708. alloc[tmp.parent].right = n
  709. }
  710. alloc[n].item = tmp.item
  711. alloc[n].parent = tmp.parent
  712. alloc[n].left = tmp.left
  713. if alloc[n].left != 0 {
  714. alloc[alloc[n].left].parent = n
  715. }
  716. alloc[n].right = tmp.right
  717. if alloc[n].right != 0 {
  718. alloc[alloc[n].right].parent = n
  719. }
  720. }
  721. alloc[n].color = tmp.color
  722. }
  723. func (tree *RBTree) deleteCase1(n uint32) {
  724. alloc := tree.storage()
  725. for true {
  726. if alloc[n].parent != 0 {
  727. if getColor(sibling(n, alloc), alloc) == red {
  728. alloc[alloc[n].parent].color = red
  729. alloc[sibling(n, alloc)].color = black
  730. if n == alloc[alloc[n].parent].left {
  731. tree.rotateLeft(alloc[n].parent)
  732. } else {
  733. tree.rotateRight(alloc[n].parent)
  734. }
  735. }
  736. if getColor(alloc[n].parent, alloc) == black &&
  737. getColor(sibling(n, alloc), alloc) == black &&
  738. getColor(alloc[sibling(n, alloc)].left, alloc) == black &&
  739. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  740. alloc[sibling(n, alloc)].color = red
  741. n = alloc[n].parent
  742. continue
  743. } else {
  744. // case 4
  745. if getColor(alloc[n].parent, alloc) == red &&
  746. getColor(sibling(n, alloc), alloc) == black &&
  747. getColor(alloc[sibling(n, alloc)].left, alloc) == black &&
  748. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  749. alloc[sibling(n, alloc)].color = red
  750. alloc[alloc[n].parent].color = black
  751. } else {
  752. tree.deleteCase5(n)
  753. }
  754. }
  755. }
  756. break
  757. }
  758. }
  759. func (tree *RBTree) deleteCase5(n uint32) {
  760. alloc := tree.storage()
  761. if n == alloc[alloc[n].parent].left &&
  762. getColor(sibling(n, alloc), alloc) == black &&
  763. getColor(alloc[sibling(n, alloc)].left, alloc) == red &&
  764. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  765. alloc[sibling(n, alloc)].color = red
  766. alloc[alloc[sibling(n, alloc)].left].color = black
  767. tree.rotateRight(sibling(n, alloc))
  768. } else if n == alloc[alloc[n].parent].right &&
  769. getColor(sibling(n, alloc), alloc) == black &&
  770. getColor(alloc[sibling(n, alloc)].right, alloc) == red &&
  771. getColor(alloc[sibling(n, alloc)].left, alloc) == black {
  772. alloc[sibling(n, alloc)].color = red
  773. alloc[alloc[sibling(n, alloc)].right].color = black
  774. tree.rotateLeft(sibling(n, alloc))
  775. }
  776. // case 6
  777. alloc[sibling(n, alloc)].color = getColor(alloc[n].parent, alloc)
  778. alloc[alloc[n].parent].color = black
  779. if n == alloc[alloc[n].parent].left {
  780. doAssert(getColor(alloc[sibling(n, alloc)].right, alloc) == red)
  781. alloc[alloc[sibling(n, alloc)].right].color = black
  782. tree.rotateLeft(alloc[n].parent)
  783. } else {
  784. doAssert(getColor(alloc[sibling(n, alloc)].left, alloc) == red)
  785. alloc[alloc[sibling(n, alloc)].left].color = black
  786. tree.rotateRight(alloc[n].parent)
  787. }
  788. }
  789. func (tree *RBTree) replaceNode(oldn, newn uint32) {
  790. alloc := tree.storage()
  791. if alloc[oldn].parent == 0 {
  792. tree.root = newn
  793. } else {
  794. if oldn == alloc[alloc[oldn].parent].left {
  795. alloc[alloc[oldn].parent].left = newn
  796. } else {
  797. alloc[alloc[oldn].parent].right = newn
  798. }
  799. }
  800. if newn != 0 {
  801. alloc[newn].parent = alloc[oldn].parent
  802. }
  803. }
  804. /*
  805. X Y
  806. A Y => X C
  807. B C A B
  808. */
  809. func (tree *RBTree) rotateLeft(x uint32) {
  810. alloc := tree.storage()
  811. y := alloc[x].right
  812. alloc[x].right = alloc[y].left
  813. if alloc[y].left != 0 {
  814. alloc[alloc[y].left].parent = x
  815. }
  816. alloc[y].parent = alloc[x].parent
  817. if alloc[x].parent == 0 {
  818. tree.root = y
  819. } else {
  820. if isLeftChild(x, alloc) {
  821. alloc[alloc[x].parent].left = y
  822. } else {
  823. alloc[alloc[x].parent].right = y
  824. }
  825. }
  826. alloc[y].left = x
  827. alloc[x].parent = y
  828. }
  829. /*
  830. Y X
  831. X C => A Y
  832. A B B C
  833. */
  834. func (tree *RBTree) rotateRight(y uint32) {
  835. alloc := tree.storage()
  836. x := alloc[y].left
  837. // Move "B"
  838. alloc[y].left = alloc[x].right
  839. if alloc[x].right != 0 {
  840. alloc[alloc[x].right].parent = y
  841. }
  842. alloc[x].parent = alloc[y].parent
  843. if alloc[y].parent == 0 {
  844. tree.root = x
  845. } else {
  846. if isLeftChild(y, alloc) {
  847. alloc[alloc[y].parent].left = x
  848. } else {
  849. alloc[alloc[y].parent].right = x
  850. }
  851. }
  852. alloc[x].right = y
  853. alloc[y].parent = x
  854. }