rbtree.go 26 KB

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