github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/interval/btree_based_interval.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  //
    11  // This code is based on: https://github.com/google/btree.
    12  
    13  package interval
    14  
    15  import (
    16  	"sort"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    19  	"github.com/cockroachdb/errors"
    20  )
    21  
    22  const (
    23  	// DefaultBTreeMinimumDegree is the default B-tree minimum degree. Benchmarks
    24  	// show that the interval tree performs best with this minimum degree.
    25  	DefaultBTreeMinimumDegree = 32
    26  	// DefaultBTreeFreeListSize is the default size of a B-tree's freelist.
    27  	DefaultBTreeFreeListSize = 32
    28  )
    29  
    30  var (
    31  	nilItems    = make(items, 16)
    32  	nilChildren = make(children, 16)
    33  )
    34  
    35  // FreeList represents a free list of btree nodes. By default each
    36  // BTree has its own FreeList, but multiple BTrees can share the same
    37  // FreeList.
    38  // Two Btrees using the same freelist are safe for concurrent write access.
    39  type FreeList struct {
    40  	mu       syncutil.Mutex
    41  	freelist []*node
    42  }
    43  
    44  // NewFreeList creates a new free list.
    45  // size is the maximum size of the returned free list.
    46  func NewFreeList(size int) *FreeList {
    47  	return &FreeList{freelist: make([]*node, 0, size)}
    48  }
    49  
    50  func (f *FreeList) newNode() (n *node) {
    51  	f.mu.Lock()
    52  	index := len(f.freelist) - 1
    53  	if index < 0 {
    54  		f.mu.Unlock()
    55  		return new(node)
    56  	}
    57  	n = f.freelist[index]
    58  	f.freelist[index] = nil
    59  	f.freelist = f.freelist[:index]
    60  	f.mu.Unlock()
    61  	return
    62  }
    63  
    64  // freeNode adds the given node to the list, returning true if it was added
    65  // and false if it was discarded.
    66  func (f *FreeList) freeNode(n *node) (out bool) {
    67  	f.mu.Lock()
    68  	if len(f.freelist) < cap(f.freelist) {
    69  		f.freelist = append(f.freelist, n)
    70  		out = true
    71  	}
    72  	f.mu.Unlock()
    73  	return
    74  }
    75  
    76  // newBTree creates a new interval tree with the given overlapper function and
    77  // the default B-tree minimum degree.
    78  func newBTree(overlapper Overlapper) *btree {
    79  	return newBTreeWithDegree(overlapper, DefaultBTreeMinimumDegree)
    80  }
    81  
    82  // newBTreeWithDegree creates a new interval tree with the given overlapper
    83  // function and the given minimum degree. A minimum degree less than 2 will
    84  // cause a panic.
    85  //
    86  // newBTreeWithDegree(overlapper, 2), for example, will create a 2-3-4 tree (each
    87  // node contains 1-3 Interfaces and 2-4 children).
    88  func newBTreeWithDegree(overlapper Overlapper, minimumDegree int) *btree {
    89  	if minimumDegree < 2 {
    90  		panic("bad minimum degree")
    91  	}
    92  	f := NewFreeList(DefaultBTreeFreeListSize)
    93  	return &btree{
    94  		minimumDegree: minimumDegree,
    95  		overlapper:    overlapper,
    96  		cow:           &copyOnWriteContext{freelist: f},
    97  	}
    98  }
    99  
   100  func isValidInterface(a Interface) error {
   101  	if a == nil {
   102  		// Note: Newf instead of New so that the error message is revealed
   103  		// in redact calls.
   104  		return errors.Newf("nil interface")
   105  	}
   106  	r := a.Range()
   107  	return rangeError(r)
   108  }
   109  
   110  // interfaces stores Interfaces sorted by Range().End in a node.
   111  type items []Interface
   112  
   113  // insertAt inserts a value into the given index, pushing all subsequent values
   114  // forward.
   115  func (s *items) insertAt(index int, e Interface) {
   116  	oldLen := len(*s)
   117  	*s = append(*s, nil)
   118  	if index < oldLen {
   119  		copy((*s)[index+1:], (*s)[index:])
   120  	}
   121  	(*s)[index] = e
   122  }
   123  
   124  // removeAt removes a value at a given index, pulling all subsequent values
   125  // back.
   126  func (s *items) removeAt(index int) Interface {
   127  	e := (*s)[index]
   128  	copy((*s)[index:], (*s)[index+1:])
   129  	(*s)[len(*s)-1] = nil
   130  	*s = (*s)[:len(*s)-1]
   131  	return e
   132  }
   133  
   134  // pop removes and returns the last element in the list.
   135  func (s *items) pop() (out Interface) {
   136  	index := len(*s) - 1
   137  	out = (*s)[index]
   138  	(*s)[index] = nil
   139  	*s = (*s)[:index]
   140  	return
   141  }
   142  
   143  // truncate truncates this instance at index so that it contains only the
   144  // first index items. index must be less than or equal to length.
   145  func (s *items) truncate(index int) {
   146  	var toClear items
   147  	*s, toClear = (*s)[:index], (*s)[index:]
   148  	for len(toClear) > 0 {
   149  		toClear = toClear[copy(toClear, nilItems):]
   150  	}
   151  }
   152  
   153  // find returns the index where the given Interface should be inserted into this
   154  // list. 'found' is true if the interface already exists in the list at the
   155  // given index.
   156  func (s items) find(e Interface) (index int, found bool) {
   157  	i := sort.Search(len(s), func(i int) bool {
   158  		return Compare(e, s[i]) < 0
   159  	})
   160  	if i > 0 && Equal(s[i-1], e) {
   161  		return i - 1, true
   162  	}
   163  	return i, false
   164  }
   165  
   166  // children stores child nodes sorted by Range.End in a node.
   167  type children []*node
   168  
   169  // insertAt inserts a value into the given index, pushing all subsequent values
   170  // forward.
   171  func (s *children) insertAt(index int, n *node) {
   172  	oldLen := len(*s)
   173  	*s = append(*s, nil)
   174  	if index < oldLen {
   175  		copy((*s)[index+1:], (*s)[index:])
   176  	}
   177  	(*s)[index] = n
   178  }
   179  
   180  // removeAt removes a value at a given index, pulling all subsequent values
   181  // back.
   182  func (s *children) removeAt(index int) *node {
   183  	n := (*s)[index]
   184  	copy((*s)[index:], (*s)[index+1:])
   185  	(*s)[len(*s)-1] = nil
   186  	*s = (*s)[:len(*s)-1]
   187  	return n
   188  }
   189  
   190  // pop removes and returns the last element in the list.
   191  func (s *children) pop() (out *node) {
   192  	index := len(*s) - 1
   193  	out = (*s)[index]
   194  	(*s)[index] = nil
   195  	*s = (*s)[:index]
   196  	return
   197  }
   198  
   199  // truncate truncates this instance at index so that it contains only the
   200  // first index children. index must be less than or equal to length.
   201  func (s *children) truncate(index int) {
   202  	var toClear children
   203  	*s, toClear = (*s)[:index], (*s)[index:]
   204  	for len(toClear) > 0 {
   205  		toClear = toClear[copy(toClear, nilChildren):]
   206  	}
   207  }
   208  
   209  // node is an internal node in a tree.
   210  //
   211  // It must at all times maintain the invariant that either
   212  //   * len(children) == 0, len(interfaces) unconstrained
   213  //   * len(children) == len(interfaces) + 1
   214  type node struct {
   215  	// Range is the node range which covers all the ranges in the subtree rooted
   216  	// at the node. Range.Start is the leftmost position. Range.End is the
   217  	// rightmost position. Here we follow the approach employed by
   218  	// https://github.com/biogo/store/tree/master/interval since it make it easy
   219  	// to analyze the traversal of intervals which overlaps with a given interval.
   220  	// CLRS only uses Range.End.
   221  	Range    Range
   222  	items    items
   223  	children children
   224  	cow      *copyOnWriteContext
   225  }
   226  
   227  func (n *node) mutableFor(cow *copyOnWriteContext) *node {
   228  	if n.cow == cow {
   229  		return n
   230  	}
   231  	out := cow.newNode()
   232  	out.Range = n.Range
   233  	if cap(out.items) >= len(n.items) {
   234  		out.items = out.items[:len(n.items)]
   235  	} else {
   236  		out.items = make(items, len(n.items), cap(n.items))
   237  	}
   238  	copy(out.items, n.items)
   239  	// Copy children
   240  	if cap(out.children) >= len(n.children) {
   241  		out.children = out.children[:len(n.children)]
   242  	} else {
   243  		out.children = make(children, len(n.children), cap(n.children))
   244  	}
   245  	copy(out.children, n.children)
   246  	return out
   247  }
   248  
   249  func (n *node) mutableChild(i int) *node {
   250  	c := n.children[i].mutableFor(n.cow)
   251  	n.children[i] = c
   252  	return c
   253  }
   254  
   255  // split splits the given node at the given index. The current node shrinks, and
   256  // this function returns the Interface that existed at that index and a new node
   257  // containing all interfaces/children after it. Before splitting:
   258  //
   259  //          +-----------+
   260  //          |   x y z   |
   261  //          ---/-/-\-\--+
   262  //
   263  // After splitting:
   264  //
   265  //          +-----------+
   266  //          |     y     |
   267  //          -----/-\----+
   268  //              /   \
   269  //             v     v
   270  // +-----------+     +-----------+
   271  // |         x |     | z         |
   272  // +-----------+     +-----------+
   273  //
   274  func (n *node) split(i int, fast bool) (Interface, *node) {
   275  	e := n.items[i]
   276  	second := n.cow.newNode()
   277  	second.items = append(second.items, n.items[i+1:]...)
   278  	n.items.truncate(i)
   279  	if len(n.children) > 0 {
   280  		second.children = append(second.children, n.children[i+1:]...)
   281  		n.children.truncate(i + 1)
   282  	}
   283  	if !fast {
   284  		// adjust range for the first split part
   285  		oldRangeEnd := n.Range.End
   286  		n.Range.End = n.rangeEnd()
   287  
   288  		// adjust range for the second split part
   289  		second.Range.Start = second.rangeStart()
   290  		if n.Range.End.Equal(oldRangeEnd) || e.Range().End.Equal(oldRangeEnd) {
   291  			second.Range.End = second.rangeEnd()
   292  		} else {
   293  			second.Range.End = oldRangeEnd
   294  		}
   295  	}
   296  	return e, second
   297  }
   298  
   299  // maybeSplitChild checks if a child should be split, and if so splits it.
   300  // Returns whether or not a split occurred.
   301  func (n *node) maybeSplitChild(i, maxItems int, fast bool) bool {
   302  	if len(n.children[i].items) < maxItems {
   303  		return false
   304  	}
   305  	first := n.mutableChild(i)
   306  	e, second := first.split(maxItems/2, fast)
   307  	n.items.insertAt(i, e)
   308  	n.children.insertAt(i+1, second)
   309  	return true
   310  }
   311  
   312  // insert inserts an Interface into the subtree rooted at this node, making sure
   313  // no nodes in the subtree exceed maxItems Interfaces.
   314  func (n *node) insert(e Interface, maxItems int, fast bool) (out Interface, extended bool) {
   315  	i, found := n.items.find(e)
   316  	if found {
   317  		out = n.items[i]
   318  		n.items[i] = e
   319  		return
   320  	}
   321  	if len(n.children) == 0 {
   322  		n.items.insertAt(i, e)
   323  		out = nil
   324  		if !fast {
   325  			if i == 0 {
   326  				extended = true
   327  				n.Range.Start = n.items[0].Range().Start
   328  			}
   329  			if n.items[i].Range().End.Compare(n.Range.End) > 0 {
   330  				extended = true
   331  				n.Range.End = n.items[i].Range().End
   332  			}
   333  		}
   334  		return
   335  	}
   336  	if n.maybeSplitChild(i, maxItems, fast) {
   337  		inTree := n.items[i]
   338  		switch Compare(e, inTree) {
   339  		case -1:
   340  			// no change, we want first split node
   341  		case 1:
   342  			i++ // we want second split node
   343  		default:
   344  			out = n.items[i]
   345  			n.items[i] = e
   346  			return
   347  		}
   348  	}
   349  	out, extended = n.mutableChild(i).insert(e, maxItems, fast)
   350  	if !fast && extended {
   351  		extended = false
   352  		if i == 0 && n.children[0].Range.Start.Compare(n.Range.Start) < 0 {
   353  			extended = true
   354  			n.Range.Start = n.children[0].Range.Start
   355  		}
   356  		if n.children[i].Range.End.Compare(n.Range.End) > 0 {
   357  			extended = true
   358  			n.Range.End = n.children[i].Range.End
   359  		}
   360  	}
   361  	return
   362  }
   363  
   364  func (t *btree) isEmpty() bool {
   365  	return t.root == nil || len(t.root.items) == 0
   366  }
   367  
   368  func (t *btree) Get(r Range) (o []Interface) {
   369  	return t.GetWithOverlapper(r, t.overlapper)
   370  }
   371  
   372  func (t *btree) GetWithOverlapper(r Range, overlapper Overlapper) (o []Interface) {
   373  	if err := rangeError(r); err != nil {
   374  		return
   375  	}
   376  	if !t.overlappable(r) {
   377  		return
   378  	}
   379  	t.root.doMatch(func(e Interface) (done bool) { o = append(o, e); return }, r, overlapper)
   380  	return
   381  }
   382  
   383  func (t *btree) DoMatching(fn Operation, r Range) bool {
   384  	if err := rangeError(r); err != nil {
   385  		return false
   386  	}
   387  	if !t.overlappable(r) {
   388  		return false
   389  	}
   390  	return t.root.doMatch(fn, r, t.overlapper)
   391  }
   392  
   393  func (t *btree) overlappable(r Range) bool {
   394  	if t.isEmpty() || !t.overlapper.Overlap(r, t.root.Range) {
   395  		return false
   396  	}
   397  	return true
   398  }
   399  
   400  // benchmarks show that if Comparable.Compare is invoked directly instead of
   401  // through an indirection with Overlapper, Insert, Delete and a traversal to
   402  // visit overlapped intervals have a noticeable speed-up. So two versions of
   403  // doMatch are created. One is for InclusiveOverlapper. The other is for
   404  // ExclusiveOverlapper.
   405  func (n *node) doMatch(fn Operation, r Range, overlapper Overlapper) (done bool) {
   406  	if overlapper == InclusiveOverlapper {
   407  		return n.inclusiveDoMatch(fn, r, overlapper)
   408  	}
   409  	return n.exclusiveDoMatch(fn, r, overlapper)
   410  }
   411  
   412  // doMatch for InclusiveOverlapper.
   413  func (n *node) inclusiveDoMatch(fn Operation, r Range, overlapper Overlapper) (done bool) {
   414  	length := sort.Search(len(n.items), func(i int) bool {
   415  		return n.items[i].Range().Start.Compare(r.End) > 0
   416  	})
   417  
   418  	if len(n.children) == 0 {
   419  		for _, e := range n.items[:length] {
   420  			if r.Start.Compare(e.Range().End) <= 0 {
   421  				if done = fn(e); done {
   422  					return
   423  				}
   424  			}
   425  		}
   426  		return
   427  	}
   428  
   429  	for i := 0; i < length; i++ {
   430  		c := n.children[i]
   431  		if r.Start.Compare(c.Range.End) <= 0 {
   432  			if done = c.inclusiveDoMatch(fn, r, overlapper); done {
   433  				return
   434  			}
   435  		}
   436  		e := n.items[i]
   437  		if r.Start.Compare(e.Range().End) <= 0 {
   438  			if done = fn(e); done {
   439  				return
   440  			}
   441  		}
   442  	}
   443  
   444  	if overlapper.Overlap(r, n.children[length].Range) {
   445  		done = n.children[length].inclusiveDoMatch(fn, r, overlapper)
   446  	}
   447  	return
   448  }
   449  
   450  // doMatch for ExclusiveOverlapper.
   451  func (n *node) exclusiveDoMatch(fn Operation, r Range, overlapper Overlapper) (done bool) {
   452  	length := sort.Search(len(n.items), func(i int) bool {
   453  		return n.items[i].Range().Start.Compare(r.End) >= 0
   454  	})
   455  
   456  	if len(n.children) == 0 {
   457  		for _, e := range n.items[:length] {
   458  			if r.Start.Compare(e.Range().End) < 0 {
   459  				if done = fn(e); done {
   460  					return
   461  				}
   462  			}
   463  		}
   464  		return
   465  	}
   466  
   467  	for i := 0; i < length; i++ {
   468  		c := n.children[i]
   469  		if r.Start.Compare(c.Range.End) < 0 {
   470  			if done = c.exclusiveDoMatch(fn, r, overlapper); done {
   471  				return
   472  			}
   473  		}
   474  		e := n.items[i]
   475  		if r.Start.Compare(e.Range().End) < 0 {
   476  			if done = fn(e); done {
   477  				return
   478  			}
   479  		}
   480  	}
   481  
   482  	if overlapper.Overlap(r, n.children[length].Range) {
   483  		done = n.children[length].exclusiveDoMatch(fn, r, overlapper)
   484  	}
   485  	return
   486  }
   487  
   488  func (t *btree) Do(fn Operation) bool {
   489  	if t.root == nil {
   490  		return false
   491  	}
   492  	return t.root.do(fn)
   493  }
   494  
   495  func (n *node) do(fn Operation) (done bool) {
   496  	cLen := len(n.children)
   497  	if cLen == 0 {
   498  		for _, e := range n.items {
   499  			if done = fn(e); done {
   500  				return
   501  			}
   502  		}
   503  		return
   504  	}
   505  
   506  	for i := 0; i < cLen-1; i++ {
   507  		c := n.children[i]
   508  		if done = c.do(fn); done {
   509  			return
   510  		}
   511  		e := n.items[i]
   512  		if done = fn(e); done {
   513  			return
   514  		}
   515  	}
   516  	done = n.children[cLen-1].do(fn)
   517  	return
   518  }
   519  
   520  // toRemove details what interface to remove in a node.remove call.
   521  type toRemove int
   522  
   523  const (
   524  	removeItem toRemove = iota // removes the given interface
   525  	removeMin                  // removes smallest interface in the subtree
   526  	removeMax                  // removes largest interface in the subtree
   527  )
   528  
   529  // remove removes an interface from the subtree rooted at this node.
   530  func (n *node) remove(
   531  	e Interface, minItems int, typ toRemove, fast bool,
   532  ) (out Interface, shrunk bool) {
   533  	var i int
   534  	var found bool
   535  	switch typ {
   536  	case removeMax:
   537  		if len(n.children) == 0 {
   538  			return n.removeFromLeaf(len(n.items)-1, fast)
   539  		}
   540  		i = len(n.items)
   541  	case removeMin:
   542  		if len(n.children) == 0 {
   543  			return n.removeFromLeaf(0, fast)
   544  		}
   545  		i = 0
   546  	case removeItem:
   547  		i, found = n.items.find(e)
   548  		if len(n.children) == 0 {
   549  			if found {
   550  				return n.removeFromLeaf(i, fast)
   551  			}
   552  			return
   553  		}
   554  	default:
   555  		panic("invalid remove type")
   556  	}
   557  	// If we get to here, we have children.
   558  	if len(n.children[i].items) <= minItems {
   559  		out, shrunk = n.growChildAndRemove(i, e, minItems, typ, fast)
   560  		return
   561  	}
   562  	child := n.mutableChild(i)
   563  	// Either we had enough interfaces to begin with, or we've done some
   564  	// merging/stealing, because we've got enough now and we're ready to return
   565  	// stuff.
   566  	if found {
   567  		// The interface exists at index 'i', and the child we've selected can give
   568  		// us a predecessor, since if we've gotten here it's got > minItems
   569  		// interfaces in it.
   570  		out = n.items[i]
   571  		// We use our special-case 'remove' call with typ=removeMax to pull the
   572  		// predecessor of interface i (the rightmost leaf of our immediate left
   573  		// child) and set it into where we pulled the interface from.
   574  		n.items[i], _ = child.remove(nil, minItems, removeMax, fast)
   575  		if !fast {
   576  			shrunk = n.adjustRangeEndForRemoval(out, nil)
   577  		}
   578  		return
   579  	}
   580  	// Final recursive call. Once we're here, we know that the interface isn't in
   581  	// this node and that the child is big enough to remove from.
   582  	out, shrunk = child.remove(e, minItems, typ, fast)
   583  	if !fast && shrunk {
   584  		shrunkOnStart := false
   585  		if i == 0 {
   586  			if n.Range.Start.Compare(child.Range.Start) < 0 {
   587  				shrunkOnStart = true
   588  				n.Range.Start = child.Range.Start
   589  			}
   590  		}
   591  		shrunkOnEnd := n.adjustRangeEndForRemoval(out, nil)
   592  		shrunk = shrunkOnStart || shrunkOnEnd
   593  	}
   594  	return
   595  }
   596  
   597  // adjustRangeEndForRemoval adjusts Range.End for the node after an interface
   598  // and/or a child is removed.
   599  func (n *node) adjustRangeEndForRemoval(e Interface, c *node) (decreased bool) {
   600  	if (e != nil && e.Range().End.Equal(n.Range.End)) || (c != nil && c.Range.End.Equal(n.Range.End)) {
   601  		newEnd := n.rangeEnd()
   602  		if n.Range.End.Compare(newEnd) > 0 {
   603  			decreased = true
   604  			n.Range.End = newEnd
   605  		}
   606  	}
   607  	return
   608  }
   609  
   610  // removeFromLeaf removes children[i] from the leaf node.
   611  func (n *node) removeFromLeaf(i int, fast bool) (out Interface, shrunk bool) {
   612  	if i == len(n.items)-1 {
   613  		out = n.items.pop()
   614  	} else {
   615  		out = n.items.removeAt(i)
   616  	}
   617  	if !fast && len(n.items) > 0 {
   618  		shrunkOnStart := false
   619  		if i == 0 {
   620  			oldStart := n.Range.Start
   621  			n.Range.Start = n.items[0].Range().Start
   622  			if !n.Range.Start.Equal(oldStart) {
   623  				shrunkOnStart = true
   624  			}
   625  		}
   626  		shrunkOnEnd := n.adjustRangeEndForRemoval(out, nil)
   627  		shrunk = shrunkOnStart || shrunkOnEnd
   628  	}
   629  	return
   630  }
   631  
   632  // growChildAndRemove grows child 'i' to make sure it's possible to remove an
   633  // Interface from it while keeping it at minItems, then calls remove to
   634  // actually remove it.
   635  //
   636  // Most documentation says we have to do two sets of special casing:
   637  //   1) interface is in this node
   638  //   2) interface is in child
   639  // In both cases, we need to handle the two subcases:
   640  //   A) node has enough values that it can spare one
   641  //   B) node doesn't have enough values
   642  // For the latter, we have to check:
   643  //   a) left sibling has node to spare
   644  //   b) right sibling has node to spare
   645  //   c) we must merge
   646  // To simplify our code here, we handle cases #1 and #2 the same:
   647  // If a node doesn't have enough Interfaces, we make sure it does (using a,b,c).
   648  // We then simply redo our remove call, and the second time (regardless of
   649  // whether we're in case 1 or 2), we'll have enough Interfaces and can guarantee
   650  // that we hit case A.
   651  func (n *node) growChildAndRemove(
   652  	i int, e Interface, minItems int, typ toRemove, fast bool,
   653  ) (out Interface, shrunk bool) {
   654  	if i > 0 && len(n.children[i-1].items) > minItems {
   655  		n.stealFromLeftChild(i, fast)
   656  	} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
   657  		n.stealFromRightChild(i, fast)
   658  	} else {
   659  		if i >= len(n.items) {
   660  			i--
   661  		}
   662  		n.mergeWithRightChild(i, fast)
   663  	}
   664  	return n.remove(e, minItems, typ, fast)
   665  }
   666  
   667  // Steal from left child. Before stealing:
   668  //
   669  //          +-----------+
   670  //          |     y     |
   671  //          -----/-\----+
   672  //              /   \
   673  //             v     v
   674  // +-----------+     +-----------+
   675  // |         x |     |           |
   676  // +----------\+     +-----------+
   677  //             \
   678  //              v
   679  //              a
   680  //
   681  // After stealing:
   682  //
   683  //          +-----------+
   684  //          |     x     |
   685  //          -----/-\----+
   686  //              /   \
   687  //             v     v
   688  // +-----------+     +-----------+
   689  // |           |     | y         |
   690  // +-----------+     +/----------+
   691  //                   /
   692  //                  v
   693  //                  a
   694  //
   695  func (n *node) stealFromLeftChild(i int, fast bool) {
   696  	// steal
   697  	stealTo := n.mutableChild(i)
   698  	stealFrom := n.mutableChild(i - 1)
   699  	x := stealFrom.items.pop()
   700  	y := n.items[i-1]
   701  	stealTo.items.insertAt(0, y)
   702  	n.items[i-1] = x
   703  	var a *node
   704  	if len(stealFrom.children) > 0 {
   705  		a = stealFrom.children.pop()
   706  		stealTo.children.insertAt(0, a)
   707  	}
   708  
   709  	if !fast {
   710  		// adjust range for stealFrom
   711  		stealFrom.adjustRangeEndForRemoval(x, a)
   712  
   713  		// adjust range for stealTo
   714  		stealTo.Range.Start = stealTo.rangeStart()
   715  		if y.Range().End.Compare(stealTo.Range.End) > 0 {
   716  			stealTo.Range.End = y.Range().End
   717  		}
   718  		if a != nil && a.Range.End.Compare(stealTo.Range.End) > 0 {
   719  			stealTo.Range.End = a.Range.End
   720  		}
   721  	}
   722  }
   723  
   724  // Steal from right child. Before stealing:
   725  //
   726  //          +-----------+
   727  //          |     y     |
   728  //          -----/-\----+
   729  //              /   \
   730  //             v     v
   731  // +-----------+     +-----------+
   732  // |           |     | x         |
   733  // +---------- +     +/----------+
   734  //                   /
   735  //                  v
   736  //                  a
   737  //
   738  // After stealing:
   739  //
   740  //          +-----------+
   741  //          |     x     |
   742  //          -----/-\----+
   743  //              /   \
   744  //             v     v
   745  // +-----------+     +-----------+
   746  // |         y |     |           |
   747  // +----------\+     +-----------+
   748  //             \
   749  //              v
   750  //              a
   751  //
   752  func (n *node) stealFromRightChild(i int, fast bool) {
   753  	// steal
   754  	stealTo := n.mutableChild(i)
   755  	stealFrom := n.mutableChild(i + 1)
   756  	x := stealFrom.items.removeAt(0)
   757  	y := n.items[i]
   758  	stealTo.items = append(stealTo.items, y)
   759  	n.items[i] = x
   760  	var a *node
   761  	if len(stealFrom.children) > 0 {
   762  		a = stealFrom.children.removeAt(0)
   763  		stealTo.children = append(stealTo.children, a)
   764  	}
   765  
   766  	if !fast {
   767  		// adjust range for stealFrom
   768  		stealFrom.Range.Start = stealFrom.rangeStart()
   769  		stealFrom.adjustRangeEndForRemoval(x, a)
   770  
   771  		// adjust range for stealTo
   772  		if y.Range().End.Compare(stealTo.Range.End) > 0 {
   773  			stealTo.Range.End = y.Range().End
   774  		}
   775  		if a != nil && a.Range.End.Compare(stealTo.Range.End) > 0 {
   776  			stealTo.Range.End = a.Range.End
   777  		}
   778  	}
   779  }
   780  
   781  // Merge with right child. Before merging:
   782  //
   783  //          +-----------+
   784  //          |   u y v   |
   785  //          -----/-\----+
   786  //              /   \
   787  //             v     v
   788  // +-----------+     +-----------+
   789  // |         x |     | z         |
   790  // +---------- +     +-----------+
   791  //
   792  // After merging:
   793  //
   794  //          +-----------+
   795  //          |    u v    |
   796  //          ------|-----+
   797  //                |
   798  //                v
   799  //          +-----------+
   800  //          |   x y z   |
   801  //          +---------- +
   802  //
   803  func (n *node) mergeWithRightChild(i int, fast bool) {
   804  	// merge
   805  	child := n.mutableChild(i)
   806  	mergeItem := n.items.removeAt(i)
   807  	mergeChild := n.children.removeAt(i + 1)
   808  	child.items = append(child.items, mergeItem)
   809  	child.items = append(child.items, mergeChild.items...)
   810  	child.children = append(child.children, mergeChild.children...)
   811  
   812  	if !fast {
   813  		if mergeItem.Range().End.Compare(child.Range.End) > 0 {
   814  			child.Range.End = mergeItem.Range().End
   815  		}
   816  		if mergeChild.Range.End.Compare(child.Range.End) > 0 {
   817  			child.Range.End = mergeChild.Range.End
   818  		}
   819  	}
   820  	n.cow.freeNode(mergeChild)
   821  }
   822  
   823  var _ Tree = (*btree)(nil)
   824  
   825  // btree is an interval tree based on an augmented BTree.
   826  //
   827  // Tree stores Instances in an ordered structure, allowing easy insertion,
   828  // removal, and iteration.
   829  //
   830  // Write operations are not safe for concurrent mutation by multiple
   831  // goroutines, but Read operations are.
   832  type btree struct {
   833  	length        int
   834  	minimumDegree int
   835  	overlapper    Overlapper
   836  	root          *node
   837  	cow           *copyOnWriteContext
   838  }
   839  
   840  // copyOnWriteContext pointers determine node ownership... a tree with a write
   841  // context equivalent to a node's write context is allowed to modify that node.
   842  // A tree whose write context does not match a node's is not allowed to modify
   843  // it, and must create a new, writable copy (IE: it's a Clone).
   844  //
   845  // When doing any write operation, we maintain the invariant that the current
   846  // node's context is equal to the context of the tree that requested the write.
   847  // We do this by, before we descend into any node, creating a copy with the
   848  // correct context if the contexts don't match.
   849  //
   850  // Since the node we're currently visiting on any write has the requesting
   851  // tree's context, that node is modifiable in place.  Children of that node may
   852  // not share context, but before we descend into them, we'll make a mutable
   853  // copy.
   854  type copyOnWriteContext struct {
   855  	freelist *FreeList
   856  }
   857  
   858  // cloneInternal clones the btree, lazily.  Clone should not be called concurrently,
   859  // but the original tree (t) and the new tree (t2) can be used concurrently
   860  // once the Clone call completes.
   861  //
   862  // The internal tree structure of b is marked read-only and shared between t and
   863  // t2.  Writes to both t and t2 use copy-on-write logic, creating new nodes
   864  // whenever one of b's original nodes would have been modified.  Read operations
   865  // should have no performance degredation.  Write operations for both t and t2
   866  // will initially experience minor slow-downs caused by additional allocs and
   867  // copies due to the aforementioned copy-on-write logic, but should converge to
   868  // the original performance characteristics of the original tree.
   869  func (t *btree) cloneInternal() (t2 *btree) {
   870  	// Create two entirely new copy-on-write contexts.
   871  	// This operation effectively creates three trees:
   872  	//   the original, shared nodes (old b.cow)
   873  	//   the new b.cow nodes
   874  	//   the new out.cow nodes
   875  	cow1, cow2 := *t.cow, *t.cow
   876  	out := *t
   877  	t.cow = &cow1
   878  	out.cow = &cow2
   879  	return &out
   880  }
   881  
   882  // Clone clones the btree, lazily.
   883  func (t *btree) Clone() Tree {
   884  	return t.cloneInternal()
   885  }
   886  
   887  // adjustRange sets the Range to the maximum extent of the childrens' Range
   888  // spans and its range spans.
   889  func (n *node) adjustRange() {
   890  	n.Range.Start = n.rangeStart()
   891  	n.Range.End = n.rangeEnd()
   892  }
   893  
   894  // rangeStart returns the leftmost position for the node range, assuming that
   895  // its children have correct range extents.
   896  func (n *node) rangeStart() Comparable {
   897  	minStart := n.items[0].Range().Start
   898  	if len(n.children) > 0 {
   899  		minStart = n.children[0].Range.Start
   900  	}
   901  	return minStart
   902  }
   903  
   904  // rangeEnd returns the rightmost position for the node range, assuming that its
   905  // children have correct range extents.
   906  func (n *node) rangeEnd() Comparable {
   907  	if len(n.items) == 0 {
   908  		maxEnd := n.children[0].Range.End
   909  		for _, c := range n.children[1:] {
   910  			if end := c.Range.End; maxEnd.Compare(end) < 0 {
   911  				maxEnd = end
   912  			}
   913  		}
   914  		return maxEnd
   915  	}
   916  	maxEnd := n.items[0].Range().End
   917  	for _, e := range n.items[1:] {
   918  		if end := e.Range().End; maxEnd.Compare(end) < 0 {
   919  			maxEnd = end
   920  		}
   921  	}
   922  	for _, c := range n.children {
   923  		if end := c.Range.End; maxEnd.Compare(end) < 0 {
   924  			maxEnd = end
   925  		}
   926  	}
   927  	return maxEnd
   928  }
   929  
   930  func (t *btree) AdjustRanges() {
   931  	if t.isEmpty() {
   932  		return
   933  	}
   934  	t.root.adjustRanges(t.root.cow)
   935  }
   936  
   937  func (n *node) adjustRanges(c *copyOnWriteContext) {
   938  	if n.cow != c {
   939  		// Could not have been modified.
   940  		return
   941  	}
   942  	for _, child := range n.children {
   943  		child.adjustRanges(c)
   944  	}
   945  	n.adjustRange()
   946  }
   947  
   948  // maxItems returns the max number of Interfaces to allow per node.
   949  func (t *btree) maxItems() int {
   950  	return t.minimumDegree*2 - 1
   951  }
   952  
   953  // minItems returns the min number of Interfaces to allow per node (ignored
   954  // for the root node).
   955  func (t *btree) minItems() int {
   956  	return t.minimumDegree - 1
   957  }
   958  
   959  func (c *copyOnWriteContext) newNode() (n *node) {
   960  	n = c.freelist.newNode()
   961  	n.cow = c
   962  	return
   963  }
   964  
   965  type freeType int
   966  
   967  const (
   968  	ftFreelistFull freeType = iota // node was freed (available for GC, not stored in freelist)
   969  	ftStored                       // node was stored in the freelist for later use
   970  	ftNotOwned                     // node was ignored by COW, since it's owned by another one
   971  )
   972  
   973  // freeNode frees a node within a given COW context, if it's owned by that
   974  // context.  It returns what happened to the node (see freeType const
   975  // documentation).
   976  func (c *copyOnWriteContext) freeNode(n *node) freeType {
   977  	if n.cow == c {
   978  		// clear to allow GC
   979  		n.items.truncate(0)
   980  		n.children.truncate(0)
   981  		n.cow = nil // clear to allow GC
   982  		if c.freelist.freeNode(n) {
   983  			return ftStored
   984  		}
   985  		return ftFreelistFull
   986  	}
   987  	return ftNotOwned
   988  }
   989  
   990  func (t *btree) Insert(e Interface, fast bool) (err error) {
   991  	// t.metrics("Insert")
   992  	if err = isValidInterface(e); err != nil {
   993  		return
   994  	}
   995  
   996  	if t.root == nil {
   997  		t.root = t.cow.newNode()
   998  		t.root.items = append(t.root.items, e)
   999  		t.length++
  1000  		if !fast {
  1001  			t.root.Range.Start = e.Range().Start
  1002  			t.root.Range.End = e.Range().End
  1003  		}
  1004  		return nil
  1005  	}
  1006  
  1007  	t.root = t.root.mutableFor(t.cow)
  1008  	if len(t.root.items) >= t.maxItems() {
  1009  		oldroot := t.root
  1010  		t.root = t.cow.newNode()
  1011  		if !fast {
  1012  			t.root.Range.Start = oldroot.Range.Start
  1013  			t.root.Range.End = oldroot.Range.End
  1014  		}
  1015  		e2, second := oldroot.split(t.maxItems()/2, fast)
  1016  		t.root.items = append(t.root.items, e2)
  1017  		t.root.children = append(t.root.children, oldroot, second)
  1018  	}
  1019  
  1020  	out, _ := t.root.insert(e, t.maxItems(), fast)
  1021  	if out == nil {
  1022  		t.length++
  1023  	}
  1024  	return
  1025  }
  1026  
  1027  func (t *btree) Delete(e Interface, fast bool) (err error) {
  1028  	// t.metrics("Delete")
  1029  	if err = isValidInterface(e); err != nil {
  1030  		return
  1031  	}
  1032  	if !t.overlappable(e.Range()) {
  1033  		return
  1034  	}
  1035  	t.delete(e, removeItem, fast)
  1036  	return
  1037  }
  1038  
  1039  func (t *btree) delete(e Interface, typ toRemove, fast bool) Interface {
  1040  	t.root = t.root.mutableFor(t.cow)
  1041  	out, _ := t.root.remove(e, t.minItems(), typ, fast)
  1042  	if len(t.root.items) == 0 && len(t.root.children) > 0 {
  1043  		oldroot := t.root
  1044  		t.root = t.root.children[0]
  1045  		t.cow.freeNode(oldroot)
  1046  	}
  1047  	if out != nil {
  1048  		t.length--
  1049  	}
  1050  	return out
  1051  }
  1052  
  1053  func (t *btree) Len() int {
  1054  	return t.length
  1055  }
  1056  
  1057  type stackElem struct {
  1058  	node  *node
  1059  	index int
  1060  }
  1061  
  1062  type btreeIterator struct {
  1063  	stack []*stackElem
  1064  }
  1065  
  1066  func (ti *btreeIterator) Next() (i Interface, ok bool) {
  1067  	if len(ti.stack) == 0 {
  1068  		return nil, false
  1069  	}
  1070  	elem := ti.stack[len(ti.stack)-1]
  1071  	curItem := elem.node.items[elem.index]
  1072  	elem.index++
  1073  	if elem.index >= len(elem.node.items) {
  1074  		ti.stack = ti.stack[:len(ti.stack)-1]
  1075  	}
  1076  	if len(elem.node.children) > 0 {
  1077  		for r := elem.node.children[elem.index]; r != nil; r = r.children[0] {
  1078  			ti.stack = append(ti.stack, &stackElem{r, 0})
  1079  			if len(r.children) == 0 {
  1080  				break
  1081  			}
  1082  		}
  1083  	}
  1084  	return curItem, true
  1085  }
  1086  
  1087  func (t *btree) Iterator() TreeIterator {
  1088  	var ti btreeIterator
  1089  	for n := t.root; n != nil; n = n.children[0] {
  1090  		ti.stack = append(ti.stack, &stackElem{n, 0})
  1091  		if len(n.children) == 0 {
  1092  			break
  1093  		}
  1094  	}
  1095  	return &ti
  1096  }
  1097  
  1098  // ClearWithOpt removes all items from the btree.  If addNodesToFreelist is
  1099  // true, t's nodes are added to its freelist as part of this call, until the
  1100  // freelist is full.  Otherwise, the root node is simply dereferenced and the
  1101  // subtree left to Go's normal GC processes.
  1102  //
  1103  // This can be much faster than calling Delete on all elements, because that
  1104  // requires finding/removing each element in the tree and updating the tree
  1105  // accordingly.  It also is somewhat faster than creating a new tree to replace
  1106  // the old one, because nodes from the old tree are reclaimed into the freelist
  1107  // for use by the new one, instead of being lost to the garbage collector.
  1108  //
  1109  // This call takes:
  1110  //   O(1): when addNodesToFreelist is false, this is a single operation.
  1111  //   O(1): when the freelist is already full, it breaks out immediately
  1112  //   O(freelist size):  when the freelist is empty and the nodes are all owned
  1113  //       by this tree, nodes are added to the freelist until full.
  1114  //   O(tree size):  when all nodes are owned by another tree, all nodes are
  1115  //       iterated over looking for nodes to add to the freelist, and due to
  1116  //       ownership, none are.
  1117  func (t *btree) ClearWithOpt(addNodesToFreelist bool) {
  1118  	if t.root != nil && addNodesToFreelist {
  1119  		t.root.reset(t.cow)
  1120  	}
  1121  	t.root, t.length = nil, 0
  1122  }
  1123  
  1124  func (t *btree) Clear() {
  1125  	t.ClearWithOpt(true)
  1126  }
  1127  
  1128  // reset returns a subtree to the freelist.  It breaks out immediately if the
  1129  // freelist is full, since the only benefit of iterating is to fill that
  1130  // freelist up.  Returns true if parent reset call should continue.
  1131  func (n *node) reset(c *copyOnWriteContext) bool {
  1132  	if n.cow != c {
  1133  		return false
  1134  	}
  1135  	for _, child := range n.children {
  1136  		if !child.reset(c) {
  1137  			return false
  1138  		}
  1139  	}
  1140  	return c.freeNode(n) != ftFreelistFull
  1141  }