go.etcd.io/etcd@v3.3.27+incompatible/pkg/adt/interval_tree.go (about)

     1  // Copyright 2016 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package adt
    16  
    17  import (
    18  	"bytes"
    19  	"fmt"
    20  	"math"
    21  	"strings"
    22  )
    23  
    24  // Comparable is an interface for trichotomic comparisons.
    25  type Comparable interface {
    26  	// Compare gives the result of a 3-way comparison
    27  	// a.Compare(b) = 1 => a > b
    28  	// a.Compare(b) = 0 => a == b
    29  	// a.Compare(b) = -1 => a < b
    30  	Compare(c Comparable) int
    31  }
    32  
    33  type rbcolor int
    34  
    35  const (
    36  	black rbcolor = iota
    37  	red
    38  )
    39  
    40  func (c rbcolor) String() string {
    41  	switch c {
    42  	case black:
    43  		return "black"
    44  	case red:
    45  		return "black"
    46  	default:
    47  		panic(fmt.Errorf("unknown color %d", c))
    48  	}
    49  }
    50  
    51  // Interval implements a Comparable interval [begin, end)
    52  // TODO: support different sorts of intervals: (a,b), [a,b], (a, b]
    53  type Interval struct {
    54  	Begin Comparable
    55  	End   Comparable
    56  }
    57  
    58  // Compare on an interval gives == if the interval overlaps.
    59  func (ivl *Interval) Compare(c Comparable) int {
    60  	ivl2 := c.(*Interval)
    61  	ivbCmpBegin := ivl.Begin.Compare(ivl2.Begin)
    62  	ivbCmpEnd := ivl.Begin.Compare(ivl2.End)
    63  	iveCmpBegin := ivl.End.Compare(ivl2.Begin)
    64  
    65  	// ivl is left of ivl2
    66  	if ivbCmpBegin < 0 && iveCmpBegin <= 0 {
    67  		return -1
    68  	}
    69  
    70  	// iv is right of iv2
    71  	if ivbCmpEnd >= 0 {
    72  		return 1
    73  	}
    74  
    75  	return 0
    76  }
    77  
    78  type intervalNode struct {
    79  	// iv is the interval-value pair entry.
    80  	iv IntervalValue
    81  	// max endpoint of all descendent nodes.
    82  	max Comparable
    83  	// left and right are sorted by low endpoint of key interval
    84  	left, right *intervalNode
    85  	// parent is the direct ancestor of the node
    86  	parent *intervalNode
    87  	c      rbcolor
    88  }
    89  
    90  func (x *intervalNode) color(sentinel *intervalNode) rbcolor {
    91  	if x == sentinel {
    92  		return black
    93  	}
    94  	return x.c
    95  }
    96  
    97  func (x *intervalNode) height(sentinel *intervalNode) int {
    98  	if x == sentinel {
    99  		return 0
   100  	}
   101  	ld := x.left.height(sentinel)
   102  	rd := x.right.height(sentinel)
   103  	if ld < rd {
   104  		return rd + 1
   105  	}
   106  	return ld + 1
   107  }
   108  
   109  func (x *intervalNode) min(sentinel *intervalNode) *intervalNode {
   110  	for x.left != sentinel {
   111  		x = x.left
   112  	}
   113  	return x
   114  }
   115  
   116  // successor is the next in-order node in the tree
   117  func (x *intervalNode) successor(sentinel *intervalNode) *intervalNode {
   118  	if x.right != sentinel {
   119  		return x.right.min(sentinel)
   120  	}
   121  	y := x.parent
   122  	for y != sentinel && x == y.right {
   123  		x = y
   124  		y = y.parent
   125  	}
   126  	return y
   127  }
   128  
   129  // updateMax updates the maximum values for a node and its ancestors
   130  func (x *intervalNode) updateMax(sentinel *intervalNode) {
   131  	for x != sentinel {
   132  		oldmax := x.max
   133  		max := x.iv.Ivl.End
   134  		if x.left != sentinel && x.left.max.Compare(max) > 0 {
   135  			max = x.left.max
   136  		}
   137  		if x.right != sentinel && x.right.max.Compare(max) > 0 {
   138  			max = x.right.max
   139  		}
   140  		if oldmax.Compare(max) == 0 {
   141  			break
   142  		}
   143  		x.max = max
   144  		x = x.parent
   145  	}
   146  }
   147  
   148  type nodeVisitor func(n *intervalNode) bool
   149  
   150  // visit will call a node visitor on each node that overlaps the given interval
   151  func (x *intervalNode) visit(iv *Interval, sentinel *intervalNode, nv nodeVisitor) bool {
   152  	if x == sentinel {
   153  		return true
   154  	}
   155  	v := iv.Compare(&x.iv.Ivl)
   156  	switch {
   157  	case v < 0:
   158  		if !x.left.visit(iv, sentinel, nv) {
   159  			return false
   160  		}
   161  	case v > 0:
   162  		maxiv := Interval{x.iv.Ivl.Begin, x.max}
   163  		if maxiv.Compare(iv) == 0 {
   164  			if !x.left.visit(iv, sentinel, nv) || !x.right.visit(iv, sentinel, nv) {
   165  				return false
   166  			}
   167  		}
   168  	default:
   169  		if !x.left.visit(iv, sentinel, nv) || !nv(x) || !x.right.visit(iv, sentinel, nv) {
   170  			return false
   171  		}
   172  	}
   173  	return true
   174  }
   175  
   176  // IntervalValue represents a range tree node that contains a range and a value.
   177  type IntervalValue struct {
   178  	Ivl Interval
   179  	Val interface{}
   180  }
   181  
   182  // IntervalTree represents a (mostly) textbook implementation of the
   183  // "Introduction to Algorithms" (Cormen et al, 3rd ed.) chapter 13 red-black tree
   184  // and chapter 14.3 interval tree with search supporting "stabbing queries".
   185  type IntervalTree interface {
   186  	// Insert adds a node with the given interval into the tree.
   187  	Insert(ivl Interval, val interface{})
   188  	// Delete removes the node with the given interval from the tree, returning
   189  	// true if a node is in fact removed.
   190  	Delete(ivl Interval) bool
   191  	// Len gives the number of elements in the tree.
   192  	Len() int
   193  	// Height is the number of levels in the tree; one node has height 1.
   194  	Height() int
   195  	// MaxHeight is the expected maximum tree height given the number of nodes.
   196  	MaxHeight() int
   197  	// Visit calls a visitor function on every tree node intersecting the given interval.
   198  	// It will visit each interval [x, y) in ascending order sorted on x.
   199  	Visit(ivl Interval, ivv IntervalVisitor)
   200  	// Find gets the IntervalValue for the node matching the given interval
   201  	Find(ivl Interval) *IntervalValue
   202  	// Intersects returns true if there is some tree node intersecting the given interval.
   203  	Intersects(iv Interval) bool
   204  	// Contains returns true if the interval tree's keys cover the entire given interval.
   205  	Contains(ivl Interval) bool
   206  	// Stab returns a slice with all elements in the tree intersecting the interval.
   207  	Stab(iv Interval) []*IntervalValue
   208  	// Union merges a given interval tree into the receiver.
   209  	Union(inIvt IntervalTree, ivl Interval)
   210  }
   211  
   212  // NewIntervalTree returns a new interval tree.
   213  func NewIntervalTree() IntervalTree {
   214  	sentinel := &intervalNode{
   215  		iv:     IntervalValue{},
   216  		max:    nil,
   217  		left:   nil,
   218  		right:  nil,
   219  		parent: nil,
   220  		c:      black,
   221  	}
   222  	return &intervalTree{
   223  		root:     sentinel,
   224  		count:    0,
   225  		sentinel: sentinel,
   226  	}
   227  }
   228  
   229  type intervalTree struct {
   230  	root  *intervalNode
   231  	count int
   232  
   233  	// red-black NIL node
   234  	// use 'sentinel' as a dummy object to simplify boundary conditions
   235  	// use the sentinel to treat a nil child of a node x as an ordinary node whose parent is x
   236  	// use one shared sentinel to represent all nil leaves and the root's parent
   237  	sentinel *intervalNode
   238  }
   239  
   240  // TODO: make this consistent with textbook implementation
   241  //
   242  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.4, p324
   243  //
   244  //	 0. RB-DELETE(T, z)
   245  //	 1.
   246  //	 2. y = z
   247  //	 3. y-original-color = y.color
   248  //	 4.
   249  //	 5. if z.left == T.nil
   250  //	 6. 	x = z.right
   251  //	 7. 	RB-TRANSPLANT(T, z, z.right)
   252  //	 8. else if z.right == T.nil
   253  //	 9. 	x = z.left
   254  //	10. 	RB-TRANSPLANT(T, z, z.left)
   255  //	11. else
   256  //	12. 	y = TREE-MINIMUM(z.right)
   257  //	13. 	y-original-color = y.color
   258  //	14. 	x = y.right
   259  //	15. 	if y.p == z
   260  //	16. 		x.p = y
   261  //	17. 	else
   262  //	18. 		RB-TRANSPLANT(T, y, y.right)
   263  //	19. 		y.right = z.right
   264  //	20. 		y.right.p = y
   265  //	21. 	RB-TRANSPLANT(T, z, y)
   266  //	22. 	y.left = z.left
   267  //	23. 	y.left.p = y
   268  //	24. 	y.color = z.color
   269  //	25.
   270  //	26. if y-original-color == BLACK
   271  //	27. 	RB-DELETE-FIXUP(T, x)
   272  
   273  // Delete removes the node with the given interval from the tree, returning
   274  // true if a node is in fact removed.
   275  func (ivt *intervalTree) Delete(ivl Interval) bool {
   276  	z := ivt.find(ivl)
   277  	if z == ivt.sentinel {
   278  		return false
   279  	}
   280  
   281  	y := z
   282  	if z.left != ivt.sentinel && z.right != ivt.sentinel {
   283  		y = z.successor(ivt.sentinel)
   284  	}
   285  
   286  	x := ivt.sentinel
   287  	if y.left != ivt.sentinel {
   288  		x = y.left
   289  	} else if y.right != ivt.sentinel {
   290  		x = y.right
   291  	}
   292  
   293  	x.parent = y.parent
   294  
   295  	if y.parent == ivt.sentinel {
   296  		ivt.root = x
   297  	} else {
   298  		if y == y.parent.left {
   299  			y.parent.left = x
   300  		} else {
   301  			y.parent.right = x
   302  		}
   303  		y.parent.updateMax(ivt.sentinel)
   304  	}
   305  	if y != z {
   306  		z.iv = y.iv
   307  		z.updateMax(ivt.sentinel)
   308  	}
   309  
   310  	if y.color(ivt.sentinel) == black {
   311  		ivt.deleteFixup(x)
   312  	}
   313  
   314  	ivt.count--
   315  	return true
   316  }
   317  
   318  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.4, p326
   319  //
   320  //	 0. RB-DELETE-FIXUP(T, z)
   321  //	 1.
   322  //	 2. while x ≠ T.root and x.color == BLACK
   323  //	 3. 	if x == x.p.left
   324  //	 4. 		w = x.p.right
   325  //	 5. 		if w.color == RED
   326  //	 6. 			w.color = BLACK
   327  //	 7. 			x.p.color = RED
   328  //	 8. 			LEFT-ROTATE(T, x, p)
   329  //	 9. 		if w.left.color == BLACK and w.right.color == BLACK
   330  //	10. 			w.color = RED
   331  //	11. 			x = x.p
   332  //	12. 		else if w.right.color == BLACK
   333  //	13. 				w.left.color = BLACK
   334  //	14. 				w.color = RED
   335  //	15. 				RIGHT-ROTATE(T, w)
   336  //	16. 				w = w.p.right
   337  //	17. 			w.color = x.p.color
   338  //	18. 			x.p.color = BLACK
   339  //	19. 			LEFT-ROTATE(T, w.p)
   340  //	20. 			x = T.root
   341  //	21. 	else
   342  //	22. 		w = x.p.left
   343  //	23. 		if w.color == RED
   344  //	24. 			w.color = BLACK
   345  //	25. 			x.p.color = RED
   346  //	26. 			RIGHT-ROTATE(T, x, p)
   347  //	27. 		if w.right.color == BLACK and w.left.color == BLACK
   348  //	28. 			w.color = RED
   349  //	29. 			x = x.p
   350  //	30. 		else if w.left.color == BLACK
   351  //	31. 				w.right.color = BLACK
   352  //	32. 				w.color = RED
   353  //	33. 				LEFT-ROTATE(T, w)
   354  //	34. 				w = w.p.left
   355  //	35. 			w.color = x.p.color
   356  //	36. 			x.p.color = BLACK
   357  //	37. 			RIGHT-ROTATE(T, w.p)
   358  //	38. 			x = T.root
   359  //	39.
   360  //	40. x.color = BLACK
   361  //
   362  func (ivt *intervalTree) deleteFixup(x *intervalNode) {
   363  	for x != ivt.root && x.color(ivt.sentinel) == black {
   364  		if x == x.parent.left { // line 3-20
   365  			w := x.parent.right
   366  			if w.color(ivt.sentinel) == red {
   367  				w.c = black
   368  				x.parent.c = red
   369  				ivt.rotateLeft(x.parent)
   370  				w = x.parent.right
   371  			}
   372  			if w == nil {
   373  				break
   374  			}
   375  			if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
   376  				w.c = red
   377  				x = x.parent
   378  			} else {
   379  				if w.right.color(ivt.sentinel) == black {
   380  					w.left.c = black
   381  					w.c = red
   382  					ivt.rotateRight(w)
   383  					w = x.parent.right
   384  				}
   385  				w.c = x.parent.color(ivt.sentinel)
   386  				x.parent.c = black
   387  				w.right.c = black
   388  				ivt.rotateLeft(x.parent)
   389  				x = ivt.root
   390  			}
   391  		} else { // line 22-38
   392  			// same as above but with left and right exchanged
   393  			w := x.parent.left
   394  			if w.color(ivt.sentinel) == red {
   395  				w.c = black
   396  				x.parent.c = red
   397  				ivt.rotateRight(x.parent)
   398  				w = x.parent.left
   399  			}
   400  			if w == nil {
   401  				break
   402  			}
   403  			if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
   404  				w.c = red
   405  				x = x.parent
   406  			} else {
   407  				if w.left.color(ivt.sentinel) == black {
   408  					w.right.c = black
   409  					w.c = red
   410  					ivt.rotateLeft(w)
   411  					w = x.parent.left
   412  				}
   413  				w.c = x.parent.color(ivt.sentinel)
   414  				x.parent.c = black
   415  				w.left.c = black
   416  				ivt.rotateRight(x.parent)
   417  				x = ivt.root
   418  			}
   419  		}
   420  	}
   421  
   422  	if x != nil {
   423  		x.c = black
   424  	}
   425  }
   426  
   427  func (ivt *intervalTree) createIntervalNode(ivl Interval, val interface{}) *intervalNode {
   428  	return &intervalNode{
   429  		iv:     IntervalValue{ivl, val},
   430  		max:    ivl.End,
   431  		c:      red,
   432  		left:   ivt.sentinel,
   433  		right:  ivt.sentinel,
   434  		parent: ivt.sentinel,
   435  	}
   436  }
   437  
   438  // TODO: make this consistent with textbook implementation
   439  //
   440  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p315
   441  //
   442  //	 0. RB-INSERT(T, z)
   443  //	 1.
   444  //	 2. y = T.nil
   445  //	 3. x = T.root
   446  //	 4.
   447  //	 5. while x ≠ T.nil
   448  //	 6. 	y = x
   449  //	 7. 	if z.key < x.key
   450  //	 8. 		x = x.left
   451  //	 9. 	else
   452  //	10. 		x = x.right
   453  //	11.
   454  //	12. z.p = y
   455  //	13.
   456  //	14. if y == T.nil
   457  //	15. 	T.root = z
   458  //	16. else if z.key < y.key
   459  //	17. 	y.left = z
   460  //	18. else
   461  //	19. 	y.right = z
   462  //	20.
   463  //	21. z.left = T.nil
   464  //	22. z.right = T.nil
   465  //	23. z.color = RED
   466  //	24.
   467  //	25. RB-INSERT-FIXUP(T, z)
   468  
   469  // Insert adds a node with the given interval into the tree.
   470  func (ivt *intervalTree) Insert(ivl Interval, val interface{}) {
   471  	y := ivt.sentinel
   472  	z := ivt.createIntervalNode(ivl, val)
   473  	x := ivt.root
   474  	for x != ivt.sentinel {
   475  		y = x
   476  		if z.iv.Ivl.Begin.Compare(x.iv.Ivl.Begin) < 0 {
   477  			x = x.left
   478  		} else {
   479  			x = x.right
   480  		}
   481  	}
   482  
   483  	z.parent = y
   484  	if y == ivt.sentinel {
   485  		ivt.root = z
   486  	} else {
   487  		if z.iv.Ivl.Begin.Compare(y.iv.Ivl.Begin) < 0 {
   488  			y.left = z
   489  		} else {
   490  			y.right = z
   491  		}
   492  		y.updateMax(ivt.sentinel)
   493  	}
   494  	z.c = red
   495  
   496  	ivt.insertFixup(z)
   497  	ivt.count++
   498  }
   499  
   500  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p316
   501  //
   502  //	 0. RB-INSERT-FIXUP(T, z)
   503  //	 1.
   504  //	 2. while z.p.color == RED
   505  //	 3. 	if z.p == z.p.p.left
   506  //	 4. 		y = z.p.p.right
   507  //	 5. 		if y.color == RED
   508  //	 6. 			z.p.color = BLACK
   509  //	 7. 			y.color = BLACK
   510  //	 8. 			z.p.p.color = RED
   511  //	 9. 			z = z.p.p
   512  //	10. 		else if z == z.p.right
   513  //	11. 				z = z.p
   514  //	12. 				LEFT-ROTATE(T, z)
   515  //	13. 			z.p.color = BLACK
   516  //	14. 			z.p.p.color = RED
   517  //	15. 			RIGHT-ROTATE(T, z.p.p)
   518  //	16. 	else
   519  //	17. 		y = z.p.p.left
   520  //	18. 		if y.color == RED
   521  //	19. 			z.p.color = BLACK
   522  //	20. 			y.color = BLACK
   523  //	21. 			z.p.p.color = RED
   524  //	22. 			z = z.p.p
   525  //	23. 		else if z == z.p.right
   526  //	24. 				z = z.p
   527  //	25. 				RIGHT-ROTATE(T, z)
   528  //	26. 			z.p.color = BLACK
   529  //	27. 			z.p.p.color = RED
   530  //	28. 			LEFT-ROTATE(T, z.p.p)
   531  //	29.
   532  //	30. T.root.color = BLACK
   533  //
   534  func (ivt *intervalTree) insertFixup(z *intervalNode) {
   535  	for z.parent.color(ivt.sentinel) == red {
   536  		if z.parent == z.parent.parent.left { // line 3-15
   537  
   538  			y := z.parent.parent.right
   539  			if y.color(ivt.sentinel) == red {
   540  				y.c = black
   541  				z.parent.c = black
   542  				z.parent.parent.c = red
   543  				z = z.parent.parent
   544  			} else {
   545  				if z == z.parent.right {
   546  					z = z.parent
   547  					ivt.rotateLeft(z)
   548  				}
   549  				z.parent.c = black
   550  				z.parent.parent.c = red
   551  				ivt.rotateRight(z.parent.parent)
   552  			}
   553  		} else { // line 16-28
   554  			// same as then with left/right exchanged
   555  			y := z.parent.parent.left
   556  			if y.color(ivt.sentinel) == red {
   557  				y.c = black
   558  				z.parent.c = black
   559  				z.parent.parent.c = red
   560  				z = z.parent.parent
   561  			} else {
   562  				if z == z.parent.left {
   563  					z = z.parent
   564  					ivt.rotateRight(z)
   565  				}
   566  				z.parent.c = black
   567  				z.parent.parent.c = red
   568  				ivt.rotateLeft(z.parent.parent)
   569  			}
   570  		}
   571  	}
   572  
   573  	// line 30
   574  	ivt.root.c = black
   575  }
   576  
   577  // rotateLeft moves x so it is left of its right child
   578  //
   579  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.2, p313
   580  //
   581  //	 0. LEFT-ROTATE(T, x)
   582  //	 1.
   583  //	 2. y = x.right
   584  //	 3. x.right = y.left
   585  //	 4.
   586  //	 5. if y.left ≠ T.nil
   587  //	 6. 	y.left.p = x
   588  //	 7.
   589  //	 8. y.p = x.p
   590  //	 9.
   591  //	10. if x.p == T.nil
   592  //	11. 	T.root = y
   593  //	12. else if x == x.p.left
   594  //	13. 	x.p.left = y
   595  //	14. else
   596  //	15. 	x.p.right = y
   597  //	16.
   598  //	17. y.left = x
   599  //	18. x.p = y
   600  //
   601  func (ivt *intervalTree) rotateLeft(x *intervalNode) {
   602  	// rotateLeft x must have right child
   603  	if x.right == ivt.sentinel {
   604  		return
   605  	}
   606  
   607  	// line 2-3
   608  	y := x.right
   609  	x.right = y.left
   610  
   611  	// line 5-6
   612  	if y.left != ivt.sentinel {
   613  		y.left.parent = x
   614  	}
   615  	x.updateMax(ivt.sentinel)
   616  
   617  	// line 10-15, 18
   618  	ivt.replaceParent(x, y)
   619  
   620  	// line 17
   621  	y.left = x
   622  	y.updateMax(ivt.sentinel)
   623  }
   624  
   625  // rotateRight moves x so it is right of its left child
   626  //
   627  //	 0. RIGHT-ROTATE(T, x)
   628  //	 1.
   629  //	 2. y = x.left
   630  //	 3. x.left = y.right
   631  //	 4.
   632  //	 5. if y.right ≠ T.nil
   633  //	 6. 	y.right.p = x
   634  //	 7.
   635  //	 8. y.p = x.p
   636  //	 9.
   637  //	10. if x.p == T.nil
   638  //	11. 	T.root = y
   639  //	12. else if x == x.p.right
   640  //	13. 	x.p.right = y
   641  //	14. else
   642  //	15. 	x.p.left = y
   643  //	16.
   644  //	17. y.right = x
   645  //	18. x.p = y
   646  //
   647  func (ivt *intervalTree) rotateRight(x *intervalNode) {
   648  	// rotateRight x must have left child
   649  	if x.left == ivt.sentinel {
   650  		return
   651  	}
   652  
   653  	// line 2-3
   654  	y := x.left
   655  	x.left = y.right
   656  
   657  	// line 5-6
   658  	if y.right != ivt.sentinel {
   659  		y.right.parent = x
   660  	}
   661  	x.updateMax(ivt.sentinel)
   662  
   663  	// line 10-15, 18
   664  	ivt.replaceParent(x, y)
   665  
   666  	// line 17
   667  	y.right = x
   668  	y.updateMax(ivt.sentinel)
   669  }
   670  
   671  // replaceParent replaces x's parent with y
   672  func (ivt *intervalTree) replaceParent(x *intervalNode, y *intervalNode) {
   673  	y.parent = x.parent
   674  	if x.parent == ivt.sentinel {
   675  		ivt.root = y
   676  	} else {
   677  		if x == x.parent.left {
   678  			x.parent.left = y
   679  		} else {
   680  			x.parent.right = y
   681  		}
   682  		x.parent.updateMax(ivt.sentinel)
   683  	}
   684  	x.parent = y
   685  }
   686  
   687  // Len gives the number of elements in the tree
   688  func (ivt *intervalTree) Len() int { return ivt.count }
   689  
   690  // Height is the number of levels in the tree; one node has height 1.
   691  func (ivt *intervalTree) Height() int { return ivt.root.height(ivt.sentinel) }
   692  
   693  // MaxHeight is the expected maximum tree height given the number of nodes
   694  func (ivt *intervalTree) MaxHeight() int {
   695  	return int((2 * math.Log2(float64(ivt.Len()+1))) + 0.5)
   696  }
   697  
   698  // IntervalVisitor is used on tree searches; return false to stop searching.
   699  type IntervalVisitor func(n *IntervalValue) bool
   700  
   701  // Visit calls a visitor function on every tree node intersecting the given interval.
   702  // It will visit each interval [x, y) in ascending order sorted on x.
   703  func (ivt *intervalTree) Visit(ivl Interval, ivv IntervalVisitor) {
   704  	ivt.root.visit(&ivl, ivt.sentinel, func(n *intervalNode) bool { return ivv(&n.iv) })
   705  }
   706  
   707  // find the exact node for a given interval
   708  func (ivt *intervalTree) find(ivl Interval) *intervalNode {
   709  	ret := ivt.sentinel
   710  	f := func(n *intervalNode) bool {
   711  		if n.iv.Ivl != ivl {
   712  			return true
   713  		}
   714  		ret = n
   715  		return false
   716  	}
   717  	ivt.root.visit(&ivl, ivt.sentinel, f)
   718  	return ret
   719  }
   720  
   721  // Find gets the IntervalValue for the node matching the given interval
   722  func (ivt *intervalTree) Find(ivl Interval) (ret *IntervalValue) {
   723  	n := ivt.find(ivl)
   724  	if n == ivt.sentinel {
   725  		return nil
   726  	}
   727  	return &n.iv
   728  }
   729  
   730  // Intersects returns true if there is some tree node intersecting the given interval.
   731  func (ivt *intervalTree) Intersects(iv Interval) bool {
   732  	x := ivt.root
   733  	for x != ivt.sentinel && iv.Compare(&x.iv.Ivl) != 0 {
   734  		if x.left != ivt.sentinel && x.left.max.Compare(iv.Begin) > 0 {
   735  			x = x.left
   736  		} else {
   737  			x = x.right
   738  		}
   739  	}
   740  	return x != ivt.sentinel
   741  }
   742  
   743  // Contains returns true if the interval tree's keys cover the entire given interval.
   744  func (ivt *intervalTree) Contains(ivl Interval) bool {
   745  	var maxEnd, minBegin Comparable
   746  
   747  	isContiguous := true
   748  	ivt.Visit(ivl, func(n *IntervalValue) bool {
   749  		if minBegin == nil {
   750  			minBegin = n.Ivl.Begin
   751  			maxEnd = n.Ivl.End
   752  			return true
   753  		}
   754  		if maxEnd.Compare(n.Ivl.Begin) < 0 {
   755  			isContiguous = false
   756  			return false
   757  		}
   758  		if n.Ivl.End.Compare(maxEnd) > 0 {
   759  			maxEnd = n.Ivl.End
   760  		}
   761  		return true
   762  	})
   763  
   764  	return isContiguous && minBegin != nil && maxEnd.Compare(ivl.End) >= 0 && minBegin.Compare(ivl.Begin) <= 0
   765  }
   766  
   767  // Stab returns a slice with all elements in the tree intersecting the interval.
   768  func (ivt *intervalTree) Stab(iv Interval) (ivs []*IntervalValue) {
   769  	if ivt.count == 0 {
   770  		return nil
   771  	}
   772  	f := func(n *IntervalValue) bool { ivs = append(ivs, n); return true }
   773  	ivt.Visit(iv, f)
   774  	return ivs
   775  }
   776  
   777  // Union merges a given interval tree into the receiver.
   778  func (ivt *intervalTree) Union(inIvt IntervalTree, ivl Interval) {
   779  	f := func(n *IntervalValue) bool {
   780  		ivt.Insert(n.Ivl, n.Val)
   781  		return true
   782  	}
   783  	inIvt.Visit(ivl, f)
   784  }
   785  
   786  type visitedInterval struct {
   787  	root  Interval
   788  	left  Interval
   789  	right Interval
   790  	color rbcolor
   791  	depth int
   792  }
   793  
   794  func (vi visitedInterval) String() string {
   795  	bd := new(strings.Builder)
   796  	bd.WriteString(fmt.Sprintf("root [%v,%v,%v], left [%v,%v], right [%v,%v], depth %d",
   797  		vi.root.Begin, vi.root.End, vi.color,
   798  		vi.left.Begin, vi.left.End,
   799  		vi.right.Begin, vi.right.End,
   800  		vi.depth,
   801  	))
   802  	return bd.String()
   803  }
   804  
   805  // visitLevel traverses tree in level order.
   806  // used for testing
   807  func (ivt *intervalTree) visitLevel() []visitedInterval {
   808  	if ivt.root == ivt.sentinel {
   809  		return nil
   810  	}
   811  
   812  	rs := make([]visitedInterval, 0, ivt.Len())
   813  
   814  	type pair struct {
   815  		node  *intervalNode
   816  		depth int
   817  	}
   818  	queue := []pair{{ivt.root, 0}}
   819  	for len(queue) > 0 {
   820  		f := queue[0]
   821  		queue = queue[1:]
   822  
   823  		vi := visitedInterval{
   824  			root:  f.node.iv.Ivl,
   825  			color: f.node.color(ivt.sentinel),
   826  			depth: f.depth,
   827  		}
   828  		if f.node.left != ivt.sentinel {
   829  			vi.left = f.node.left.iv.Ivl
   830  			queue = append(queue, pair{f.node.left, f.depth + 1})
   831  		}
   832  		if f.node.right != ivt.sentinel {
   833  			vi.right = f.node.right.iv.Ivl
   834  			queue = append(queue, pair{f.node.right, f.depth + 1})
   835  		}
   836  
   837  		rs = append(rs, vi)
   838  	}
   839  
   840  	return rs
   841  }
   842  
   843  type StringComparable string
   844  
   845  func (s StringComparable) Compare(c Comparable) int {
   846  	sc := c.(StringComparable)
   847  	if s < sc {
   848  		return -1
   849  	}
   850  	if s > sc {
   851  		return 1
   852  	}
   853  	return 0
   854  }
   855  
   856  func NewStringInterval(begin, end string) Interval {
   857  	return Interval{StringComparable(begin), StringComparable(end)}
   858  }
   859  
   860  func NewStringPoint(s string) Interval {
   861  	return Interval{StringComparable(s), StringComparable(s + "\x00")}
   862  }
   863  
   864  // StringAffineComparable treats "" as > all other strings
   865  type StringAffineComparable string
   866  
   867  func (s StringAffineComparable) Compare(c Comparable) int {
   868  	sc := c.(StringAffineComparable)
   869  
   870  	if len(s) == 0 {
   871  		if len(sc) == 0 {
   872  			return 0
   873  		}
   874  		return 1
   875  	}
   876  	if len(sc) == 0 {
   877  		return -1
   878  	}
   879  
   880  	if s < sc {
   881  		return -1
   882  	}
   883  	if s > sc {
   884  		return 1
   885  	}
   886  	return 0
   887  }
   888  
   889  func NewStringAffineInterval(begin, end string) Interval {
   890  	return Interval{StringAffineComparable(begin), StringAffineComparable(end)}
   891  }
   892  
   893  func NewStringAffinePoint(s string) Interval {
   894  	return NewStringAffineInterval(s, s+"\x00")
   895  }
   896  
   897  func NewInt64Interval(a int64, b int64) Interval {
   898  	return Interval{Int64Comparable(a), Int64Comparable(b)}
   899  }
   900  
   901  func newInt64EmptyInterval() Interval {
   902  	return Interval{Begin: nil, End: nil}
   903  }
   904  
   905  func NewInt64Point(a int64) Interval {
   906  	return Interval{Int64Comparable(a), Int64Comparable(a + 1)}
   907  }
   908  
   909  type Int64Comparable int64
   910  
   911  func (v Int64Comparable) Compare(c Comparable) int {
   912  	vc := c.(Int64Comparable)
   913  	cmp := v - vc
   914  	if cmp < 0 {
   915  		return -1
   916  	}
   917  	if cmp > 0 {
   918  		return 1
   919  	}
   920  	return 0
   921  }
   922  
   923  // BytesAffineComparable treats empty byte arrays as > all other byte arrays
   924  type BytesAffineComparable []byte
   925  
   926  func (b BytesAffineComparable) Compare(c Comparable) int {
   927  	bc := c.(BytesAffineComparable)
   928  
   929  	if len(b) == 0 {
   930  		if len(bc) == 0 {
   931  			return 0
   932  		}
   933  		return 1
   934  	}
   935  	if len(bc) == 0 {
   936  		return -1
   937  	}
   938  
   939  	return bytes.Compare(b, bc)
   940  }
   941  
   942  func NewBytesAffineInterval(begin, end []byte) Interval {
   943  	return Interval{BytesAffineComparable(begin), BytesAffineComparable(end)}
   944  }
   945  
   946  func NewBytesAffinePoint(b []byte) Interval {
   947  	be := make([]byte, len(b)+1)
   948  	copy(be, b)
   949  	be[len(b)] = 0
   950  	return NewBytesAffineInterval(b, be)
   951  }