Source file src/go.etcd.io/etcd/pkg/v3/adt/interval_tree.go

Documentation: go.etcd.io/etcd/pkg/v3/adt

     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 "red"
    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  //	    RB-DELETE(T, z)
   245  //
   246  //	    y = z
   247  //	    y-original-color = y.color
   248  //
   249  //	    if z.left == T.nil
   250  //	    	x = z.right
   251  //	    	RB-TRANSPLANT(T, z, z.right)
   252  //	    else if z.right == T.nil
   253  //	    	x = z.left
   254  //	    	RB-TRANSPLANT(T, z, z.left)
   255  //	    else
   256  //	    	y = TREE-MINIMUM(z.right)
   257  //	    	y-original-color = y.color
   258  //	    	x = y.right
   259  //	    	if y.p == z
   260  //	    		x.p = y
   261  //	    	else
   262  //	    		RB-TRANSPLANT(T, y, y.right)
   263  //	    		y.right = z.right
   264  //	    		y.right.p = y
   265  //	    	RB-TRANSPLANT(T, z, y)
   266  //	    	y.left = z.left
   267  //	    	y.left.p = y
   268  //	    	y.color = z.color
   269  //
   270  //	    if y-original-color == BLACK
   271  //	    	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  //	RB-DELETE-FIXUP(T, z)
   321  //
   322  //	while x ≠ T.root and x.color == BLACK
   323  //		if x == x.p.left
   324  //			w = x.p.right
   325  //			if w.color == RED
   326  //				w.color = BLACK
   327  //				x.p.color = RED
   328  //				LEFT-ROTATE(T, x, p)
   329  //			if w.left.color == BLACK and w.right.color == BLACK
   330  //				w.color = RED
   331  //				x = x.p
   332  //			else if w.right.color == BLACK
   333  //					w.left.color = BLACK
   334  //					w.color = RED
   335  //					RIGHT-ROTATE(T, w)
   336  //					w = w.p.right
   337  //				w.color = x.p.color
   338  //				x.p.color = BLACK
   339  //				LEFT-ROTATE(T, w.p)
   340  //				x = T.root
   341  //		else
   342  //			w = x.p.left
   343  //			if w.color == RED
   344  //				w.color = BLACK
   345  //				x.p.color = RED
   346  //				RIGHT-ROTATE(T, x, p)
   347  //			if w.right.color == BLACK and w.left.color == BLACK
   348  //				w.color = RED
   349  //				x = x.p
   350  //			else if w.left.color == BLACK
   351  //					w.right.color = BLACK
   352  //					w.color = RED
   353  //					LEFT-ROTATE(T, w)
   354  //					w = w.p.left
   355  //				w.color = x.p.color
   356  //				x.p.color = BLACK
   357  //				RIGHT-ROTATE(T, w.p)
   358  //				x = T.root
   359  //
   360  //	x.color = BLACK
   361  func (ivt *intervalTree) deleteFixup(x *intervalNode) {
   362  	for x != ivt.root && x.color(ivt.sentinel) == black {
   363  		if x == x.parent.left { // line 3-20
   364  			w := x.parent.right
   365  			if w.color(ivt.sentinel) == red {
   366  				w.c = black
   367  				x.parent.c = red
   368  				ivt.rotateLeft(x.parent)
   369  				w = x.parent.right
   370  			}
   371  			if w == nil {
   372  				break
   373  			}
   374  			if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
   375  				w.c = red
   376  				x = x.parent
   377  			} else {
   378  				if w.right.color(ivt.sentinel) == black {
   379  					w.left.c = black
   380  					w.c = red
   381  					ivt.rotateRight(w)
   382  					w = x.parent.right
   383  				}
   384  				w.c = x.parent.color(ivt.sentinel)
   385  				x.parent.c = black
   386  				w.right.c = black
   387  				ivt.rotateLeft(x.parent)
   388  				x = ivt.root
   389  			}
   390  		} else { // line 22-38
   391  			// same as above but with left and right exchanged
   392  			w := x.parent.left
   393  			if w.color(ivt.sentinel) == red {
   394  				w.c = black
   395  				x.parent.c = red
   396  				ivt.rotateRight(x.parent)
   397  				w = x.parent.left
   398  			}
   399  			if w == nil {
   400  				break
   401  			}
   402  			if w.left.color(ivt.sentinel) == black && w.right.color(ivt.sentinel) == black {
   403  				w.c = red
   404  				x = x.parent
   405  			} else {
   406  				if w.left.color(ivt.sentinel) == black {
   407  					w.right.c = black
   408  					w.c = red
   409  					ivt.rotateLeft(w)
   410  					w = x.parent.left
   411  				}
   412  				w.c = x.parent.color(ivt.sentinel)
   413  				x.parent.c = black
   414  				w.left.c = black
   415  				ivt.rotateRight(x.parent)
   416  				x = ivt.root
   417  			}
   418  		}
   419  	}
   420  
   421  	if x != nil {
   422  		x.c = black
   423  	}
   424  }
   425  
   426  func (ivt *intervalTree) createIntervalNode(ivl Interval, val interface{}) *intervalNode {
   427  	return &intervalNode{
   428  		iv:     IntervalValue{ivl, val},
   429  		max:    ivl.End,
   430  		c:      red,
   431  		left:   ivt.sentinel,
   432  		right:  ivt.sentinel,
   433  		parent: ivt.sentinel,
   434  	}
   435  }
   436  
   437  // TODO: make this consistent with textbook implementation
   438  //
   439  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p315
   440  //
   441  //	    RB-INSERT(T, z)
   442  //
   443  //	    y = T.nil
   444  //	    x = T.root
   445  //
   446  //	    while x ≠ T.nil
   447  //	    	y = x
   448  //	    	if z.key < x.key
   449  //	    		x = x.left
   450  //	    	else
   451  //	    		x = x.right
   452  //
   453  //	    z.p = y
   454  //
   455  //	    if y == T.nil
   456  //	    	T.root = z
   457  //	    else if z.key < y.key
   458  //	    	y.left = z
   459  //	    else
   460  //	    	y.right = z
   461  //
   462  //	    z.left = T.nil
   463  //	    z.right = T.nil
   464  //	    z.color = RED
   465  //
   466  //	    RB-INSERT-FIXUP(T, z)
   467  
   468  // Insert adds a node with the given interval into the tree.
   469  func (ivt *intervalTree) Insert(ivl Interval, val interface{}) {
   470  	y := ivt.sentinel
   471  	z := ivt.createIntervalNode(ivl, val)
   472  	x := ivt.root
   473  	for x != ivt.sentinel {
   474  		y = x
   475  		if z.iv.Ivl.Begin.Compare(x.iv.Ivl.Begin) < 0 {
   476  			x = x.left
   477  		} else {
   478  			x = x.right
   479  		}
   480  	}
   481  
   482  	z.parent = y
   483  	if y == ivt.sentinel {
   484  		ivt.root = z
   485  	} else {
   486  		if z.iv.Ivl.Begin.Compare(y.iv.Ivl.Begin) < 0 {
   487  			y.left = z
   488  		} else {
   489  			y.right = z
   490  		}
   491  		y.updateMax(ivt.sentinel)
   492  	}
   493  	z.c = red
   494  
   495  	ivt.insertFixup(z)
   496  	ivt.count++
   497  }
   498  
   499  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.3, p316
   500  //
   501  //	RB-INSERT-FIXUP(T, z)
   502  //
   503  //	while z.p.color == RED
   504  //		if z.p == z.p.p.left
   505  //			y = z.p.p.right
   506  //			if y.color == RED
   507  //				z.p.color = BLACK
   508  //				y.color = BLACK
   509  //				z.p.p.color = RED
   510  //				z = z.p.p
   511  //			else if z == z.p.right
   512  //					z = z.p
   513  //					LEFT-ROTATE(T, z)
   514  //				z.p.color = BLACK
   515  //				z.p.p.color = RED
   516  //				RIGHT-ROTATE(T, z.p.p)
   517  //		else
   518  //			y = z.p.p.left
   519  //			if y.color == RED
   520  //				z.p.color = BLACK
   521  //				y.color = BLACK
   522  //				z.p.p.color = RED
   523  //				z = z.p.p
   524  //			else if z == z.p.right
   525  //					z = z.p
   526  //					RIGHT-ROTATE(T, z)
   527  //				z.p.color = BLACK
   528  //				z.p.p.color = RED
   529  //				LEFT-ROTATE(T, z.p.p)
   530  //
   531  //	T.root.color = BLACK
   532  func (ivt *intervalTree) insertFixup(z *intervalNode) {
   533  	for z.parent.color(ivt.sentinel) == red {
   534  		if z.parent == z.parent.parent.left { // line 3-15
   535  
   536  			y := z.parent.parent.right
   537  			if y.color(ivt.sentinel) == red {
   538  				y.c = black
   539  				z.parent.c = black
   540  				z.parent.parent.c = red
   541  				z = z.parent.parent
   542  			} else {
   543  				if z == z.parent.right {
   544  					z = z.parent
   545  					ivt.rotateLeft(z)
   546  				}
   547  				z.parent.c = black
   548  				z.parent.parent.c = red
   549  				ivt.rotateRight(z.parent.parent)
   550  			}
   551  		} else { // line 16-28
   552  			// same as then with left/right exchanged
   553  			y := z.parent.parent.left
   554  			if y.color(ivt.sentinel) == red {
   555  				y.c = black
   556  				z.parent.c = black
   557  				z.parent.parent.c = red
   558  				z = z.parent.parent
   559  			} else {
   560  				if z == z.parent.left {
   561  					z = z.parent
   562  					ivt.rotateRight(z)
   563  				}
   564  				z.parent.c = black
   565  				z.parent.parent.c = red
   566  				ivt.rotateLeft(z.parent.parent)
   567  			}
   568  		}
   569  	}
   570  
   571  	// line 30
   572  	ivt.root.c = black
   573  }
   574  
   575  // rotateLeft moves x so it is left of its right child
   576  //
   577  // "Introduction to Algorithms" (Cormen et al, 3rd ed.), chapter 13.2, p313
   578  //
   579  //	LEFT-ROTATE(T, x)
   580  //
   581  //	y = x.right
   582  //	x.right = y.left
   583  //
   584  //	if y.left ≠ T.nil
   585  //		y.left.p = x
   586  //
   587  //	y.p = x.p
   588  //
   589  //	if x.p == T.nil
   590  //		T.root = y
   591  //	else if x == x.p.left
   592  //		x.p.left = y
   593  //	else
   594  //		x.p.right = y
   595  //
   596  //	y.left = x
   597  //	x.p = y
   598  func (ivt *intervalTree) rotateLeft(x *intervalNode) {
   599  	// rotateLeft x must have right child
   600  	if x.right == ivt.sentinel {
   601  		return
   602  	}
   603  
   604  	// line 2-3
   605  	y := x.right
   606  	x.right = y.left
   607  
   608  	// line 5-6
   609  	if y.left != ivt.sentinel {
   610  		y.left.parent = x
   611  	}
   612  	x.updateMax(ivt.sentinel)
   613  
   614  	// line 10-15, 18
   615  	ivt.replaceParent(x, y)
   616  
   617  	// line 17
   618  	y.left = x
   619  	y.updateMax(ivt.sentinel)
   620  }
   621  
   622  // rotateRight moves x so it is right of its left child
   623  //
   624  //	RIGHT-ROTATE(T, x)
   625  //
   626  //	y = x.left
   627  //	x.left = y.right
   628  //
   629  //	if y.right ≠ T.nil
   630  //		y.right.p = x
   631  //
   632  //	y.p = x.p
   633  //
   634  //	if x.p == T.nil
   635  //		T.root = y
   636  //	else if x == x.p.right
   637  //		x.p.right = y
   638  //	else
   639  //		x.p.left = y
   640  //
   641  //	y.right = x
   642  //	x.p = y
   643  func (ivt *intervalTree) rotateRight(x *intervalNode) {
   644  	// rotateRight x must have left child
   645  	if x.left == ivt.sentinel {
   646  		return
   647  	}
   648  
   649  	// line 2-3
   650  	y := x.left
   651  	x.left = y.right
   652  
   653  	// line 5-6
   654  	if y.right != ivt.sentinel {
   655  		y.right.parent = x
   656  	}
   657  	x.updateMax(ivt.sentinel)
   658  
   659  	// line 10-15, 18
   660  	ivt.replaceParent(x, y)
   661  
   662  	// line 17
   663  	y.right = x
   664  	y.updateMax(ivt.sentinel)
   665  }
   666  
   667  // replaceParent replaces x's parent with y
   668  func (ivt *intervalTree) replaceParent(x *intervalNode, y *intervalNode) {
   669  	y.parent = x.parent
   670  	if x.parent == ivt.sentinel {
   671  		ivt.root = y
   672  	} else {
   673  		if x == x.parent.left {
   674  			x.parent.left = y
   675  		} else {
   676  			x.parent.right = y
   677  		}
   678  		x.parent.updateMax(ivt.sentinel)
   679  	}
   680  	x.parent = y
   681  }
   682  
   683  // Len gives the number of elements in the tree
   684  func (ivt *intervalTree) Len() int { return ivt.count }
   685  
   686  // Height is the number of levels in the tree; one node has height 1.
   687  func (ivt *intervalTree) Height() int { return ivt.root.height(ivt.sentinel) }
   688  
   689  // MaxHeight is the expected maximum tree height given the number of nodes
   690  func (ivt *intervalTree) MaxHeight() int {
   691  	return int((2 * math.Log2(float64(ivt.Len()+1))) + 0.5)
   692  }
   693  
   694  // IntervalVisitor is used on tree searches; return false to stop searching.
   695  type IntervalVisitor func(n *IntervalValue) bool
   696  
   697  // Visit calls a visitor function on every tree node intersecting the given interval.
   698  // It will visit each interval [x, y) in ascending order sorted on x.
   699  func (ivt *intervalTree) Visit(ivl Interval, ivv IntervalVisitor) {
   700  	ivt.root.visit(&ivl, ivt.sentinel, func(n *intervalNode) bool { return ivv(&n.iv) })
   701  }
   702  
   703  // find the exact node for a given interval
   704  func (ivt *intervalTree) find(ivl Interval) *intervalNode {
   705  	ret := ivt.sentinel
   706  	f := func(n *intervalNode) bool {
   707  		if n.iv.Ivl != ivl {
   708  			return true
   709  		}
   710  		ret = n
   711  		return false
   712  	}
   713  	ivt.root.visit(&ivl, ivt.sentinel, f)
   714  	return ret
   715  }
   716  
   717  // Find gets the IntervalValue for the node matching the given interval
   718  func (ivt *intervalTree) Find(ivl Interval) (ret *IntervalValue) {
   719  	n := ivt.find(ivl)
   720  	if n == ivt.sentinel {
   721  		return nil
   722  	}
   723  	return &n.iv
   724  }
   725  
   726  // Intersects returns true if there is some tree node intersecting the given interval.
   727  func (ivt *intervalTree) Intersects(iv Interval) bool {
   728  	x := ivt.root
   729  	for x != ivt.sentinel && iv.Compare(&x.iv.Ivl) != 0 {
   730  		if x.left != ivt.sentinel && x.left.max.Compare(iv.Begin) > 0 {
   731  			x = x.left
   732  		} else {
   733  			x = x.right
   734  		}
   735  	}
   736  	return x != ivt.sentinel
   737  }
   738  
   739  // Contains returns true if the interval tree's keys cover the entire given interval.
   740  func (ivt *intervalTree) Contains(ivl Interval) bool {
   741  	var maxEnd, minBegin Comparable
   742  
   743  	isContiguous := true
   744  	ivt.Visit(ivl, func(n *IntervalValue) bool {
   745  		if minBegin == nil {
   746  			minBegin = n.Ivl.Begin
   747  			maxEnd = n.Ivl.End
   748  			return true
   749  		}
   750  		if maxEnd.Compare(n.Ivl.Begin) < 0 {
   751  			isContiguous = false
   752  			return false
   753  		}
   754  		if n.Ivl.End.Compare(maxEnd) > 0 {
   755  			maxEnd = n.Ivl.End
   756  		}
   757  		return true
   758  	})
   759  
   760  	return isContiguous && minBegin != nil && maxEnd.Compare(ivl.End) >= 0 && minBegin.Compare(ivl.Begin) <= 0
   761  }
   762  
   763  // Stab returns a slice with all elements in the tree intersecting the interval.
   764  func (ivt *intervalTree) Stab(iv Interval) (ivs []*IntervalValue) {
   765  	if ivt.count == 0 {
   766  		return nil
   767  	}
   768  	f := func(n *IntervalValue) bool { ivs = append(ivs, n); return true }
   769  	ivt.Visit(iv, f)
   770  	return ivs
   771  }
   772  
   773  // Union merges a given interval tree into the receiver.
   774  func (ivt *intervalTree) Union(inIvt IntervalTree, ivl Interval) {
   775  	f := func(n *IntervalValue) bool {
   776  		ivt.Insert(n.Ivl, n.Val)
   777  		return true
   778  	}
   779  	inIvt.Visit(ivl, f)
   780  }
   781  
   782  type visitedInterval struct {
   783  	root  Interval
   784  	left  Interval
   785  	right Interval
   786  	color rbcolor
   787  	depth int
   788  }
   789  
   790  func (vi visitedInterval) String() string {
   791  	bd := new(strings.Builder)
   792  	bd.WriteString(fmt.Sprintf("root [%v,%v,%v], left [%v,%v], right [%v,%v], depth %d",
   793  		vi.root.Begin, vi.root.End, vi.color,
   794  		vi.left.Begin, vi.left.End,
   795  		vi.right.Begin, vi.right.End,
   796  		vi.depth,
   797  	))
   798  	return bd.String()
   799  }
   800  
   801  // visitLevel traverses tree in level order.
   802  // used for testing
   803  func (ivt *intervalTree) visitLevel() []visitedInterval {
   804  	if ivt.root == ivt.sentinel {
   805  		return nil
   806  	}
   807  
   808  	rs := make([]visitedInterval, 0, ivt.Len())
   809  
   810  	type pair struct {
   811  		node  *intervalNode
   812  		depth int
   813  	}
   814  	queue := []pair{{ivt.root, 0}}
   815  	for len(queue) > 0 {
   816  		f := queue[0]
   817  		queue = queue[1:]
   818  
   819  		vi := visitedInterval{
   820  			root:  f.node.iv.Ivl,
   821  			color: f.node.color(ivt.sentinel),
   822  			depth: f.depth,
   823  		}
   824  		if f.node.left != ivt.sentinel {
   825  			vi.left = f.node.left.iv.Ivl
   826  			queue = append(queue, pair{f.node.left, f.depth + 1})
   827  		}
   828  		if f.node.right != ivt.sentinel {
   829  			vi.right = f.node.right.iv.Ivl
   830  			queue = append(queue, pair{f.node.right, f.depth + 1})
   831  		}
   832  
   833  		rs = append(rs, vi)
   834  	}
   835  
   836  	return rs
   837  }
   838  
   839  type StringComparable string
   840  
   841  func (s StringComparable) Compare(c Comparable) int {
   842  	sc := c.(StringComparable)
   843  	if s < sc {
   844  		return -1
   845  	}
   846  	if s > sc {
   847  		return 1
   848  	}
   849  	return 0
   850  }
   851  
   852  func NewStringInterval(begin, end string) Interval {
   853  	return Interval{StringComparable(begin), StringComparable(end)}
   854  }
   855  
   856  func NewStringPoint(s string) Interval {
   857  	return Interval{StringComparable(s), StringComparable(s + "\x00")}
   858  }
   859  
   860  // StringAffineComparable treats "" as > all other strings
   861  type StringAffineComparable string
   862  
   863  func (s StringAffineComparable) Compare(c Comparable) int {
   864  	sc := c.(StringAffineComparable)
   865  
   866  	if len(s) == 0 {
   867  		if len(sc) == 0 {
   868  			return 0
   869  		}
   870  		return 1
   871  	}
   872  	if len(sc) == 0 {
   873  		return -1
   874  	}
   875  
   876  	if s < sc {
   877  		return -1
   878  	}
   879  	if s > sc {
   880  		return 1
   881  	}
   882  	return 0
   883  }
   884  
   885  func NewStringAffineInterval(begin, end string) Interval {
   886  	return Interval{StringAffineComparable(begin), StringAffineComparable(end)}
   887  }
   888  
   889  func NewStringAffinePoint(s string) Interval {
   890  	return NewStringAffineInterval(s, s+"\x00")
   891  }
   892  
   893  func NewInt64Interval(a int64, b int64) Interval {
   894  	return Interval{Int64Comparable(a), Int64Comparable(b)}
   895  }
   896  
   897  func newInt64EmptyInterval() Interval {
   898  	return Interval{Begin: nil, End: nil}
   899  }
   900  
   901  func NewInt64Point(a int64) Interval {
   902  	return Interval{Int64Comparable(a), Int64Comparable(a + 1)}
   903  }
   904  
   905  type Int64Comparable int64
   906  
   907  func (v Int64Comparable) Compare(c Comparable) int {
   908  	vc := c.(Int64Comparable)
   909  	cmp := v - vc
   910  	if cmp < 0 {
   911  		return -1
   912  	}
   913  	if cmp > 0 {
   914  		return 1
   915  	}
   916  	return 0
   917  }
   918  
   919  // BytesAffineComparable treats empty byte arrays as > all other byte arrays
   920  type BytesAffineComparable []byte
   921  
   922  func (b BytesAffineComparable) Compare(c Comparable) int {
   923  	bc := c.(BytesAffineComparable)
   924  
   925  	if len(b) == 0 {
   926  		if len(bc) == 0 {
   927  			return 0
   928  		}
   929  		return 1
   930  	}
   931  	if len(bc) == 0 {
   932  		return -1
   933  	}
   934  
   935  	return bytes.Compare(b, bc)
   936  }
   937  
   938  func NewBytesAffineInterval(begin, end []byte) Interval {
   939  	return Interval{BytesAffineComparable(begin), BytesAffineComparable(end)}
   940  }
   941  
   942  func NewBytesAffinePoint(b []byte) Interval {
   943  	be := make([]byte, len(b)+1)
   944  	copy(be, b)
   945  	be[len(b)] = 0
   946  	return NewBytesAffineInterval(b, be)
   947  }
   948  

View as plain text