rbtree.go 26 KB

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