rbtree.go 26 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063
  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. _, exists := allocator.gaps[n]
  260. doAssert(!exists)
  261. allocator.storage[n] = node{}
  262. allocator.gaps[n] = true
  263. }
  264. // RBTree is a red-black tree with an API similar to C++ STL's.
  265. //
  266. // The implementation is inspired (read: stolen) from:
  267. // http://en.literateprograms.org/Red-black_tree_(C)#chunk use:private function prototypes.
  268. //
  269. // The code was optimized for the simple integer types of Key and Value.
  270. // The code was further optimized for using allocators.
  271. // Credits: Yaz Saito.
  272. type RBTree struct {
  273. // Root of the tree
  274. root uint32
  275. // The minimum and maximum nodes under the tree.
  276. minNode, maxNode uint32
  277. // Number of nodes under root, including the root
  278. count int32
  279. // Nodes allocator
  280. allocator *Allocator
  281. }
  282. // NewRBTree creates a new red-black binary tree.
  283. func NewRBTree(allocator *Allocator) *RBTree {
  284. return &RBTree{allocator: allocator}
  285. }
  286. func (tree RBTree) storage() []node {
  287. return tree.allocator.storage
  288. }
  289. // Allocator returns the bound nodes allocator.
  290. func (tree RBTree) Allocator() *Allocator {
  291. return tree.allocator
  292. }
  293. // Len returns the number of elements in the tree.
  294. func (tree RBTree) Len() int {
  295. return int(tree.count)
  296. }
  297. // CloneShallow performs a shallow copy of the tree - the nodes are assumed to already exist in the allocator.
  298. func (tree RBTree) CloneShallow(allocator *Allocator) *RBTree {
  299. clone := tree
  300. clone.allocator = allocator
  301. return &clone
  302. }
  303. // CloneDeep performs a deep copy of the tree - the nodes are created from scratch.
  304. func (tree RBTree) CloneDeep(allocator *Allocator) *RBTree {
  305. clone := &RBTree{
  306. count: tree.count,
  307. allocator: allocator,
  308. }
  309. nodeMap := map[uint32]uint32{0: 0}
  310. originStorage := tree.storage()
  311. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  312. newNode := allocator.malloc()
  313. cloneNode := &allocator.storage[newNode]
  314. cloneNode.item = *iter.Item()
  315. cloneNode.color = originStorage[iter.node].color
  316. nodeMap[iter.node] = newNode
  317. }
  318. cloneStorage := allocator.storage
  319. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  320. cloneNode := &cloneStorage[nodeMap[iter.node]]
  321. originNode := originStorage[iter.node]
  322. cloneNode.left = nodeMap[originNode.left]
  323. cloneNode.right = nodeMap[originNode.right]
  324. cloneNode.parent = nodeMap[originNode.parent]
  325. }
  326. clone.root = nodeMap[tree.root]
  327. clone.minNode = nodeMap[tree.minNode]
  328. clone.maxNode = nodeMap[tree.maxNode]
  329. return clone
  330. }
  331. // Erase removes all the nodes from the tree.
  332. func (tree *RBTree) Erase() {
  333. nodes := make([]uint32, 0, tree.count)
  334. for iter := tree.Min(); !iter.Limit(); iter = iter.Next() {
  335. nodes = append(nodes, iter.node)
  336. }
  337. for _, node := range nodes {
  338. tree.allocator.free(node)
  339. }
  340. tree.root = 0
  341. tree.minNode = 0
  342. tree.maxNode = 0
  343. tree.count = 0
  344. }
  345. // Get is a convenience function for finding an element equal to Key. Returns
  346. // nil if not found.
  347. func (tree RBTree) Get(key uint32) *uint32 {
  348. n, exact := tree.findGE(key)
  349. if exact {
  350. return &tree.storage()[n].item.Value
  351. }
  352. return nil
  353. }
  354. // Min creates an iterator that points to the minimum item in the tree.
  355. // If the tree is empty, returns Limit()
  356. func (tree *RBTree) Min() Iterator {
  357. return Iterator{tree, tree.minNode}
  358. }
  359. // Max creates an iterator that points at the maximum item in the tree.
  360. //
  361. // If the tree is empty, returns NegativeLimit().
  362. func (tree *RBTree) Max() Iterator {
  363. if tree.maxNode == 0 {
  364. return Iterator{tree, negativeLimitNode}
  365. }
  366. return Iterator{tree, tree.maxNode}
  367. }
  368. // Limit creates an iterator that points beyond the maximum item in the tree.
  369. func (tree *RBTree) Limit() Iterator {
  370. return Iterator{tree, 0}
  371. }
  372. // NegativeLimit creates an iterator that points before the minimum item in the tree.
  373. func (tree *RBTree) NegativeLimit() Iterator {
  374. return Iterator{tree, negativeLimitNode}
  375. }
  376. // FindGE finds the smallest element N such that N >= Key, and returns the
  377. // iterator pointing to the element. If no such element is found,
  378. // returns tree.Limit().
  379. func (tree *RBTree) FindGE(key uint32) Iterator {
  380. n, _ := tree.findGE(key)
  381. return Iterator{tree, n}
  382. }
  383. // FindLE finds the largest element N such that N <= Key, and returns the
  384. // iterator pointing to the element. If no such element is found,
  385. // returns iter.NegativeLimit().
  386. func (tree *RBTree) FindLE(key uint32) Iterator {
  387. n, exact := tree.findGE(key)
  388. if exact {
  389. return Iterator{tree, n}
  390. }
  391. if n != 0 {
  392. return Iterator{tree, doPrev(n, tree.storage())}
  393. }
  394. if tree.maxNode == 0 {
  395. return Iterator{tree, negativeLimitNode}
  396. }
  397. return Iterator{tree, tree.maxNode}
  398. }
  399. // Insert an item. If the item is already in the tree, do nothing and
  400. // return false. Else return true.
  401. func (tree *RBTree) Insert(item Item) (bool, Iterator) {
  402. // TODO: delay creating n until it is found to be inserted
  403. n := tree.doInsert(item)
  404. if n == 0 {
  405. return false, Iterator{}
  406. }
  407. alloc := tree.storage()
  408. insN := n
  409. alloc[n].color = red
  410. for true {
  411. // Case 1: N is at the root
  412. if alloc[n].parent == 0 {
  413. alloc[n].color = black
  414. break
  415. }
  416. // Case 2: The parent is black, so the tree already
  417. // satisfies the RB properties
  418. if alloc[alloc[n].parent].color == black {
  419. break
  420. }
  421. // Case 3: parent and uncle are both red.
  422. // Then paint both black and make grandparent red.
  423. grandparent := alloc[alloc[n].parent].parent
  424. var uncle uint32
  425. if isLeftChild(alloc[n].parent, alloc) {
  426. uncle = alloc[grandparent].right
  427. } else {
  428. uncle = alloc[grandparent].left
  429. }
  430. if uncle != 0 && alloc[uncle].color == red {
  431. alloc[alloc[n].parent].color = black
  432. alloc[uncle].color = black
  433. alloc[grandparent].color = red
  434. n = grandparent
  435. continue
  436. }
  437. // Case 4: parent is red, uncle is black (1)
  438. if isRightChild(n, alloc) && isLeftChild(alloc[n].parent, alloc) {
  439. tree.rotateLeft(alloc[n].parent)
  440. n = alloc[n].left
  441. continue
  442. }
  443. if isLeftChild(n, alloc) && isRightChild(alloc[n].parent, alloc) {
  444. tree.rotateRight(alloc[n].parent)
  445. n = alloc[n].right
  446. continue
  447. }
  448. // Case 5: parent is read, uncle is black (2)
  449. alloc[alloc[n].parent].color = black
  450. alloc[grandparent].color = red
  451. if isLeftChild(n, alloc) {
  452. tree.rotateRight(grandparent)
  453. } else {
  454. tree.rotateLeft(grandparent)
  455. }
  456. break
  457. }
  458. return true, Iterator{tree, insN}
  459. }
  460. // DeleteWithKey deletes an item with the given Key. Returns true iff the item was
  461. // found.
  462. func (tree *RBTree) DeleteWithKey(key uint32) bool {
  463. n, exact := tree.findGE(key)
  464. if exact {
  465. tree.doDelete(n)
  466. return true
  467. }
  468. return false
  469. }
  470. // DeleteWithIterator deletes the current item.
  471. //
  472. // REQUIRES: !iter.Limit() && !iter.NegativeLimit()
  473. func (tree *RBTree) DeleteWithIterator(iter Iterator) {
  474. doAssert(!iter.Limit() && !iter.NegativeLimit())
  475. tree.doDelete(iter.node)
  476. }
  477. // Iterator allows scanning tree elements in sort order.
  478. //
  479. // Iterator invalidation rule is the same as C++ std::map<>'s. That
  480. // is, if you delete the element that an iterator points to, the
  481. // iterator becomes invalid. For other operation types, the iterator
  482. // remains valid.
  483. type Iterator struct {
  484. tree *RBTree
  485. node uint32
  486. }
  487. // Equal checks for the underlying nodes equality.
  488. func (iter Iterator) Equal(other Iterator) bool {
  489. return iter.node == other.node
  490. }
  491. // Limit checks if the iterator points beyond the max element in the tree.
  492. func (iter Iterator) Limit() bool {
  493. return iter.node == 0
  494. }
  495. // Min checks if the iterator points to the minimum element in the tree.
  496. func (iter Iterator) Min() bool {
  497. return iter.node == iter.tree.minNode
  498. }
  499. // Max checks if the iterator points to the maximum element in the tree.
  500. func (iter Iterator) Max() bool {
  501. return iter.node == iter.tree.maxNode
  502. }
  503. // NegativeLimit checks if the iterator points before the minimum element in the tree.
  504. func (iter Iterator) NegativeLimit() bool {
  505. return iter.node == negativeLimitNode
  506. }
  507. // Item returns the current element. Allows mutating the node
  508. // (key to be changed with care!).
  509. //
  510. // The result is nil if iter.Limit() || iter.NegativeLimit().
  511. func (iter Iterator) Item() *Item {
  512. if iter.Limit() || iter.NegativeLimit() {
  513. return nil
  514. }
  515. return &iter.tree.storage()[iter.node].item
  516. }
  517. // Next creates a new iterator that points to the successor of the current element.
  518. //
  519. // REQUIRES: !iter.Limit()
  520. func (iter Iterator) Next() Iterator {
  521. doAssert(!iter.Limit())
  522. if iter.NegativeLimit() {
  523. return Iterator{iter.tree, iter.tree.minNode}
  524. }
  525. return Iterator{iter.tree, doNext(iter.node, iter.tree.storage())}
  526. }
  527. // Prev creates a new iterator that points to the predecessor of the current
  528. // node.
  529. //
  530. // REQUIRES: !iter.NegativeLimit()
  531. func (iter Iterator) Prev() Iterator {
  532. doAssert(!iter.NegativeLimit())
  533. if !iter.Limit() {
  534. return Iterator{iter.tree, doPrev(iter.node, iter.tree.storage())}
  535. }
  536. if iter.tree.maxNode == 0 {
  537. return Iterator{iter.tree, negativeLimitNode}
  538. }
  539. return Iterator{iter.tree, iter.tree.maxNode}
  540. }
  541. func doAssert(b bool) {
  542. if !b {
  543. panic("rbtree internal assertion failed")
  544. }
  545. }
  546. const (
  547. red = false
  548. black = true
  549. negativeLimitNode = math.MaxUint32
  550. )
  551. type node struct {
  552. item Item
  553. parent, left, right uint32
  554. color bool // black or red
  555. }
  556. //
  557. // Internal node attribute accessors
  558. //
  559. func getColor(n uint32, allocator []node) bool {
  560. if n == 0 {
  561. return black
  562. }
  563. return allocator[n].color
  564. }
  565. func isLeftChild(n uint32, allocator []node) bool {
  566. return n == allocator[allocator[n].parent].left
  567. }
  568. func isRightChild(n uint32, allocator []node) bool {
  569. return n == allocator[allocator[n].parent].right
  570. }
  571. func sibling(n uint32, allocator []node) uint32 {
  572. doAssert(allocator[n].parent != 0)
  573. if isLeftChild(n, allocator) {
  574. return allocator[allocator[n].parent].right
  575. }
  576. return allocator[allocator[n].parent].left
  577. }
  578. // Return the minimum node that's larger than N. Return nil if no such
  579. // node is found.
  580. func doNext(n uint32, allocator []node) uint32 {
  581. if allocator[n].right != 0 {
  582. m := allocator[n].right
  583. for allocator[m].left != 0 {
  584. m = allocator[m].left
  585. }
  586. return m
  587. }
  588. for n != 0 {
  589. p := allocator[n].parent
  590. if p == 0 {
  591. return 0
  592. }
  593. if isLeftChild(n, allocator) {
  594. return p
  595. }
  596. n = p
  597. }
  598. return 0
  599. }
  600. // Return the maximum node that's smaller than N. Return nil if no
  601. // such node is found.
  602. func doPrev(n uint32, allocator []node) uint32 {
  603. if allocator[n].left != 0 {
  604. return maxPredecessor(n, allocator)
  605. }
  606. for n != 0 {
  607. p := allocator[n].parent
  608. if p == 0 {
  609. break
  610. }
  611. if isRightChild(n, allocator) {
  612. return p
  613. }
  614. n = p
  615. }
  616. return negativeLimitNode
  617. }
  618. // Return the predecessor of "n".
  619. func maxPredecessor(n uint32, allocator []node) uint32 {
  620. doAssert(allocator[n].left != 0)
  621. m := allocator[n].left
  622. for allocator[m].right != 0 {
  623. m = allocator[m].right
  624. }
  625. return m
  626. }
  627. //
  628. // Tree methods
  629. //
  630. //
  631. // Private methods
  632. //
  633. func (tree *RBTree) recomputeMinNode() {
  634. alloc := tree.storage()
  635. tree.minNode = tree.root
  636. if tree.minNode != 0 {
  637. for alloc[tree.minNode].left != 0 {
  638. tree.minNode = alloc[tree.minNode].left
  639. }
  640. }
  641. }
  642. func (tree *RBTree) recomputeMaxNode() {
  643. alloc := tree.storage()
  644. tree.maxNode = tree.root
  645. if tree.maxNode != 0 {
  646. for alloc[tree.maxNode].right != 0 {
  647. tree.maxNode = alloc[tree.maxNode].right
  648. }
  649. }
  650. }
  651. func (tree *RBTree) maybeSetMinNode(n uint32) {
  652. alloc := tree.storage()
  653. if tree.minNode == 0 {
  654. tree.minNode = n
  655. tree.maxNode = n
  656. } else if alloc[n].item.Key < alloc[tree.minNode].item.Key {
  657. tree.minNode = n
  658. }
  659. }
  660. func (tree *RBTree) maybeSetMaxNode(n uint32) {
  661. alloc := tree.storage()
  662. if tree.maxNode == 0 {
  663. tree.minNode = n
  664. tree.maxNode = n
  665. } else if alloc[n].item.Key > alloc[tree.maxNode].item.Key {
  666. tree.maxNode = n
  667. }
  668. }
  669. // Try inserting "item" into the tree. Return nil if the item is
  670. // already in the tree. Otherwise return a new (leaf) node.
  671. func (tree *RBTree) doInsert(item Item) uint32 {
  672. if tree.root == 0 {
  673. n := tree.allocator.malloc()
  674. tree.storage()[n].item = item
  675. tree.root = n
  676. tree.minNode = n
  677. tree.maxNode = n
  678. tree.count++
  679. return n
  680. }
  681. parent := tree.root
  682. storage := tree.storage()
  683. for true {
  684. parentNode := storage[parent]
  685. comp := int(item.Key) - int(parentNode.item.Key)
  686. if comp == 0 {
  687. return 0
  688. } else if comp < 0 {
  689. if parentNode.left == 0 {
  690. n := tree.allocator.malloc()
  691. storage = tree.storage()
  692. newNode := &storage[n]
  693. newNode.item = item
  694. newNode.parent = parent
  695. storage[parent].left = n
  696. tree.count++
  697. tree.maybeSetMinNode(n)
  698. return n
  699. }
  700. parent = parentNode.left
  701. } else {
  702. if parentNode.right == 0 {
  703. n := tree.allocator.malloc()
  704. storage = tree.storage()
  705. newNode := &storage[n]
  706. newNode.item = item
  707. newNode.parent = parent
  708. storage[parent].right = n
  709. tree.count++
  710. tree.maybeSetMaxNode(n)
  711. return n
  712. }
  713. parent = parentNode.right
  714. }
  715. }
  716. panic("should not reach here")
  717. }
  718. // Find a node whose item >= Key. The 2nd return Value is true iff the
  719. // node.item==Key. Returns (nil, false) if all nodes in the tree are <
  720. // Key.
  721. func (tree RBTree) findGE(key uint32) (uint32, bool) {
  722. alloc := tree.storage()
  723. n := tree.root
  724. for true {
  725. if n == 0 {
  726. return 0, false
  727. }
  728. comp := int(key) - int(alloc[n].item.Key)
  729. if comp == 0 {
  730. return n, true
  731. } else if comp < 0 {
  732. if alloc[n].left != 0 {
  733. n = alloc[n].left
  734. } else {
  735. return n, false
  736. }
  737. } else {
  738. if alloc[n].right != 0 {
  739. n = alloc[n].right
  740. } else {
  741. succ := doNext(n, alloc)
  742. if succ == 0 {
  743. return 0, false
  744. }
  745. return succ, key == alloc[succ].item.Key
  746. }
  747. }
  748. }
  749. panic("should not reach here")
  750. }
  751. // Delete N from the tree.
  752. func (tree *RBTree) doDelete(n uint32) {
  753. alloc := tree.storage()
  754. if alloc[n].left != 0 && alloc[n].right != 0 {
  755. pred := maxPredecessor(n, alloc)
  756. tree.swapNodes(n, pred)
  757. }
  758. doAssert(alloc[n].left == 0 || alloc[n].right == 0)
  759. child := alloc[n].right
  760. if child == 0 {
  761. child = alloc[n].left
  762. }
  763. if alloc[n].color == black {
  764. alloc[n].color = getColor(child, alloc)
  765. tree.deleteCase1(n)
  766. }
  767. tree.replaceNode(n, child)
  768. if alloc[n].parent == 0 && child != 0 {
  769. alloc[child].color = black
  770. }
  771. tree.allocator.free(n)
  772. tree.count--
  773. if tree.count == 0 {
  774. tree.minNode = 0
  775. tree.maxNode = 0
  776. } else {
  777. if tree.minNode == n {
  778. tree.recomputeMinNode()
  779. }
  780. if tree.maxNode == n {
  781. tree.recomputeMaxNode()
  782. }
  783. }
  784. }
  785. // Move n to the pred's place, and vice versa
  786. //
  787. func (tree *RBTree) swapNodes(n, pred uint32) {
  788. doAssert(pred != n)
  789. alloc := tree.storage()
  790. isLeft := isLeftChild(pred, alloc)
  791. tmp := alloc[pred]
  792. tree.replaceNode(n, pred)
  793. alloc[pred].color = alloc[n].color
  794. if tmp.parent == n {
  795. // swap the positions of n and pred
  796. if isLeft {
  797. alloc[pred].left = n
  798. alloc[pred].right = alloc[n].right
  799. if alloc[pred].right != 0 {
  800. alloc[alloc[pred].right].parent = pred
  801. }
  802. } else {
  803. alloc[pred].left = alloc[n].left
  804. if alloc[pred].left != 0 {
  805. alloc[alloc[pred].left].parent = pred
  806. }
  807. alloc[pred].right = n
  808. }
  809. alloc[n].item = tmp.item
  810. alloc[n].parent = pred
  811. alloc[n].left = tmp.left
  812. if alloc[n].left != 0 {
  813. alloc[alloc[n].left].parent = n
  814. }
  815. alloc[n].right = tmp.right
  816. if alloc[n].right != 0 {
  817. alloc[alloc[n].right].parent = n
  818. }
  819. } else {
  820. alloc[pred].left = alloc[n].left
  821. if alloc[pred].left != 0 {
  822. alloc[alloc[pred].left].parent = pred
  823. }
  824. alloc[pred].right = alloc[n].right
  825. if alloc[pred].right != 0 {
  826. alloc[alloc[pred].right].parent = pred
  827. }
  828. if isLeft {
  829. alloc[tmp.parent].left = n
  830. } else {
  831. alloc[tmp.parent].right = n
  832. }
  833. alloc[n].item = tmp.item
  834. alloc[n].parent = tmp.parent
  835. alloc[n].left = tmp.left
  836. if alloc[n].left != 0 {
  837. alloc[alloc[n].left].parent = n
  838. }
  839. alloc[n].right = tmp.right
  840. if alloc[n].right != 0 {
  841. alloc[alloc[n].right].parent = n
  842. }
  843. }
  844. alloc[n].color = tmp.color
  845. }
  846. func (tree *RBTree) deleteCase1(n uint32) {
  847. alloc := tree.storage()
  848. for true {
  849. if alloc[n].parent != 0 {
  850. if getColor(sibling(n, alloc), alloc) == red {
  851. alloc[alloc[n].parent].color = red
  852. alloc[sibling(n, alloc)].color = black
  853. if n == alloc[alloc[n].parent].left {
  854. tree.rotateLeft(alloc[n].parent)
  855. } else {
  856. tree.rotateRight(alloc[n].parent)
  857. }
  858. }
  859. if getColor(alloc[n].parent, alloc) == black &&
  860. getColor(sibling(n, alloc), alloc) == black &&
  861. getColor(alloc[sibling(n, alloc)].left, alloc) == black &&
  862. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  863. alloc[sibling(n, alloc)].color = red
  864. n = alloc[n].parent
  865. continue
  866. } else {
  867. // case 4
  868. if getColor(alloc[n].parent, alloc) == red &&
  869. getColor(sibling(n, alloc), alloc) == black &&
  870. getColor(alloc[sibling(n, alloc)].left, alloc) == black &&
  871. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  872. alloc[sibling(n, alloc)].color = red
  873. alloc[alloc[n].parent].color = black
  874. } else {
  875. tree.deleteCase5(n)
  876. }
  877. }
  878. }
  879. break
  880. }
  881. }
  882. func (tree *RBTree) deleteCase5(n uint32) {
  883. alloc := tree.storage()
  884. if n == alloc[alloc[n].parent].left &&
  885. getColor(sibling(n, alloc), alloc) == black &&
  886. getColor(alloc[sibling(n, alloc)].left, alloc) == red &&
  887. getColor(alloc[sibling(n, alloc)].right, alloc) == black {
  888. alloc[sibling(n, alloc)].color = red
  889. alloc[alloc[sibling(n, alloc)].left].color = black
  890. tree.rotateRight(sibling(n, alloc))
  891. } else if n == alloc[alloc[n].parent].right &&
  892. getColor(sibling(n, alloc), alloc) == black &&
  893. getColor(alloc[sibling(n, alloc)].right, alloc) == red &&
  894. getColor(alloc[sibling(n, alloc)].left, alloc) == black {
  895. alloc[sibling(n, alloc)].color = red
  896. alloc[alloc[sibling(n, alloc)].right].color = black
  897. tree.rotateLeft(sibling(n, alloc))
  898. }
  899. // case 6
  900. alloc[sibling(n, alloc)].color = getColor(alloc[n].parent, alloc)
  901. alloc[alloc[n].parent].color = black
  902. if n == alloc[alloc[n].parent].left {
  903. doAssert(getColor(alloc[sibling(n, alloc)].right, alloc) == red)
  904. alloc[alloc[sibling(n, alloc)].right].color = black
  905. tree.rotateLeft(alloc[n].parent)
  906. } else {
  907. doAssert(getColor(alloc[sibling(n, alloc)].left, alloc) == red)
  908. alloc[alloc[sibling(n, alloc)].left].color = black
  909. tree.rotateRight(alloc[n].parent)
  910. }
  911. }
  912. func (tree *RBTree) replaceNode(oldn, newn uint32) {
  913. alloc := tree.storage()
  914. if alloc[oldn].parent == 0 {
  915. tree.root = newn
  916. } else {
  917. if oldn == alloc[alloc[oldn].parent].left {
  918. alloc[alloc[oldn].parent].left = newn
  919. } else {
  920. alloc[alloc[oldn].parent].right = newn
  921. }
  922. }
  923. if newn != 0 {
  924. alloc[newn].parent = alloc[oldn].parent
  925. }
  926. }
  927. /*
  928. X Y
  929. A Y => X C
  930. B C A B
  931. */
  932. func (tree *RBTree) rotateLeft(x uint32) {
  933. alloc := tree.storage()
  934. y := alloc[x].right
  935. alloc[x].right = alloc[y].left
  936. if alloc[y].left != 0 {
  937. alloc[alloc[y].left].parent = x
  938. }
  939. alloc[y].parent = alloc[x].parent
  940. if alloc[x].parent == 0 {
  941. tree.root = y
  942. } else {
  943. if isLeftChild(x, alloc) {
  944. alloc[alloc[x].parent].left = y
  945. } else {
  946. alloc[alloc[x].parent].right = y
  947. }
  948. }
  949. alloc[y].left = x
  950. alloc[x].parent = y
  951. }
  952. /*
  953. Y X
  954. X C => A Y
  955. A B B C
  956. */
  957. func (tree *RBTree) rotateRight(y uint32) {
  958. alloc := tree.storage()
  959. x := alloc[y].left
  960. // Move "B"
  961. alloc[y].left = alloc[x].right
  962. if alloc[x].right != 0 {
  963. alloc[alloc[x].right].parent = y
  964. }
  965. alloc[x].parent = alloc[y].parent
  966. if alloc[y].parent == 0 {
  967. tree.root = x
  968. } else {
  969. if isLeftChild(y, alloc) {
  970. alloc[alloc[y].parent].left = x
  971. } else {
  972. alloc[alloc[y].parent].right = x
  973. }
  974. }
  975. alloc[x].right = y
  976. alloc[y].parent = x
  977. }