rbtree.go 23 KB

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