gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/segment/set.go (about)

     1  // Copyright 2018 The gVisor 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 segment provides tools for working with collections of segments. A
    16  // segment is a key-value mapping, where the key is a non-empty contiguous
    17  // range of values of type Key, and the value is a single value of type Value.
    18  //
    19  // Clients using this package must use the go_template_instance rule in
    20  // tools/go_generics/defs.bzl to create an instantiation of this
    21  // template package, providing types to use in place of Key, Range, Value, and
    22  // Functions. See pkg/segment/test/BUILD for a usage example.
    23  package segment
    24  
    25  import (
    26  	"bytes"
    27  	"fmt"
    28  )
    29  
    30  // Key is a required type parameter that must be an integral type.
    31  type Key uint64
    32  
    33  // Range is a required type parameter equivalent to Range<Key>.
    34  type Range any
    35  
    36  // Value is a required type parameter.
    37  type Value any
    38  
    39  // trackGaps is an optional parameter.
    40  //
    41  // If trackGaps is 1, the Set will track maximum gap size recursively,
    42  // enabling the GapIterator.{Prev,Next}LargeEnoughGap functions. In this
    43  // case, Key must be an unsigned integer.
    44  //
    45  // trackGaps must be 0 or 1.
    46  const trackGaps = 0
    47  
    48  var _ = uint8(trackGaps << 7) // Will fail if not zero or one.
    49  
    50  // dynamicGap is a type that disappears if trackGaps is 0.
    51  type dynamicGap [trackGaps]Key
    52  
    53  // Get returns the value of the gap.
    54  //
    55  // Precondition: trackGaps must be non-zero.
    56  func (d *dynamicGap) Get() Key {
    57  	return d[:][0]
    58  }
    59  
    60  // Set sets the value of the gap.
    61  //
    62  // Precondition: trackGaps must be non-zero.
    63  func (d *dynamicGap) Set(v Key) {
    64  	d[:][0] = v
    65  }
    66  
    67  // Functions is a required type parameter that must be a struct implementing
    68  // the methods defined by Functions.
    69  type Functions interface {
    70  	// MinKey returns the minimum allowed key.
    71  	MinKey() Key
    72  
    73  	// MaxKey returns the maximum allowed key + 1.
    74  	MaxKey() Key
    75  
    76  	// ClearValue deinitializes the given value. (For example, if Value is a
    77  	// pointer or interface type, ClearValue should set it to nil.)
    78  	ClearValue(*Value)
    79  
    80  	// Merge attempts to merge the values corresponding to two consecutive
    81  	// segments. If successful, Merge returns (merged value, true). Otherwise,
    82  	// it returns (unspecified, false).
    83  	//
    84  	// Preconditions: r1.End == r2.Start.
    85  	//
    86  	// Postconditions: If merging succeeds, val1 and val2 are invalidated.
    87  	Merge(r1 Range, val1 Value, r2 Range, val2 Value) (Value, bool)
    88  
    89  	// Split splits a segment's value at a key within its range, such that the
    90  	// first returned value corresponds to the range [r.Start, split) and the
    91  	// second returned value corresponds to the range [split, r.End).
    92  	//
    93  	// Preconditions: r.Start < split < r.End.
    94  	//
    95  	// Postconditions: The original value val is invalidated.
    96  	Split(r Range, val Value, split Key) (Value, Value)
    97  }
    98  
    99  const (
   100  	// minDegree is the minimum degree of an internal node in a Set B-tree.
   101  	//
   102  	//	- Any non-root node has at least minDegree-1 segments.
   103  	//
   104  	//	- Any non-root internal (non-leaf) node has at least minDegree children.
   105  	//
   106  	//	- The root node may have fewer than minDegree-1 segments, but it may
   107  	// only have 0 segments if the tree is empty.
   108  	//
   109  	// Our implementation requires minDegree >= 3. Higher values of minDegree
   110  	// usually improve performance, but increase memory usage for small sets.
   111  	minDegree = 3
   112  
   113  	maxDegree = 2 * minDegree
   114  )
   115  
   116  // A Set is a mapping of segments with non-overlapping Range keys. The zero
   117  // value for a Set is an empty set. Set values are not safely movable nor
   118  // copyable. Set is thread-compatible.
   119  //
   120  // +stateify savable
   121  type Set struct {
   122  	root node `state:".([]FlatSegment)"`
   123  }
   124  
   125  // IsEmpty returns true if the set contains no segments.
   126  func (s *Set) IsEmpty() bool {
   127  	return s.root.nrSegments == 0
   128  }
   129  
   130  // IsEmptyRange returns true iff no segments in the set overlap the given
   131  // range. This is semantically equivalent to s.SpanRange(r) == 0, but may be
   132  // more efficient.
   133  func (s *Set) IsEmptyRange(r Range) bool {
   134  	switch {
   135  	case r.Length() < 0:
   136  		panic(fmt.Sprintf("invalid range %v", r))
   137  	case r.Length() == 0:
   138  		return true
   139  	}
   140  	_, gap := s.Find(r.Start)
   141  	if !gap.Ok() {
   142  		return false
   143  	}
   144  	return r.End <= gap.End()
   145  }
   146  
   147  // Span returns the total size of all segments in the set.
   148  func (s *Set) Span() Key {
   149  	var sz Key
   150  	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
   151  		sz += seg.Range().Length()
   152  	}
   153  	return sz
   154  }
   155  
   156  // SpanRange returns the total size of the intersection of segments in the set
   157  // with the given range.
   158  func (s *Set) SpanRange(r Range) Key {
   159  	switch {
   160  	case r.Length() < 0:
   161  		panic(fmt.Sprintf("invalid range %v", r))
   162  	case r.Length() == 0:
   163  		return 0
   164  	}
   165  	var sz Key
   166  	for seg := s.LowerBoundSegment(r.Start); seg.Ok() && seg.Start() < r.End; seg = seg.NextSegment() {
   167  		sz += seg.Range().Intersect(r).Length()
   168  	}
   169  	return sz
   170  }
   171  
   172  // FirstSegment returns the first segment in the set. If the set is empty,
   173  // FirstSegment returns a terminal iterator.
   174  func (s *Set) FirstSegment() Iterator {
   175  	if s.root.nrSegments == 0 {
   176  		return Iterator{}
   177  	}
   178  	return s.root.firstSegment()
   179  }
   180  
   181  // LastSegment returns the last segment in the set. If the set is empty,
   182  // LastSegment returns a terminal iterator.
   183  func (s *Set) LastSegment() Iterator {
   184  	if s.root.nrSegments == 0 {
   185  		return Iterator{}
   186  	}
   187  	return s.root.lastSegment()
   188  }
   189  
   190  // FirstGap returns the first gap in the set.
   191  func (s *Set) FirstGap() GapIterator {
   192  	n := &s.root
   193  	for n.hasChildren {
   194  		n = n.children[0]
   195  	}
   196  	return GapIterator{n, 0}
   197  }
   198  
   199  // LastGap returns the last gap in the set.
   200  func (s *Set) LastGap() GapIterator {
   201  	n := &s.root
   202  	for n.hasChildren {
   203  		n = n.children[n.nrSegments]
   204  	}
   205  	return GapIterator{n, n.nrSegments}
   206  }
   207  
   208  // Find returns the segment or gap whose range contains the given key. If a
   209  // segment is found, the returned Iterator is non-terminal and the
   210  // returned GapIterator is terminal. Otherwise, the returned Iterator is
   211  // terminal and the returned GapIterator is non-terminal.
   212  func (s *Set) Find(key Key) (Iterator, GapIterator) {
   213  	n := &s.root
   214  	for {
   215  		// Binary search invariant: the correct value of i lies within [lower,
   216  		// upper].
   217  		lower := 0
   218  		upper := n.nrSegments
   219  		for lower < upper {
   220  			i := lower + (upper-lower)/2
   221  			if r := n.keys[i]; key < r.End {
   222  				if key >= r.Start {
   223  					return Iterator{n, i}, GapIterator{}
   224  				}
   225  				upper = i
   226  			} else {
   227  				lower = i + 1
   228  			}
   229  		}
   230  		i := lower
   231  		if !n.hasChildren {
   232  			return Iterator{}, GapIterator{n, i}
   233  		}
   234  		n = n.children[i]
   235  	}
   236  }
   237  
   238  // FindSegment returns the segment whose range contains the given key. If no
   239  // such segment exists, FindSegment returns a terminal iterator.
   240  func (s *Set) FindSegment(key Key) Iterator {
   241  	seg, _ := s.Find(key)
   242  	return seg
   243  }
   244  
   245  // LowerBoundSegment returns the segment with the lowest range that contains a
   246  // key greater than or equal to min. If no such segment exists,
   247  // LowerBoundSegment returns a terminal iterator.
   248  func (s *Set) LowerBoundSegment(min Key) Iterator {
   249  	seg, gap := s.Find(min)
   250  	if seg.Ok() {
   251  		return seg
   252  	}
   253  	return gap.NextSegment()
   254  }
   255  
   256  // UpperBoundSegment returns the segment with the highest range that contains a
   257  // key less than or equal to max. If no such segment exists, UpperBoundSegment
   258  // returns a terminal iterator.
   259  func (s *Set) UpperBoundSegment(max Key) Iterator {
   260  	seg, gap := s.Find(max)
   261  	if seg.Ok() {
   262  		return seg
   263  	}
   264  	return gap.PrevSegment()
   265  }
   266  
   267  // FindGap returns the gap containing the given key. If no such gap exists
   268  // (i.e. the set contains a segment containing that key), FindGap returns a
   269  // terminal iterator.
   270  func (s *Set) FindGap(key Key) GapIterator {
   271  	_, gap := s.Find(key)
   272  	return gap
   273  }
   274  
   275  // LowerBoundGap returns the gap with the lowest range that is greater than or
   276  // equal to min.
   277  func (s *Set) LowerBoundGap(min Key) GapIterator {
   278  	seg, gap := s.Find(min)
   279  	if gap.Ok() {
   280  		return gap
   281  	}
   282  	return seg.NextGap()
   283  }
   284  
   285  // UpperBoundGap returns the gap with the highest range that is less than or
   286  // equal to max.
   287  func (s *Set) UpperBoundGap(max Key) GapIterator {
   288  	seg, gap := s.Find(max)
   289  	if gap.Ok() {
   290  		return gap
   291  	}
   292  	return seg.PrevGap()
   293  }
   294  
   295  // FirstLargeEnoughGap returns the first gap in the set with at least the given
   296  // length. If no such gap exists, FirstLargeEnoughGap returns a terminal
   297  // iterator.
   298  //
   299  // Precondition: trackGaps must be 1.
   300  func (s *Set) FirstLargeEnoughGap(minSize Key) GapIterator {
   301  	if trackGaps != 1 {
   302  		panic("set is not tracking gaps")
   303  	}
   304  	gap := s.FirstGap()
   305  	if gap.Range().Length() >= minSize {
   306  		return gap
   307  	}
   308  	return gap.NextLargeEnoughGap(minSize)
   309  }
   310  
   311  // LastLargeEnoughGap returns the last gap in the set with at least the given
   312  // length. If no such gap exists, LastLargeEnoughGap returns a terminal
   313  // iterator.
   314  //
   315  // Precondition: trackGaps must be 1.
   316  func (s *Set) LastLargeEnoughGap(minSize Key) GapIterator {
   317  	if trackGaps != 1 {
   318  		panic("set is not tracking gaps")
   319  	}
   320  	gap := s.LastGap()
   321  	if gap.Range().Length() >= minSize {
   322  		return gap
   323  	}
   324  	return gap.PrevLargeEnoughGap(minSize)
   325  }
   326  
   327  // LowerBoundLargeEnoughGap returns the first gap in the set with at least the
   328  // given length and whose range contains a key greater than or equal to min. If
   329  // no such gap exists, LowerBoundLargeEnoughGap returns a terminal iterator.
   330  //
   331  // Precondition: trackGaps must be 1.
   332  func (s *Set) LowerBoundLargeEnoughGap(min, minSize Key) GapIterator {
   333  	if trackGaps != 1 {
   334  		panic("set is not tracking gaps")
   335  	}
   336  	gap := s.LowerBoundGap(min)
   337  	if gap.Range().Length() >= minSize {
   338  		return gap
   339  	}
   340  	return gap.NextLargeEnoughGap(minSize)
   341  }
   342  
   343  // UpperBoundLargeEnoughGap returns the last gap in the set with at least the
   344  // given length and whose range contains a key less than or equal to max. If no
   345  // such gap exists, UpperBoundLargeEnoughGap returns a terminal iterator.
   346  //
   347  // Precondition: trackGaps must be 1.
   348  func (s *Set) UpperBoundLargeEnoughGap(max, minSize Key) GapIterator {
   349  	if trackGaps != 1 {
   350  		panic("set is not tracking gaps")
   351  	}
   352  	gap := s.UpperBoundGap(max)
   353  	if gap.Range().Length() >= minSize {
   354  		return gap
   355  	}
   356  	return gap.PrevLargeEnoughGap(minSize)
   357  }
   358  
   359  // Insert inserts the given segment into the given gap. If the new segment can
   360  // be merged with adjacent segments, Insert will do so. Insert returns an
   361  // iterator to the segment containing the inserted value (which may have been
   362  // merged with other values). All existing iterators (including gap, but not
   363  // including the returned iterator) are invalidated.
   364  //
   365  // If the gap cannot accommodate the segment, or if r is invalid, Insert panics.
   366  //
   367  // Insert is semantically equivalent to a InsertWithoutMerging followed by a
   368  // Merge, but may be more efficient. Note that there is no unchecked variant of
   369  // Insert since Insert must retrieve and inspect gap's predecessor and
   370  // successor segments regardless.
   371  func (s *Set) Insert(gap GapIterator, r Range, val Value) Iterator {
   372  	if r.Length() <= 0 {
   373  		panic(fmt.Sprintf("invalid segment range %v", r))
   374  	}
   375  	prev, next := gap.PrevSegment(), gap.NextSegment()
   376  	if prev.Ok() && prev.End() > r.Start {
   377  		panic(fmt.Sprintf("new segment %v overlaps predecessor %v", r, prev.Range()))
   378  	}
   379  	if next.Ok() && next.Start() < r.End {
   380  		panic(fmt.Sprintf("new segment %v overlaps successor %v", r, next.Range()))
   381  	}
   382  	if prev.Ok() && prev.End() == r.Start {
   383  		if mval, ok := (Functions{}).Merge(prev.Range(), prev.Value(), r, val); ok {
   384  			shrinkMaxGap := trackGaps != 0 && gap.Range().Length() == gap.node.maxGap.Get()
   385  			prev.SetEndUnchecked(r.End)
   386  			prev.SetValue(mval)
   387  			if shrinkMaxGap {
   388  				gap.node.updateMaxGapLeaf()
   389  			}
   390  			if next.Ok() && next.Start() == r.End {
   391  				val = mval
   392  				if mval, ok := (Functions{}).Merge(prev.Range(), val, next.Range(), next.Value()); ok {
   393  					prev.SetEndUnchecked(next.End())
   394  					prev.SetValue(mval)
   395  					return s.Remove(next).PrevSegment()
   396  				}
   397  			}
   398  			return prev
   399  		}
   400  	}
   401  	if next.Ok() && next.Start() == r.End {
   402  		if mval, ok := (Functions{}).Merge(r, val, next.Range(), next.Value()); ok {
   403  			shrinkMaxGap := trackGaps != 0 && gap.Range().Length() == gap.node.maxGap.Get()
   404  			next.SetStartUnchecked(r.Start)
   405  			next.SetValue(mval)
   406  			if shrinkMaxGap {
   407  				gap.node.updateMaxGapLeaf()
   408  			}
   409  			return next
   410  		}
   411  	}
   412  	// InsertWithoutMergingUnchecked will maintain maxGap if necessary.
   413  	return s.InsertWithoutMergingUnchecked(gap, r, val)
   414  }
   415  
   416  // InsertWithoutMerging inserts the given segment into the given gap and
   417  // returns an iterator to the inserted segment. All existing iterators
   418  // (including gap, but not including the returned iterator) are invalidated.
   419  //
   420  // If the gap cannot accommodate the segment, or if r is invalid,
   421  // InsertWithoutMerging panics.
   422  func (s *Set) InsertWithoutMerging(gap GapIterator, r Range, val Value) Iterator {
   423  	if r.Length() <= 0 {
   424  		panic(fmt.Sprintf("invalid segment range %v", r))
   425  	}
   426  	if gr := gap.Range(); !gr.IsSupersetOf(r) {
   427  		panic(fmt.Sprintf("cannot insert segment range %v into gap range %v", r, gr))
   428  	}
   429  	return s.InsertWithoutMergingUnchecked(gap, r, val)
   430  }
   431  
   432  // InsertWithoutMergingUnchecked inserts the given segment into the given gap
   433  // and returns an iterator to the inserted segment. All existing iterators
   434  // (including gap, but not including the returned iterator) are invalidated.
   435  //
   436  // Preconditions:
   437  //   - r.Start >= gap.Start().
   438  //   - r.End <= gap.End().
   439  func (s *Set) InsertWithoutMergingUnchecked(gap GapIterator, r Range, val Value) Iterator {
   440  	gap = gap.node.rebalanceBeforeInsert(gap)
   441  	splitMaxGap := trackGaps != 0 && (gap.node.nrSegments == 0 || gap.Range().Length() == gap.node.maxGap.Get())
   442  	copy(gap.node.keys[gap.index+1:], gap.node.keys[gap.index:gap.node.nrSegments])
   443  	copy(gap.node.values[gap.index+1:], gap.node.values[gap.index:gap.node.nrSegments])
   444  	gap.node.keys[gap.index] = r
   445  	gap.node.values[gap.index] = val
   446  	gap.node.nrSegments++
   447  	if splitMaxGap {
   448  		gap.node.updateMaxGapLeaf()
   449  	}
   450  	return Iterator{gap.node, gap.index}
   451  }
   452  
   453  // InsertRange inserts the given segment into the set. If the new segment can
   454  // be merged with adjacent segments, InsertRange will do so. InsertRange
   455  // returns an iterator to the segment containing the inserted value (which may
   456  // have been merged with other values). All existing iterators (excluding the
   457  // returned iterator) are invalidated.
   458  //
   459  // If the new segment would overlap an existing segment, or if r is invalid,
   460  // InsertRange panics.
   461  //
   462  // InsertRange searches the set to find the gap to insert into. If the caller
   463  // already has the appropriate GapIterator, or if the caller needs to do
   464  // additional work between finding the gap and insertion, use Insert instead.
   465  func (s *Set) InsertRange(r Range, val Value) Iterator {
   466  	if r.Length() <= 0 {
   467  		panic(fmt.Sprintf("invalid segment range %v", r))
   468  	}
   469  	seg, gap := s.Find(r.Start)
   470  	if seg.Ok() {
   471  		panic(fmt.Sprintf("new segment %v overlaps existing segment %v", r, seg.Range()))
   472  	}
   473  	if gap.End() < r.End {
   474  		panic(fmt.Sprintf("new segment %v overlaps existing segment %v", r, gap.NextSegment().Range()))
   475  	}
   476  	return s.Insert(gap, r, val)
   477  }
   478  
   479  // InsertWithoutMergingRange inserts the given segment into the set and returns
   480  // an iterator to the inserted segment. All existing iterators (excluding the
   481  // returned iterator) are invalidated.
   482  //
   483  // If the new segment would overlap an existing segment, or if r is invalid,
   484  // InsertWithoutMergingRange panics.
   485  //
   486  // InsertWithoutMergingRange searches the set to find the gap to insert into.
   487  // If the caller already has the appropriate GapIterator, or if the caller
   488  // needs to do additional work between finding the gap and insertion, use
   489  // InsertWithoutMerging instead.
   490  func (s *Set) InsertWithoutMergingRange(r Range, val Value) Iterator {
   491  	if r.Length() <= 0 {
   492  		panic(fmt.Sprintf("invalid segment range %v", r))
   493  	}
   494  	seg, gap := s.Find(r.Start)
   495  	if seg.Ok() {
   496  		panic(fmt.Sprintf("new segment %v overlaps existing segment %v", r, seg.Range()))
   497  	}
   498  	if gap.End() < r.End {
   499  		panic(fmt.Sprintf("new segment %v overlaps existing segment %v", r, gap.NextSegment().Range()))
   500  	}
   501  	return s.InsertWithoutMerging(gap, r, val)
   502  }
   503  
   504  // TryInsertRange attempts to insert the given segment into the set. If the new
   505  // segment can be merged with adjacent segments, TryInsertRange will do so.
   506  // TryInsertRange returns an iterator to the segment containing the inserted
   507  // value (which may have been merged with other values). All existing iterators
   508  // (excluding the returned iterator) are invalidated.
   509  //
   510  // If the new segment would overlap an existing segment, TryInsertRange does
   511  // nothing and returns a terminal iterator.
   512  //
   513  // TryInsertRange searches the set to find the gap to insert into. If the
   514  // caller already has the appropriate GapIterator, or if the caller needs to do
   515  // additional work between finding the gap and insertion, use Insert instead.
   516  func (s *Set) TryInsertRange(r Range, val Value) Iterator {
   517  	if r.Length() <= 0 {
   518  		panic(fmt.Sprintf("invalid segment range %v", r))
   519  	}
   520  	seg, gap := s.Find(r.Start)
   521  	if seg.Ok() {
   522  		return Iterator{}
   523  	}
   524  	if gap.End() < r.End {
   525  		return Iterator{}
   526  	}
   527  	return s.Insert(gap, r, val)
   528  }
   529  
   530  // TryInsertWithoutMergingRange attempts to insert the given segment into the
   531  // set. If successful, it returns an iterator to the inserted segment; all
   532  // existing iterators (excluding the returned iterator) are invalidated. If the
   533  // new segment would overlap an existing segment, TryInsertWithoutMergingRange
   534  // does nothing and returns a terminal iterator.
   535  //
   536  // TryInsertWithoutMergingRange searches the set to find the gap to insert
   537  // into. If the caller already has the appropriate GapIterator, or if the
   538  // caller needs to do additional work between finding the gap and insertion,
   539  // use InsertWithoutMerging instead.
   540  func (s *Set) TryInsertWithoutMergingRange(r Range, val Value) Iterator {
   541  	if r.Length() <= 0 {
   542  		panic(fmt.Sprintf("invalid segment range %v", r))
   543  	}
   544  	seg, gap := s.Find(r.Start)
   545  	if seg.Ok() {
   546  		return Iterator{}
   547  	}
   548  	if gap.End() < r.End {
   549  		return Iterator{}
   550  	}
   551  	return s.InsertWithoutMerging(gap, r, val)
   552  }
   553  
   554  // Remove removes the given segment and returns an iterator to the vacated gap.
   555  // All existing iterators (including seg, but not including the returned
   556  // iterator) are invalidated.
   557  func (s *Set) Remove(seg Iterator) GapIterator {
   558  	// We only want to remove directly from a leaf node.
   559  	if seg.node.hasChildren {
   560  		// Since seg.node has children, the removed segment must have a
   561  		// predecessor (at the end of the rightmost leaf of its left child
   562  		// subtree). Move the contents of that predecessor into the removed
   563  		// segment's position, and remove that predecessor instead. (We choose
   564  		// to steal the predecessor rather than the successor because removing
   565  		// from the end of a leaf node doesn't involve any copying unless
   566  		// merging is required.)
   567  		victim := seg.PrevSegment()
   568  		// This must be unchecked since until victim is removed, seg and victim
   569  		// overlap.
   570  		seg.SetRangeUnchecked(victim.Range())
   571  		seg.SetValue(victim.Value())
   572  		// Need to update the nextAdjacentNode's maxGap because the gap in between
   573  		// must have been modified by updating seg.Range() to victim.Range().
   574  		// seg.NextSegment() must exist since the last segment can't be in a
   575  		// non-leaf node.
   576  		nextAdjacentNode := seg.NextSegment().node
   577  		if trackGaps != 0 {
   578  			nextAdjacentNode.updateMaxGapLeaf()
   579  		}
   580  		return s.Remove(victim).NextGap()
   581  	}
   582  	copy(seg.node.keys[seg.index:], seg.node.keys[seg.index+1:seg.node.nrSegments])
   583  	copy(seg.node.values[seg.index:], seg.node.values[seg.index+1:seg.node.nrSegments])
   584  	Functions{}.ClearValue(&seg.node.values[seg.node.nrSegments-1])
   585  	seg.node.nrSegments--
   586  	if trackGaps != 0 {
   587  		seg.node.updateMaxGapLeaf()
   588  	}
   589  	return seg.node.rebalanceAfterRemove(GapIterator{seg.node, seg.index})
   590  }
   591  
   592  // RemoveAll removes all segments from the set. All existing iterators are
   593  // invalidated.
   594  func (s *Set) RemoveAll() {
   595  	s.root = node{}
   596  }
   597  
   598  // RemoveRange removes all segments in the given range. An iterator to the
   599  // newly formed gap is returned, and all existing iterators are invalidated.
   600  //
   601  // RemoveRange searches the set to find segments to remove. If the caller
   602  // already has an iterator to either end of the range of segments to remove, or
   603  // if the caller needs to do additional work before removing each segment,
   604  // iterate segments and call Remove in a loop instead.
   605  func (s *Set) RemoveRange(r Range) GapIterator {
   606  	seg, gap := s.Find(r.Start)
   607  	if seg.Ok() {
   608  		seg = s.Isolate(seg, r)
   609  		gap = s.Remove(seg)
   610  	}
   611  	for seg = gap.NextSegment(); seg.Ok() && seg.Start() < r.End; seg = gap.NextSegment() {
   612  		seg = s.SplitAfter(seg, r.End)
   613  		gap = s.Remove(seg)
   614  	}
   615  	return gap
   616  }
   617  
   618  // RemoveFullRange is equivalent to RemoveRange, except that if any key in the
   619  // given range does not correspond to a segment, RemoveFullRange panics.
   620  func (s *Set) RemoveFullRange(r Range) GapIterator {
   621  	seg := s.FindSegment(r.Start)
   622  	if !seg.Ok() {
   623  		panic(fmt.Sprintf("missing segment at %v", r.Start))
   624  	}
   625  	seg = s.SplitBefore(seg, r.Start)
   626  	for {
   627  		seg = s.SplitAfter(seg, r.End)
   628  		end := seg.End()
   629  		gap := s.Remove(seg)
   630  		if r.End <= end {
   631  			return gap
   632  		}
   633  		seg = gap.NextSegment()
   634  		if !seg.Ok() || seg.Start() != end {
   635  			panic(fmt.Sprintf("missing segment at %v", end))
   636  		}
   637  	}
   638  }
   639  
   640  // Merge attempts to merge two neighboring segments. If successful, Merge
   641  // returns an iterator to the merged segment, and all existing iterators are
   642  // invalidated. Otherwise, Merge returns a terminal iterator.
   643  //
   644  // If first is not the predecessor of second, Merge panics.
   645  func (s *Set) Merge(first, second Iterator) Iterator {
   646  	if first.NextSegment() != second {
   647  		panic(fmt.Sprintf("attempt to merge non-neighboring segments %v, %v", first.Range(), second.Range()))
   648  	}
   649  	return s.MergeUnchecked(first, second)
   650  }
   651  
   652  // MergeUnchecked attempts to merge two neighboring segments. If successful,
   653  // MergeUnchecked returns an iterator to the merged segment, and all existing
   654  // iterators are invalidated. Otherwise, MergeUnchecked returns a terminal
   655  // iterator.
   656  //
   657  // Precondition: first is the predecessor of second: first.NextSegment() ==
   658  // second, first == second.PrevSegment().
   659  func (s *Set) MergeUnchecked(first, second Iterator) Iterator {
   660  	if first.End() == second.Start() {
   661  		if mval, ok := (Functions{}).Merge(first.Range(), first.Value(), second.Range(), second.Value()); ok {
   662  			// N.B. This must be unchecked because until s.Remove(second), first
   663  			// overlaps second.
   664  			first.SetEndUnchecked(second.End())
   665  			first.SetValue(mval)
   666  			// Remove will handle the maxGap update if necessary.
   667  			return s.Remove(second).PrevSegment()
   668  		}
   669  	}
   670  	return Iterator{}
   671  }
   672  
   673  // MergePrev attempts to merge the given segment with its predecessor if
   674  // possible, and returns an updated iterator to the extended segment. All
   675  // existing iterators (including seg, but not including the returned iterator)
   676  // are invalidated.
   677  //
   678  // MergePrev is usually used when mutating segments while iterating them in
   679  // order of increasing keys, to attempt merging of each mutated segment with
   680  // its previously-mutated predecessor. In such cases, merging a mutated segment
   681  // with its unmutated successor would incorrectly cause the latter to be
   682  // skipped.
   683  func (s *Set) MergePrev(seg Iterator) Iterator {
   684  	if prev := seg.PrevSegment(); prev.Ok() {
   685  		if mseg := s.MergeUnchecked(prev, seg); mseg.Ok() {
   686  			seg = mseg
   687  		}
   688  	}
   689  	return seg
   690  }
   691  
   692  // MergeNext attempts to merge the given segment with its successor if
   693  // possible, and returns an updated iterator to the extended segment. All
   694  // existing iterators (including seg, but not including the returned iterator)
   695  // are invalidated.
   696  //
   697  // MergeNext is usually used when mutating segments while iterating them in
   698  // order of decreasing keys, to attempt merging of each mutated segment with
   699  // its previously-mutated successor. In such cases, merging a mutated segment
   700  // with its unmutated predecessor would incorrectly cause the latter to be
   701  // skipped.
   702  func (s *Set) MergeNext(seg Iterator) Iterator {
   703  	if next := seg.NextSegment(); next.Ok() {
   704  		if mseg := s.MergeUnchecked(seg, next); mseg.Ok() {
   705  			seg = mseg
   706  		}
   707  	}
   708  	return seg
   709  }
   710  
   711  // Unisolate attempts to merge the given segment with its predecessor and
   712  // successor if possible, and returns an updated iterator to the extended
   713  // segment. All existing iterators (including seg, but not including the
   714  // returned iterator) are invalidated.
   715  //
   716  // Unisolate is usually used in conjunction with Isolate when mutating part of
   717  // a single segment in a way that may affect its mergeability. For the reasons
   718  // described by MergePrev and MergeNext, it is usually incorrect to use the
   719  // return value of Unisolate in a loop variable.
   720  func (s *Set) Unisolate(seg Iterator) Iterator {
   721  	if prev := seg.PrevSegment(); prev.Ok() {
   722  		if mseg := s.MergeUnchecked(prev, seg); mseg.Ok() {
   723  			seg = mseg
   724  		}
   725  	}
   726  	if next := seg.NextSegment(); next.Ok() {
   727  		if mseg := s.MergeUnchecked(seg, next); mseg.Ok() {
   728  			seg = mseg
   729  		}
   730  	}
   731  	return seg
   732  }
   733  
   734  // MergeAll merges all mergeable adjacent segments in the set. All existing
   735  // iterators are invalidated.
   736  func (s *Set) MergeAll() {
   737  	seg := s.FirstSegment()
   738  	if !seg.Ok() {
   739  		return
   740  	}
   741  	next := seg.NextSegment()
   742  	for next.Ok() {
   743  		if mseg := s.MergeUnchecked(seg, next); mseg.Ok() {
   744  			seg, next = mseg, mseg.NextSegment()
   745  		} else {
   746  			seg, next = next, next.NextSegment()
   747  		}
   748  	}
   749  }
   750  
   751  // MergeInsideRange attempts to merge all adjacent segments that contain a key
   752  // in the specific range. All existing iterators are invalidated.
   753  //
   754  // MergeInsideRange only makes sense after mutating the set in a way that may
   755  // change the mergeability of modified segments; callers should prefer to use
   756  // MergePrev or MergeNext during the mutating loop instead (depending on the
   757  // direction of iteration), in order to avoid a redundant search.
   758  func (s *Set) MergeInsideRange(r Range) {
   759  	seg := s.LowerBoundSegment(r.Start)
   760  	if !seg.Ok() {
   761  		return
   762  	}
   763  	next := seg.NextSegment()
   764  	for next.Ok() && next.Start() < r.End {
   765  		if mseg := s.MergeUnchecked(seg, next); mseg.Ok() {
   766  			seg, next = mseg, mseg.NextSegment()
   767  		} else {
   768  			seg, next = next, next.NextSegment()
   769  		}
   770  	}
   771  }
   772  
   773  // MergeOutsideRange attempts to merge the segment containing r.Start with its
   774  // predecessor, and the segment containing r.End-1 with its successor.
   775  //
   776  // MergeOutsideRange only makes sense after mutating the set in a way that may
   777  // change the mergeability of modified segments; callers should prefer to use
   778  // MergePrev or MergeNext during the mutating loop instead (depending on the
   779  // direction of iteration), in order to avoid two redundant searches.
   780  func (s *Set) MergeOutsideRange(r Range) {
   781  	first := s.FindSegment(r.Start)
   782  	if first.Ok() {
   783  		if prev := first.PrevSegment(); prev.Ok() {
   784  			s.Merge(prev, first)
   785  		}
   786  	}
   787  	last := s.FindSegment(r.End - 1)
   788  	if last.Ok() {
   789  		if next := last.NextSegment(); next.Ok() {
   790  			s.Merge(last, next)
   791  		}
   792  	}
   793  }
   794  
   795  // Split splits the given segment at the given key and returns iterators to the
   796  // two resulting segments. All existing iterators (including seg, but not
   797  // including the returned iterators) are invalidated.
   798  //
   799  // If the segment cannot be split at split (because split is at the start or
   800  // end of the segment's range, so splitting would produce a segment with zero
   801  // length, or because split falls outside the segment's range altogether),
   802  // Split panics.
   803  func (s *Set) Split(seg Iterator, split Key) (Iterator, Iterator) {
   804  	if !seg.Range().CanSplitAt(split) {
   805  		panic(fmt.Sprintf("can't split %v at %v", seg.Range(), split))
   806  	}
   807  	return s.SplitUnchecked(seg, split)
   808  }
   809  
   810  // SplitUnchecked splits the given segment at the given key and returns
   811  // iterators to the two resulting segments. All existing iterators (including
   812  // seg, but not including the returned iterators) are invalidated.
   813  //
   814  // Preconditions: seg.Start() < key < seg.End().
   815  func (s *Set) SplitUnchecked(seg Iterator, split Key) (Iterator, Iterator) {
   816  	val1, val2 := (Functions{}).Split(seg.Range(), seg.Value(), split)
   817  	end2 := seg.End()
   818  	seg.SetEndUnchecked(split)
   819  	seg.SetValue(val1)
   820  	seg2 := s.InsertWithoutMergingUnchecked(seg.NextGap(), Range{split, end2}, val2)
   821  	// seg may now be invalid due to the Insert.
   822  	return seg2.PrevSegment(), seg2
   823  }
   824  
   825  // SplitBefore ensures that the given segment's start is at least start by
   826  // splitting at start if necessary, and returns an updated iterator to the
   827  // bounded segment. All existing iterators (including seg, but not including
   828  // the returned iterator) are invalidated.
   829  //
   830  // SplitBefore is usually when mutating segments in a range. In such cases,
   831  // when iterating segments in order of increasing keys, the first segment may
   832  // extend beyond the start of the range to be mutated, and needs to be
   833  // SplitBefore to ensure that only the part of the segment within the range is
   834  // mutated. When iterating segments in order of decreasing keys, SplitBefore
   835  // and SplitAfter; i.e. SplitBefore needs to be invoked on each segment, while
   836  // SplitAfter only needs to be invoked on the first.
   837  //
   838  // Preconditions: start < seg.End().
   839  func (s *Set) SplitBefore(seg Iterator, start Key) Iterator {
   840  	if seg.Range().CanSplitAt(start) {
   841  		_, seg = s.SplitUnchecked(seg, start)
   842  	}
   843  	return seg
   844  }
   845  
   846  // SplitAfter ensures that the given segment's end is at most end by splitting
   847  // at end if necessary, and returns an updated iterator to the bounded segment.
   848  // All existing iterators (including seg, but not including the returned
   849  // iterator) are invalidated.
   850  //
   851  // SplitAfter is usually used when mutating segments in a range. In such cases,
   852  // when iterating segments in order of increasing keys, each iterated segment
   853  // may extend beyond the end of the range to be mutated, and needs to be
   854  // SplitAfter to ensure that only the part of the segment within the range is
   855  // mutated. When iterating segments in order of decreasing keys, SplitBefore
   856  // and SplitAfter exchange roles; i.e. SplitBefore needs to be invoked on each
   857  // segment, while SplitAfter only needs to be invoked on the first.
   858  //
   859  // Preconditions: seg.Start() < end.
   860  func (s *Set) SplitAfter(seg Iterator, end Key) Iterator {
   861  	if seg.Range().CanSplitAt(end) {
   862  		seg, _ = s.SplitUnchecked(seg, end)
   863  	}
   864  	return seg
   865  }
   866  
   867  // Isolate ensures that the given segment's range is a subset of r by splitting
   868  // at r.Start and r.End if necessary, and returns an updated iterator to the
   869  // bounded segment. All existing iterators (including seg, but not including
   870  // the returned iterators) are invalidated.
   871  //
   872  // Isolate is usually used when mutating part of a single segment, or when
   873  // mutating segments in a range where the first segment is not necessarily
   874  // split, making use of SplitBefore/SplitAfter complex.
   875  //
   876  // Preconditions: seg.Range().Overlaps(r).
   877  func (s *Set) Isolate(seg Iterator, r Range) Iterator {
   878  	if seg.Range().CanSplitAt(r.Start) {
   879  		_, seg = s.SplitUnchecked(seg, r.Start)
   880  	}
   881  	if seg.Range().CanSplitAt(r.End) {
   882  		seg, _ = s.SplitUnchecked(seg, r.End)
   883  	}
   884  	return seg
   885  }
   886  
   887  // LowerBoundSegmentSplitBefore combines LowerBoundSegment and SplitBefore.
   888  //
   889  // LowerBoundSegmentSplitBefore is usually used when mutating segments in a
   890  // range while iterating them in order of increasing keys. In such cases,
   891  // LowerBoundSegmentSplitBefore provides an iterator to the first segment to be
   892  // mutated, suitable as the initial value for a loop variable.
   893  func (s *Set) LowerBoundSegmentSplitBefore(min Key) Iterator {
   894  	seg := s.LowerBoundSegment(min)
   895  	if seg.Ok() {
   896  		seg = s.SplitBefore(seg, min)
   897  	}
   898  	return seg
   899  }
   900  
   901  // UpperBoundSegmentSplitAfter combines UpperBoundSegment and SplitAfter.
   902  //
   903  // UpperBoundSegmentSplitAfter is usually used when mutating segments in a
   904  // range while iterating them in order of decreasing keys. In such cases,
   905  // UpperBoundSegmentSplitAfter provides an iterator to the first segment to be
   906  // mutated, suitable as the initial value for a loop variable.
   907  func (s *Set) UpperBoundSegmentSplitAfter(max Key) Iterator {
   908  	seg := s.UpperBoundSegment(max)
   909  	if seg.Ok() {
   910  		seg = s.SplitAfter(seg, max)
   911  	}
   912  	return seg
   913  }
   914  
   915  // VisitRange applies the function f to all segments intersecting the range r,
   916  // in order of ascending keys. Segments will not be split, so f may be called
   917  // on segments lying partially outside r. Non-empty gaps between segments are
   918  // skipped. If a call to f returns false, VisitRange stops iteration
   919  // immediately.
   920  //
   921  // N.B. f must not invalidate iterators into s.
   922  func (s *Set) VisitRange(r Range, f func(seg Iterator) bool) {
   923  	for seg := s.LowerBoundSegment(r.Start); seg.Ok() && seg.Start() < r.End; seg = seg.NextSegment() {
   924  		if !f(seg) {
   925  			return
   926  		}
   927  	}
   928  }
   929  
   930  // VisitFullRange is equivalent to VisitRange, except that if any key in r that
   931  // is visited before f returns false does not correspond to a segment,
   932  // VisitFullRange panics.
   933  func (s *Set) VisitFullRange(r Range, f func(seg Iterator) bool) {
   934  	pos := r.Start
   935  	seg := s.FindSegment(r.Start)
   936  	for {
   937  		if !seg.Ok() {
   938  			panic(fmt.Sprintf("missing segment at %v", pos))
   939  		}
   940  		if !f(seg) {
   941  			return
   942  		}
   943  		pos = seg.End()
   944  		if r.End <= pos {
   945  			return
   946  		}
   947  		seg, _ = seg.NextNonEmpty()
   948  	}
   949  }
   950  
   951  // MutateRange applies the function f to all segments intersecting the range r,
   952  // in order of ascending keys. Segments that lie partially outside r are split
   953  // before f is called, such that f only observes segments entirely within r.
   954  // Iterated segments are merged again after f is called. Non-empty gaps between
   955  // segments are skipped. If a call to f returns false, MutateRange stops
   956  // iteration immediately.
   957  //
   958  // MutateRange invalidates all existing iterators.
   959  //
   960  // N.B. f must not invalidate iterators into s.
   961  func (s *Set) MutateRange(r Range, f func(seg Iterator) bool) {
   962  	seg := s.LowerBoundSegmentSplitBefore(r.Start)
   963  	for seg.Ok() && seg.Start() < r.End {
   964  		seg = s.SplitAfter(seg, r.End)
   965  		cont := f(seg)
   966  		seg = s.MergePrev(seg)
   967  		if !cont {
   968  			s.MergeNext(seg)
   969  			return
   970  		}
   971  		seg = seg.NextSegment()
   972  	}
   973  	if seg.Ok() {
   974  		s.MergePrev(seg)
   975  	}
   976  }
   977  
   978  // MutateFullRange is equivalent to MutateRange, except that if any key in r
   979  // that is visited before f returns false does not correspond to a segment,
   980  // MutateFullRange panics.
   981  func (s *Set) MutateFullRange(r Range, f func(seg Iterator) bool) {
   982  	seg := s.FindSegment(r.Start)
   983  	if !seg.Ok() {
   984  		panic(fmt.Sprintf("missing segment at %v", r.Start))
   985  	}
   986  	seg = s.SplitBefore(seg, r.Start)
   987  	for {
   988  		seg = s.SplitAfter(seg, r.End)
   989  		cont := f(seg)
   990  		end := seg.End()
   991  		seg = s.MergePrev(seg)
   992  		if !cont || r.End <= end {
   993  			s.MergeNext(seg)
   994  			return
   995  		}
   996  		seg = seg.NextSegment()
   997  		if !seg.Ok() || seg.Start() != end {
   998  			panic(fmt.Sprintf("missing segment at %v", end))
   999  		}
  1000  	}
  1001  }
  1002  
  1003  // +stateify savable
  1004  type node struct {
  1005  	// An internal binary tree node looks like:
  1006  	//
  1007  	//   K
  1008  	//  / \
  1009  	// Cl Cr
  1010  	//
  1011  	// where all keys in the subtree rooted by Cl (the left subtree) are less
  1012  	// than K (the key of the parent node), and all keys in the subtree rooted
  1013  	// by Cr (the right subtree) are greater than K.
  1014  	//
  1015  	// An internal B-tree node's indexes work out to look like:
  1016  	//
  1017  	//   K0 K1 K2  ...   Kn-1
  1018  	//  / \/ \/ \  ...  /  \
  1019  	// C0 C1 C2 C3 ... Cn-1 Cn
  1020  	//
  1021  	// where n is nrSegments.
  1022  	nrSegments int
  1023  
  1024  	// parent is a pointer to this node's parent. If this node is root, parent
  1025  	// is nil.
  1026  	parent *node
  1027  
  1028  	// parentIndex is the index of this node in parent.children.
  1029  	parentIndex int
  1030  
  1031  	// Flag for internal nodes that is technically redundant with "children[0]
  1032  	// != nil", but is stored in the first cache line. "hasChildren" rather
  1033  	// than "isLeaf" because false must be the correct value for an empty root.
  1034  	hasChildren bool
  1035  
  1036  	// The longest gap within this node. If the node is a leaf, it's simply the
  1037  	// maximum gap among all the (nrSegments+1) gaps formed by its nrSegments keys
  1038  	// including the 0th and nrSegments-th gap possibly shared with its upper-level
  1039  	// nodes; if it's a non-leaf node, it's the max of all children's maxGap.
  1040  	maxGap dynamicGap
  1041  
  1042  	// Nodes store keys and values in separate arrays to maximize locality in
  1043  	// the common case (scanning keys for lookup).
  1044  	keys     [maxDegree - 1]Range
  1045  	values   [maxDegree - 1]Value
  1046  	children [maxDegree]*node
  1047  }
  1048  
  1049  // firstSegment returns the first segment in the subtree rooted by n.
  1050  //
  1051  // Preconditions: n.nrSegments != 0.
  1052  func (n *node) firstSegment() Iterator {
  1053  	for n.hasChildren {
  1054  		n = n.children[0]
  1055  	}
  1056  	return Iterator{n, 0}
  1057  }
  1058  
  1059  // lastSegment returns the last segment in the subtree rooted by n.
  1060  //
  1061  // Preconditions: n.nrSegments != 0.
  1062  func (n *node) lastSegment() Iterator {
  1063  	for n.hasChildren {
  1064  		n = n.children[n.nrSegments]
  1065  	}
  1066  	return Iterator{n, n.nrSegments - 1}
  1067  }
  1068  
  1069  func (n *node) prevSibling() *node {
  1070  	if n.parent == nil || n.parentIndex == 0 {
  1071  		return nil
  1072  	}
  1073  	return n.parent.children[n.parentIndex-1]
  1074  }
  1075  
  1076  func (n *node) nextSibling() *node {
  1077  	if n.parent == nil || n.parentIndex == n.parent.nrSegments {
  1078  		return nil
  1079  	}
  1080  	return n.parent.children[n.parentIndex+1]
  1081  }
  1082  
  1083  // rebalanceBeforeInsert splits n and its ancestors if they are full, as
  1084  // required for insertion, and returns an updated iterator to the position
  1085  // represented by gap.
  1086  func (n *node) rebalanceBeforeInsert(gap GapIterator) GapIterator {
  1087  	if n.nrSegments < maxDegree-1 {
  1088  		return gap
  1089  	}
  1090  	if n.parent != nil {
  1091  		gap = n.parent.rebalanceBeforeInsert(gap)
  1092  	}
  1093  	if n.parent == nil {
  1094  		// n is root. Move all segments before and after n's median segment
  1095  		// into new child nodes adjacent to the median segment, which is now
  1096  		// the only segment in root.
  1097  		left := &node{
  1098  			nrSegments:  minDegree - 1,
  1099  			parent:      n,
  1100  			parentIndex: 0,
  1101  			hasChildren: n.hasChildren,
  1102  		}
  1103  		right := &node{
  1104  			nrSegments:  minDegree - 1,
  1105  			parent:      n,
  1106  			parentIndex: 1,
  1107  			hasChildren: n.hasChildren,
  1108  		}
  1109  		copy(left.keys[:minDegree-1], n.keys[:minDegree-1])
  1110  		copy(left.values[:minDegree-1], n.values[:minDegree-1])
  1111  		copy(right.keys[:minDegree-1], n.keys[minDegree:])
  1112  		copy(right.values[:minDegree-1], n.values[minDegree:])
  1113  		n.keys[0], n.values[0] = n.keys[minDegree-1], n.values[minDegree-1]
  1114  		zeroValueSlice(n.values[1:])
  1115  		if n.hasChildren {
  1116  			copy(left.children[:minDegree], n.children[:minDegree])
  1117  			copy(right.children[:minDegree], n.children[minDegree:])
  1118  			zeroNodeSlice(n.children[2:])
  1119  			for i := 0; i < minDegree; i++ {
  1120  				left.children[i].parent = left
  1121  				left.children[i].parentIndex = i
  1122  				right.children[i].parent = right
  1123  				right.children[i].parentIndex = i
  1124  			}
  1125  		}
  1126  		n.nrSegments = 1
  1127  		n.hasChildren = true
  1128  		n.children[0] = left
  1129  		n.children[1] = right
  1130  		// In this case, n's maxGap won't violated as it's still the root,
  1131  		// but the left and right children should be updated locally as they
  1132  		// are newly split from n.
  1133  		if trackGaps != 0 {
  1134  			left.updateMaxGapLocal()
  1135  			right.updateMaxGapLocal()
  1136  		}
  1137  		if gap.node != n {
  1138  			return gap
  1139  		}
  1140  		if gap.index < minDegree {
  1141  			return GapIterator{left, gap.index}
  1142  		}
  1143  		return GapIterator{right, gap.index - minDegree}
  1144  	}
  1145  	// n is non-root. Move n's median segment into its parent node (which can't
  1146  	// be full because we've already invoked n.parent.rebalanceBeforeInsert)
  1147  	// and move all segments after n's median into a new sibling node (the
  1148  	// median segment's right child subtree).
  1149  	copy(n.parent.keys[n.parentIndex+1:], n.parent.keys[n.parentIndex:n.parent.nrSegments])
  1150  	copy(n.parent.values[n.parentIndex+1:], n.parent.values[n.parentIndex:n.parent.nrSegments])
  1151  	n.parent.keys[n.parentIndex], n.parent.values[n.parentIndex] = n.keys[minDegree-1], n.values[minDegree-1]
  1152  	copy(n.parent.children[n.parentIndex+2:], n.parent.children[n.parentIndex+1:n.parent.nrSegments+1])
  1153  	for i := n.parentIndex + 2; i < n.parent.nrSegments+2; i++ {
  1154  		n.parent.children[i].parentIndex = i
  1155  	}
  1156  	sibling := &node{
  1157  		nrSegments:  minDegree - 1,
  1158  		parent:      n.parent,
  1159  		parentIndex: n.parentIndex + 1,
  1160  		hasChildren: n.hasChildren,
  1161  	}
  1162  	n.parent.children[n.parentIndex+1] = sibling
  1163  	n.parent.nrSegments++
  1164  	copy(sibling.keys[:minDegree-1], n.keys[minDegree:])
  1165  	copy(sibling.values[:minDegree-1], n.values[minDegree:])
  1166  	zeroValueSlice(n.values[minDegree-1:])
  1167  	if n.hasChildren {
  1168  		copy(sibling.children[:minDegree], n.children[minDegree:])
  1169  		zeroNodeSlice(n.children[minDegree:])
  1170  		for i := 0; i < minDegree; i++ {
  1171  			sibling.children[i].parent = sibling
  1172  			sibling.children[i].parentIndex = i
  1173  		}
  1174  	}
  1175  	n.nrSegments = minDegree - 1
  1176  	// MaxGap of n's parent is not violated because the segments within is not changed.
  1177  	// n and its sibling's maxGap need to be updated locally as they are two new nodes split from old n.
  1178  	if trackGaps != 0 {
  1179  		n.updateMaxGapLocal()
  1180  		sibling.updateMaxGapLocal()
  1181  	}
  1182  	// gap.node can't be n.parent because gaps are always in leaf nodes.
  1183  	if gap.node != n {
  1184  		return gap
  1185  	}
  1186  	if gap.index < minDegree {
  1187  		return gap
  1188  	}
  1189  	return GapIterator{sibling, gap.index - minDegree}
  1190  }
  1191  
  1192  // rebalanceAfterRemove "unsplits" n and its ancestors if they are deficient
  1193  // (contain fewer segments than required by B-tree invariants), as required for
  1194  // removal, and returns an updated iterator to the position represented by gap.
  1195  //
  1196  // Precondition: n is the only node in the tree that may currently violate a
  1197  // B-tree invariant.
  1198  func (n *node) rebalanceAfterRemove(gap GapIterator) GapIterator {
  1199  	for {
  1200  		if n.nrSegments >= minDegree-1 {
  1201  			return gap
  1202  		}
  1203  		if n.parent == nil {
  1204  			// Root is allowed to be deficient.
  1205  			return gap
  1206  		}
  1207  		// There's one other thing we can do before resorting to unsplitting.
  1208  		// If either sibling node has at least minDegree segments, rotate that
  1209  		// sibling's closest segment through the segment in the parent that
  1210  		// separates us. That is, given:
  1211  		//
  1212  		//      ... D ...
  1213  		//         / \
  1214  		// ... B C]   [E ...
  1215  		//
  1216  		// where the node containing E is deficient, end up with:
  1217  		//
  1218  		//    ... C ...
  1219  		//       / \
  1220  		// ... B]   [D E ...
  1221  		//
  1222  		// As in Set.Remove, prefer rotating from the end of the sibling to the
  1223  		// left: by precondition, n.node has fewer segments (to memcpy) than
  1224  		// the sibling does.
  1225  		if sibling := n.prevSibling(); sibling != nil && sibling.nrSegments >= minDegree {
  1226  			copy(n.keys[1:], n.keys[:n.nrSegments])
  1227  			copy(n.values[1:], n.values[:n.nrSegments])
  1228  			n.keys[0] = n.parent.keys[n.parentIndex-1]
  1229  			n.values[0] = n.parent.values[n.parentIndex-1]
  1230  			n.parent.keys[n.parentIndex-1] = sibling.keys[sibling.nrSegments-1]
  1231  			n.parent.values[n.parentIndex-1] = sibling.values[sibling.nrSegments-1]
  1232  			Functions{}.ClearValue(&sibling.values[sibling.nrSegments-1])
  1233  			if n.hasChildren {
  1234  				copy(n.children[1:], n.children[:n.nrSegments+1])
  1235  				n.children[0] = sibling.children[sibling.nrSegments]
  1236  				sibling.children[sibling.nrSegments] = nil
  1237  				n.children[0].parent = n
  1238  				n.children[0].parentIndex = 0
  1239  				for i := 1; i < n.nrSegments+2; i++ {
  1240  					n.children[i].parentIndex = i
  1241  				}
  1242  			}
  1243  			n.nrSegments++
  1244  			sibling.nrSegments--
  1245  			// n's parent's maxGap does not need to be updated as its content is unmodified.
  1246  			// n and its sibling must be updated with (new) maxGap because of the shift of keys.
  1247  			if trackGaps != 0 {
  1248  				n.updateMaxGapLocal()
  1249  				sibling.updateMaxGapLocal()
  1250  			}
  1251  			if gap.node == sibling && gap.index == sibling.nrSegments {
  1252  				return GapIterator{n, 0}
  1253  			}
  1254  			if gap.node == n {
  1255  				return GapIterator{n, gap.index + 1}
  1256  			}
  1257  			return gap
  1258  		}
  1259  		if sibling := n.nextSibling(); sibling != nil && sibling.nrSegments >= minDegree {
  1260  			n.keys[n.nrSegments] = n.parent.keys[n.parentIndex]
  1261  			n.values[n.nrSegments] = n.parent.values[n.parentIndex]
  1262  			n.parent.keys[n.parentIndex] = sibling.keys[0]
  1263  			n.parent.values[n.parentIndex] = sibling.values[0]
  1264  			copy(sibling.keys[:sibling.nrSegments-1], sibling.keys[1:])
  1265  			copy(sibling.values[:sibling.nrSegments-1], sibling.values[1:])
  1266  			Functions{}.ClearValue(&sibling.values[sibling.nrSegments-1])
  1267  			if n.hasChildren {
  1268  				n.children[n.nrSegments+1] = sibling.children[0]
  1269  				copy(sibling.children[:sibling.nrSegments], sibling.children[1:])
  1270  				sibling.children[sibling.nrSegments] = nil
  1271  				n.children[n.nrSegments+1].parent = n
  1272  				n.children[n.nrSegments+1].parentIndex = n.nrSegments + 1
  1273  				for i := 0; i < sibling.nrSegments; i++ {
  1274  					sibling.children[i].parentIndex = i
  1275  				}
  1276  			}
  1277  			n.nrSegments++
  1278  			sibling.nrSegments--
  1279  			// n's parent's maxGap does not need to be updated as its content is unmodified.
  1280  			// n and its sibling must be updated with (new) maxGap because of the shift of keys.
  1281  			if trackGaps != 0 {
  1282  				n.updateMaxGapLocal()
  1283  				sibling.updateMaxGapLocal()
  1284  			}
  1285  			if gap.node == sibling {
  1286  				if gap.index == 0 {
  1287  					return GapIterator{n, n.nrSegments}
  1288  				}
  1289  				return GapIterator{sibling, gap.index - 1}
  1290  			}
  1291  			return gap
  1292  		}
  1293  		// Otherwise, we must unsplit.
  1294  		p := n.parent
  1295  		if p.nrSegments == 1 {
  1296  			// Merge all segments in both n and its sibling back into n.parent.
  1297  			// This is the reverse of the root splitting case in
  1298  			// node.rebalanceBeforeInsert. (Because we require minDegree >= 3,
  1299  			// only root can have 1 segment in this path, so this reduces the
  1300  			// height of the tree by 1, without violating the constraint that
  1301  			// all leaf nodes remain at the same depth.)
  1302  			left, right := p.children[0], p.children[1]
  1303  			p.nrSegments = left.nrSegments + right.nrSegments + 1
  1304  			p.hasChildren = left.hasChildren
  1305  			p.keys[left.nrSegments] = p.keys[0]
  1306  			p.values[left.nrSegments] = p.values[0]
  1307  			copy(p.keys[:left.nrSegments], left.keys[:left.nrSegments])
  1308  			copy(p.values[:left.nrSegments], left.values[:left.nrSegments])
  1309  			copy(p.keys[left.nrSegments+1:], right.keys[:right.nrSegments])
  1310  			copy(p.values[left.nrSegments+1:], right.values[:right.nrSegments])
  1311  			if left.hasChildren {
  1312  				copy(p.children[:left.nrSegments+1], left.children[:left.nrSegments+1])
  1313  				copy(p.children[left.nrSegments+1:], right.children[:right.nrSegments+1])
  1314  				for i := 0; i < p.nrSegments+1; i++ {
  1315  					p.children[i].parent = p
  1316  					p.children[i].parentIndex = i
  1317  				}
  1318  			} else {
  1319  				p.children[0] = nil
  1320  				p.children[1] = nil
  1321  			}
  1322  			// No need to update maxGap of p as its content is not changed.
  1323  			if gap.node == left {
  1324  				return GapIterator{p, gap.index}
  1325  			}
  1326  			if gap.node == right {
  1327  				return GapIterator{p, gap.index + left.nrSegments + 1}
  1328  			}
  1329  			return gap
  1330  		}
  1331  		// Merge n and either sibling, along with the segment separating the
  1332  		// two, into whichever of the two nodes comes first. This is the
  1333  		// reverse of the non-root splitting case in
  1334  		// node.rebalanceBeforeInsert.
  1335  		var left, right *node
  1336  		if n.parentIndex > 0 {
  1337  			left = n.prevSibling()
  1338  			right = n
  1339  		} else {
  1340  			left = n
  1341  			right = n.nextSibling()
  1342  		}
  1343  		// Fix up gap first since we need the old left.nrSegments, which
  1344  		// merging will change.
  1345  		if gap.node == right {
  1346  			gap = GapIterator{left, gap.index + left.nrSegments + 1}
  1347  		}
  1348  		left.keys[left.nrSegments] = p.keys[left.parentIndex]
  1349  		left.values[left.nrSegments] = p.values[left.parentIndex]
  1350  		copy(left.keys[left.nrSegments+1:], right.keys[:right.nrSegments])
  1351  		copy(left.values[left.nrSegments+1:], right.values[:right.nrSegments])
  1352  		if left.hasChildren {
  1353  			copy(left.children[left.nrSegments+1:], right.children[:right.nrSegments+1])
  1354  			for i := left.nrSegments + 1; i < left.nrSegments+right.nrSegments+2; i++ {
  1355  				left.children[i].parent = left
  1356  				left.children[i].parentIndex = i
  1357  			}
  1358  		}
  1359  		left.nrSegments += right.nrSegments + 1
  1360  		copy(p.keys[left.parentIndex:], p.keys[left.parentIndex+1:p.nrSegments])
  1361  		copy(p.values[left.parentIndex:], p.values[left.parentIndex+1:p.nrSegments])
  1362  		Functions{}.ClearValue(&p.values[p.nrSegments-1])
  1363  		copy(p.children[left.parentIndex+1:], p.children[left.parentIndex+2:p.nrSegments+1])
  1364  		for i := 0; i < p.nrSegments; i++ {
  1365  			p.children[i].parentIndex = i
  1366  		}
  1367  		p.children[p.nrSegments] = nil
  1368  		p.nrSegments--
  1369  		// Update maxGap of left locally, no need to change p and right because
  1370  		// p's contents is not changed and right is already invalid.
  1371  		if trackGaps != 0 {
  1372  			left.updateMaxGapLocal()
  1373  		}
  1374  		// This process robs p of one segment, so recurse into rebalancing p.
  1375  		n = p
  1376  	}
  1377  }
  1378  
  1379  // updateMaxGapLeaf updates maxGap bottom-up from the calling leaf until no
  1380  // necessary update.
  1381  //
  1382  // Preconditions: n must be a leaf node, trackGaps must be 1.
  1383  func (n *node) updateMaxGapLeaf() {
  1384  	if n.hasChildren {
  1385  		panic(fmt.Sprintf("updateMaxGapLeaf should always be called on leaf node: %v", n))
  1386  	}
  1387  	max := n.calculateMaxGapLeaf()
  1388  	if max == n.maxGap.Get() {
  1389  		// If new max equals the old maxGap, no update is needed.
  1390  		return
  1391  	}
  1392  	oldMax := n.maxGap.Get()
  1393  	n.maxGap.Set(max)
  1394  	if max > oldMax {
  1395  		// Grow ancestor maxGaps.
  1396  		for p := n.parent; p != nil; p = p.parent {
  1397  			if p.maxGap.Get() >= max {
  1398  				// p and its ancestors already contain an equal or larger gap.
  1399  				break
  1400  			}
  1401  			// Only if new maxGap is larger than parent's
  1402  			// old maxGap, propagate this update to parent.
  1403  			p.maxGap.Set(max)
  1404  		}
  1405  		return
  1406  	}
  1407  	// Shrink ancestor maxGaps.
  1408  	for p := n.parent; p != nil; p = p.parent {
  1409  		if p.maxGap.Get() > oldMax {
  1410  			// p and its ancestors still contain a larger gap.
  1411  			break
  1412  		}
  1413  		// If new max is smaller than the old maxGap, and this gap used
  1414  		// to be the maxGap of its parent, iterate parent's children
  1415  		// and calculate parent's new maxGap.(It's probable that parent
  1416  		// has two children with the old maxGap, but we need to check it anyway.)
  1417  		parentNewMax := p.calculateMaxGapInternal()
  1418  		if p.maxGap.Get() == parentNewMax {
  1419  			// p and its ancestors still contain a gap of at least equal size.
  1420  			break
  1421  		}
  1422  		// If p's new maxGap differs from the old one, propagate this update.
  1423  		p.maxGap.Set(parentNewMax)
  1424  	}
  1425  }
  1426  
  1427  // updateMaxGapLocal updates maxGap of the calling node solely with no
  1428  // propagation to ancestor nodes.
  1429  //
  1430  // Precondition: trackGaps must be 1.
  1431  func (n *node) updateMaxGapLocal() {
  1432  	if !n.hasChildren {
  1433  		// Leaf node iterates its gaps.
  1434  		n.maxGap.Set(n.calculateMaxGapLeaf())
  1435  	} else {
  1436  		// Non-leaf node iterates its children.
  1437  		n.maxGap.Set(n.calculateMaxGapInternal())
  1438  	}
  1439  }
  1440  
  1441  // calculateMaxGapLeaf iterates the gaps within a leaf node and calculate the
  1442  // max.
  1443  //
  1444  // Preconditions: n must be a leaf node.
  1445  func (n *node) calculateMaxGapLeaf() Key {
  1446  	max := GapIterator{n, 0}.Range().Length()
  1447  	for i := 1; i <= n.nrSegments; i++ {
  1448  		if current := (GapIterator{n, i}).Range().Length(); current > max {
  1449  			max = current
  1450  		}
  1451  	}
  1452  	return max
  1453  }
  1454  
  1455  // calculateMaxGapInternal iterates children's maxGap within an internal node n
  1456  // and calculate the max.
  1457  //
  1458  // Preconditions: n must be a non-leaf node.
  1459  func (n *node) calculateMaxGapInternal() Key {
  1460  	max := n.children[0].maxGap.Get()
  1461  	for i := 1; i <= n.nrSegments; i++ {
  1462  		if current := n.children[i].maxGap.Get(); current > max {
  1463  			max = current
  1464  		}
  1465  	}
  1466  	return max
  1467  }
  1468  
  1469  // searchFirstLargeEnoughGap returns the first gap having at least minSize length
  1470  // in the subtree rooted by n. If not found, return a terminal gap iterator.
  1471  func (n *node) searchFirstLargeEnoughGap(minSize Key) GapIterator {
  1472  	if n.maxGap.Get() < minSize {
  1473  		return GapIterator{}
  1474  	}
  1475  	if n.hasChildren {
  1476  		for i := 0; i <= n.nrSegments; i++ {
  1477  			if largeEnoughGap := n.children[i].searchFirstLargeEnoughGap(minSize); largeEnoughGap.Ok() {
  1478  				return largeEnoughGap
  1479  			}
  1480  		}
  1481  	} else {
  1482  		for i := 0; i <= n.nrSegments; i++ {
  1483  			currentGap := GapIterator{n, i}
  1484  			if currentGap.Range().Length() >= minSize {
  1485  				return currentGap
  1486  			}
  1487  		}
  1488  	}
  1489  	panic(fmt.Sprintf("invalid maxGap in %v", n))
  1490  }
  1491  
  1492  // searchLastLargeEnoughGap returns the last gap having at least minSize length
  1493  // in the subtree rooted by n. If not found, return a terminal gap iterator.
  1494  func (n *node) searchLastLargeEnoughGap(minSize Key) GapIterator {
  1495  	if n.maxGap.Get() < minSize {
  1496  		return GapIterator{}
  1497  	}
  1498  	if n.hasChildren {
  1499  		for i := n.nrSegments; i >= 0; i-- {
  1500  			if largeEnoughGap := n.children[i].searchLastLargeEnoughGap(minSize); largeEnoughGap.Ok() {
  1501  				return largeEnoughGap
  1502  			}
  1503  		}
  1504  	} else {
  1505  		for i := n.nrSegments; i >= 0; i-- {
  1506  			currentGap := GapIterator{n, i}
  1507  			if currentGap.Range().Length() >= minSize {
  1508  				return currentGap
  1509  			}
  1510  		}
  1511  	}
  1512  	panic(fmt.Sprintf("invalid maxGap in %v", n))
  1513  }
  1514  
  1515  // A Iterator is conceptually one of:
  1516  //
  1517  //   - A pointer to a segment in a set; or
  1518  //
  1519  //   - A terminal iterator, which is a sentinel indicating that the end of
  1520  //     iteration has been reached.
  1521  //
  1522  // Iterators are copyable values and are meaningfully equality-comparable. The
  1523  // zero value of Iterator is a terminal iterator.
  1524  //
  1525  // Unless otherwise specified, any mutation of a set invalidates all existing
  1526  // iterators into the set.
  1527  type Iterator struct {
  1528  	// node is the node containing the iterated segment. If the iterator is
  1529  	// terminal, node is nil.
  1530  	node *node
  1531  
  1532  	// index is the index of the segment in node.keys/values.
  1533  	index int
  1534  }
  1535  
  1536  // Ok returns true if the iterator is not terminal. All other methods are only
  1537  // valid for non-terminal iterators.
  1538  func (seg Iterator) Ok() bool {
  1539  	return seg.node != nil
  1540  }
  1541  
  1542  // Range returns the iterated segment's range key.
  1543  func (seg Iterator) Range() Range {
  1544  	return seg.node.keys[seg.index]
  1545  }
  1546  
  1547  // Start is equivalent to Range().Start, but should be preferred if only the
  1548  // start of the range is needed.
  1549  func (seg Iterator) Start() Key {
  1550  	return seg.node.keys[seg.index].Start
  1551  }
  1552  
  1553  // End is equivalent to Range().End, but should be preferred if only the end of
  1554  // the range is needed.
  1555  func (seg Iterator) End() Key {
  1556  	return seg.node.keys[seg.index].End
  1557  }
  1558  
  1559  // SetRangeUnchecked mutates the iterated segment's range key. This operation
  1560  // does not invalidate any iterators.
  1561  //
  1562  // Preconditions:
  1563  // - r.Length() > 0.
  1564  // - The new range must not overlap an existing one:
  1565  //   - If seg.NextSegment().Ok(), then r.end <= seg.NextSegment().Start().
  1566  //   - If seg.PrevSegment().Ok(), then r.start >= seg.PrevSegment().End().
  1567  func (seg Iterator) SetRangeUnchecked(r Range) {
  1568  	seg.node.keys[seg.index] = r
  1569  }
  1570  
  1571  // SetRange mutates the iterated segment's range key. If the new range would
  1572  // cause the iterated segment to overlap another segment, or if the new range
  1573  // is invalid, SetRange panics. This operation does not invalidate any
  1574  // iterators.
  1575  func (seg Iterator) SetRange(r Range) {
  1576  	if r.Length() <= 0 {
  1577  		panic(fmt.Sprintf("invalid segment range %v", r))
  1578  	}
  1579  	if prev := seg.PrevSegment(); prev.Ok() && r.Start < prev.End() {
  1580  		panic(fmt.Sprintf("new segment range %v overlaps segment range %v", r, prev.Range()))
  1581  	}
  1582  	if next := seg.NextSegment(); next.Ok() && r.End > next.Start() {
  1583  		panic(fmt.Sprintf("new segment range %v overlaps segment range %v", r, next.Range()))
  1584  	}
  1585  	seg.SetRangeUnchecked(r)
  1586  }
  1587  
  1588  // SetStartUnchecked mutates the iterated segment's start. This operation does
  1589  // not invalidate any iterators.
  1590  //
  1591  // Preconditions: The new start must be valid:
  1592  //   - start < seg.End()
  1593  //   - If seg.PrevSegment().Ok(), then start >= seg.PrevSegment().End().
  1594  func (seg Iterator) SetStartUnchecked(start Key) {
  1595  	seg.node.keys[seg.index].Start = start
  1596  }
  1597  
  1598  // SetStart mutates the iterated segment's start. If the new start value would
  1599  // cause the iterated segment to overlap another segment, or would result in an
  1600  // invalid range, SetStart panics. This operation does not invalidate any
  1601  // iterators.
  1602  func (seg Iterator) SetStart(start Key) {
  1603  	if start >= seg.End() {
  1604  		panic(fmt.Sprintf("new start %v would invalidate segment range %v", start, seg.Range()))
  1605  	}
  1606  	if prev := seg.PrevSegment(); prev.Ok() && start < prev.End() {
  1607  		panic(fmt.Sprintf("new start %v would cause segment range %v to overlap segment range %v", start, seg.Range(), prev.Range()))
  1608  	}
  1609  	seg.SetStartUnchecked(start)
  1610  }
  1611  
  1612  // SetEndUnchecked mutates the iterated segment's end. This operation does not
  1613  // invalidate any iterators.
  1614  //
  1615  // Preconditions: The new end must be valid:
  1616  //   - end > seg.Start().
  1617  //   - If seg.NextSegment().Ok(), then end <= seg.NextSegment().Start().
  1618  func (seg Iterator) SetEndUnchecked(end Key) {
  1619  	seg.node.keys[seg.index].End = end
  1620  }
  1621  
  1622  // SetEnd mutates the iterated segment's end. If the new end value would cause
  1623  // the iterated segment to overlap another segment, or would result in an
  1624  // invalid range, SetEnd panics. This operation does not invalidate any
  1625  // iterators.
  1626  func (seg Iterator) SetEnd(end Key) {
  1627  	if end <= seg.Start() {
  1628  		panic(fmt.Sprintf("new end %v would invalidate segment range %v", end, seg.Range()))
  1629  	}
  1630  	if next := seg.NextSegment(); next.Ok() && end > next.Start() {
  1631  		panic(fmt.Sprintf("new end %v would cause segment range %v to overlap segment range %v", end, seg.Range(), next.Range()))
  1632  	}
  1633  	seg.SetEndUnchecked(end)
  1634  }
  1635  
  1636  // Value returns a copy of the iterated segment's value.
  1637  func (seg Iterator) Value() Value {
  1638  	return seg.node.values[seg.index]
  1639  }
  1640  
  1641  // ValuePtr returns a pointer to the iterated segment's value. The pointer is
  1642  // invalidated if the iterator is invalidated. This operation does not
  1643  // invalidate any iterators.
  1644  func (seg Iterator) ValuePtr() *Value {
  1645  	return &seg.node.values[seg.index]
  1646  }
  1647  
  1648  // SetValue mutates the iterated segment's value. This operation does not
  1649  // invalidate any iterators.
  1650  func (seg Iterator) SetValue(val Value) {
  1651  	seg.node.values[seg.index] = val
  1652  }
  1653  
  1654  // PrevSegment returns the iterated segment's predecessor. If there is no
  1655  // preceding segment, PrevSegment returns a terminal iterator.
  1656  func (seg Iterator) PrevSegment() Iterator {
  1657  	if seg.node.hasChildren {
  1658  		return seg.node.children[seg.index].lastSegment()
  1659  	}
  1660  	if seg.index > 0 {
  1661  		return Iterator{seg.node, seg.index - 1}
  1662  	}
  1663  	if seg.node.parent == nil {
  1664  		return Iterator{}
  1665  	}
  1666  	return segmentBeforePosition(seg.node.parent, seg.node.parentIndex)
  1667  }
  1668  
  1669  // NextSegment returns the iterated segment's successor. If there is no
  1670  // succeeding segment, NextSegment returns a terminal iterator.
  1671  func (seg Iterator) NextSegment() Iterator {
  1672  	if seg.node.hasChildren {
  1673  		return seg.node.children[seg.index+1].firstSegment()
  1674  	}
  1675  	if seg.index < seg.node.nrSegments-1 {
  1676  		return Iterator{seg.node, seg.index + 1}
  1677  	}
  1678  	if seg.node.parent == nil {
  1679  		return Iterator{}
  1680  	}
  1681  	return segmentAfterPosition(seg.node.parent, seg.node.parentIndex)
  1682  }
  1683  
  1684  // PrevGap returns the gap immediately before the iterated segment.
  1685  func (seg Iterator) PrevGap() GapIterator {
  1686  	if seg.node.hasChildren {
  1687  		// Note that this isn't recursive because the last segment in a subtree
  1688  		// must be in a leaf node.
  1689  		return seg.node.children[seg.index].lastSegment().NextGap()
  1690  	}
  1691  	return GapIterator{seg.node, seg.index}
  1692  }
  1693  
  1694  // NextGap returns the gap immediately after the iterated segment.
  1695  func (seg Iterator) NextGap() GapIterator {
  1696  	if seg.node.hasChildren {
  1697  		return seg.node.children[seg.index+1].firstSegment().PrevGap()
  1698  	}
  1699  	return GapIterator{seg.node, seg.index + 1}
  1700  }
  1701  
  1702  // PrevNonEmpty returns the iterated segment's predecessor if it is adjacent,
  1703  // or the gap before the iterated segment otherwise. If seg.Start() ==
  1704  // Functions.MinKey(), PrevNonEmpty will return two terminal iterators.
  1705  // Otherwise, exactly one of the iterators returned by PrevNonEmpty will be
  1706  // non-terminal.
  1707  func (seg Iterator) PrevNonEmpty() (Iterator, GapIterator) {
  1708  	if prev := seg.PrevSegment(); prev.Ok() && prev.End() == seg.Start() {
  1709  		return prev, GapIterator{}
  1710  	}
  1711  	return Iterator{}, seg.PrevGap()
  1712  }
  1713  
  1714  // NextNonEmpty returns the iterated segment's successor if it is adjacent, or
  1715  // the gap after the iterated segment otherwise. If seg.End() ==
  1716  // Functions.MaxKey(), NextNonEmpty will return two terminal iterators.
  1717  // Otherwise, exactly one of the iterators returned by NextNonEmpty will be
  1718  // non-terminal.
  1719  func (seg Iterator) NextNonEmpty() (Iterator, GapIterator) {
  1720  	if next := seg.NextSegment(); next.Ok() && next.Start() == seg.End() {
  1721  		return next, GapIterator{}
  1722  	}
  1723  	return Iterator{}, seg.NextGap()
  1724  }
  1725  
  1726  // A GapIterator is conceptually one of:
  1727  //
  1728  //   - A pointer to a position between two segments, before the first segment, or
  1729  //     after the last segment in a set, called a *gap*; or
  1730  //
  1731  //   - A terminal iterator, which is a sentinel indicating that the end of
  1732  //     iteration has been reached.
  1733  //
  1734  // Note that the gap between two adjacent segments exists (iterators to it are
  1735  // non-terminal), but has a length of zero. GapIterator.IsEmpty returns true
  1736  // for such gaps. An empty set contains a single gap, spanning the entire range
  1737  // of the set's keys.
  1738  //
  1739  // GapIterators are copyable values and are meaningfully equality-comparable.
  1740  // The zero value of GapIterator is a terminal iterator.
  1741  //
  1742  // Unless otherwise specified, any mutation of a set invalidates all existing
  1743  // iterators into the set.
  1744  type GapIterator struct {
  1745  	// The representation of a GapIterator is identical to that of an Iterator,
  1746  	// except that index corresponds to positions between segments in the same
  1747  	// way as for node.children (see comment for node.nrSegments).
  1748  	node  *node
  1749  	index int
  1750  }
  1751  
  1752  // Ok returns true if the iterator is not terminal. All other methods are only
  1753  // valid for non-terminal iterators.
  1754  func (gap GapIterator) Ok() bool {
  1755  	return gap.node != nil
  1756  }
  1757  
  1758  // Range returns the range spanned by the iterated gap.
  1759  func (gap GapIterator) Range() Range {
  1760  	return Range{gap.Start(), gap.End()}
  1761  }
  1762  
  1763  // Start is equivalent to Range().Start, but should be preferred if only the
  1764  // start of the range is needed.
  1765  func (gap GapIterator) Start() Key {
  1766  	if ps := gap.PrevSegment(); ps.Ok() {
  1767  		return ps.End()
  1768  	}
  1769  	return Functions{}.MinKey()
  1770  }
  1771  
  1772  // End is equivalent to Range().End, but should be preferred if only the end of
  1773  // the range is needed.
  1774  func (gap GapIterator) End() Key {
  1775  	if ns := gap.NextSegment(); ns.Ok() {
  1776  		return ns.Start()
  1777  	}
  1778  	return Functions{}.MaxKey()
  1779  }
  1780  
  1781  // IsEmpty returns true if the iterated gap is empty (that is, the "gap" is
  1782  // between two adjacent segments.)
  1783  func (gap GapIterator) IsEmpty() bool {
  1784  	return gap.Range().Length() == 0
  1785  }
  1786  
  1787  // PrevSegment returns the segment immediately before the iterated gap. If no
  1788  // such segment exists, PrevSegment returns a terminal iterator.
  1789  func (gap GapIterator) PrevSegment() Iterator {
  1790  	return segmentBeforePosition(gap.node, gap.index)
  1791  }
  1792  
  1793  // NextSegment returns the segment immediately after the iterated gap. If no
  1794  // such segment exists, NextSegment returns a terminal iterator.
  1795  func (gap GapIterator) NextSegment() Iterator {
  1796  	return segmentAfterPosition(gap.node, gap.index)
  1797  }
  1798  
  1799  // PrevGap returns the iterated gap's predecessor. If no such gap exists,
  1800  // PrevGap returns a terminal iterator.
  1801  func (gap GapIterator) PrevGap() GapIterator {
  1802  	seg := gap.PrevSegment()
  1803  	if !seg.Ok() {
  1804  		return GapIterator{}
  1805  	}
  1806  	return seg.PrevGap()
  1807  }
  1808  
  1809  // NextGap returns the iterated gap's successor. If no such gap exists, NextGap
  1810  // returns a terminal iterator.
  1811  func (gap GapIterator) NextGap() GapIterator {
  1812  	seg := gap.NextSegment()
  1813  	if !seg.Ok() {
  1814  		return GapIterator{}
  1815  	}
  1816  	return seg.NextGap()
  1817  }
  1818  
  1819  // NextLargeEnoughGap returns the iterated gap's first next gap with larger
  1820  // length than minSize.  If not found, return a terminal gap iterator (does NOT
  1821  // include this gap itself).
  1822  //
  1823  // Precondition: trackGaps must be 1.
  1824  func (gap GapIterator) NextLargeEnoughGap(minSize Key) GapIterator {
  1825  	if trackGaps != 1 {
  1826  		panic("set is not tracking gaps")
  1827  	}
  1828  	if gap.node != nil && gap.node.hasChildren && gap.index == gap.node.nrSegments {
  1829  		// If gap is the trailing gap of an non-leaf node,
  1830  		// translate it to the equivalent gap on leaf level.
  1831  		gap.node = gap.NextSegment().node
  1832  		gap.index = 0
  1833  		return gap.nextLargeEnoughGapHelper(minSize)
  1834  	}
  1835  	return gap.nextLargeEnoughGapHelper(minSize)
  1836  }
  1837  
  1838  // nextLargeEnoughGapHelper is the helper function used by NextLargeEnoughGap
  1839  // to do the real recursions.
  1840  //
  1841  // Preconditions: gap is NOT the trailing gap of a non-leaf node.
  1842  func (gap GapIterator) nextLargeEnoughGapHelper(minSize Key) GapIterator {
  1843  	for {
  1844  		// Crawl up the tree if no large enough gap in current node or the
  1845  		// current gap is the trailing one on leaf level.
  1846  		for gap.node != nil &&
  1847  			(gap.node.maxGap.Get() < minSize || (!gap.node.hasChildren && gap.index == gap.node.nrSegments)) {
  1848  			gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1849  		}
  1850  		// If no large enough gap throughout the whole set, return a terminal
  1851  		// gap iterator.
  1852  		if gap.node == nil {
  1853  			return GapIterator{}
  1854  		}
  1855  		// Iterate subsequent gaps.
  1856  		gap.index++
  1857  		for gap.index <= gap.node.nrSegments {
  1858  			if gap.node.hasChildren {
  1859  				if largeEnoughGap := gap.node.children[gap.index].searchFirstLargeEnoughGap(minSize); largeEnoughGap.Ok() {
  1860  					return largeEnoughGap
  1861  				}
  1862  			} else {
  1863  				if gap.Range().Length() >= minSize {
  1864  					return gap
  1865  				}
  1866  			}
  1867  			gap.index++
  1868  		}
  1869  		gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1870  		if gap.node != nil && gap.index == gap.node.nrSegments {
  1871  			// If gap is the trailing gap of a non-leaf node, crawl up to
  1872  			// parent again and do recursion.
  1873  			gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1874  		}
  1875  	}
  1876  }
  1877  
  1878  // PrevLargeEnoughGap returns the iterated gap's first prev gap with larger or
  1879  // equal length than minSize.  If not found, return a terminal gap iterator
  1880  // (does NOT include this gap itself).
  1881  //
  1882  // Precondition: trackGaps must be 1.
  1883  func (gap GapIterator) PrevLargeEnoughGap(minSize Key) GapIterator {
  1884  	if trackGaps != 1 {
  1885  		panic("set is not tracking gaps")
  1886  	}
  1887  	if gap.node != nil && gap.node.hasChildren && gap.index == 0 {
  1888  		// If gap is the first gap of an non-leaf node,
  1889  		// translate it to the equivalent gap on leaf level.
  1890  		gap.node = gap.PrevSegment().node
  1891  		gap.index = gap.node.nrSegments
  1892  		return gap.prevLargeEnoughGapHelper(minSize)
  1893  	}
  1894  	return gap.prevLargeEnoughGapHelper(minSize)
  1895  }
  1896  
  1897  // prevLargeEnoughGapHelper is the helper function used by PrevLargeEnoughGap
  1898  // to do the real recursions.
  1899  //
  1900  // Preconditions: gap is NOT the first gap of a non-leaf node.
  1901  func (gap GapIterator) prevLargeEnoughGapHelper(minSize Key) GapIterator {
  1902  	for {
  1903  		// Crawl up the tree if no large enough gap in current node or the
  1904  		// current gap is the first one on leaf level.
  1905  		for gap.node != nil &&
  1906  			(gap.node.maxGap.Get() < minSize || (!gap.node.hasChildren && gap.index == 0)) {
  1907  			gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1908  		}
  1909  		// If no large enough gap throughout the whole set, return a terminal
  1910  		// gap iterator.
  1911  		if gap.node == nil {
  1912  			return GapIterator{}
  1913  		}
  1914  		// Iterate previous gaps.
  1915  		gap.index--
  1916  		for gap.index >= 0 {
  1917  			if gap.node.hasChildren {
  1918  				if largeEnoughGap := gap.node.children[gap.index].searchLastLargeEnoughGap(minSize); largeEnoughGap.Ok() {
  1919  					return largeEnoughGap
  1920  				}
  1921  			} else {
  1922  				if gap.Range().Length() >= minSize {
  1923  					return gap
  1924  				}
  1925  			}
  1926  			gap.index--
  1927  		}
  1928  		gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1929  		if gap.node != nil && gap.index == 0 {
  1930  			// If gap is the first gap of a non-leaf node, crawl up to
  1931  			// parent again and do recursion.
  1932  			gap.node, gap.index = gap.node.parent, gap.node.parentIndex
  1933  		}
  1934  	}
  1935  }
  1936  
  1937  // segmentBeforePosition returns the predecessor segment of the position given
  1938  // by n.children[i], which may or may not contain a child. If no such segment
  1939  // exists, segmentBeforePosition returns a terminal iterator.
  1940  func segmentBeforePosition(n *node, i int) Iterator {
  1941  	for i == 0 {
  1942  		if n.parent == nil {
  1943  			return Iterator{}
  1944  		}
  1945  		n, i = n.parent, n.parentIndex
  1946  	}
  1947  	return Iterator{n, i - 1}
  1948  }
  1949  
  1950  // segmentAfterPosition returns the successor segment of the position given by
  1951  // n.children[i], which may or may not contain a child. If no such segment
  1952  // exists, segmentAfterPosition returns a terminal iterator.
  1953  func segmentAfterPosition(n *node, i int) Iterator {
  1954  	for i == n.nrSegments {
  1955  		if n.parent == nil {
  1956  			return Iterator{}
  1957  		}
  1958  		n, i = n.parent, n.parentIndex
  1959  	}
  1960  	return Iterator{n, i}
  1961  }
  1962  
  1963  func zeroValueSlice(slice []Value) {
  1964  	// TODO(jamieliu): check if Go is actually smart enough to optimize a
  1965  	// ClearValue that assigns nil to a memset here.
  1966  	for i := range slice {
  1967  		Functions{}.ClearValue(&slice[i])
  1968  	}
  1969  }
  1970  
  1971  func zeroNodeSlice(slice []*node) {
  1972  	for i := range slice {
  1973  		slice[i] = nil
  1974  	}
  1975  }
  1976  
  1977  // String stringifies a Set for debugging.
  1978  func (s *Set) String() string {
  1979  	return s.root.String()
  1980  }
  1981  
  1982  // String stringifies a node (and all of its children) for debugging.
  1983  func (n *node) String() string {
  1984  	var buf bytes.Buffer
  1985  	n.writeDebugString(&buf, "")
  1986  	return buf.String()
  1987  }
  1988  
  1989  func (n *node) writeDebugString(buf *bytes.Buffer, prefix string) {
  1990  	if n.hasChildren != (n.nrSegments > 0 && n.children[0] != nil) {
  1991  		buf.WriteString(prefix)
  1992  		buf.WriteString(fmt.Sprintf("WARNING: inconsistent value of hasChildren: got %v, want %v\n", n.hasChildren, !n.hasChildren))
  1993  	}
  1994  	for i := 0; i < n.nrSegments; i++ {
  1995  		if child := n.children[i]; child != nil {
  1996  			cprefix := fmt.Sprintf("%s- % 3d ", prefix, i)
  1997  			if child.parent != n || child.parentIndex != i {
  1998  				buf.WriteString(cprefix)
  1999  				buf.WriteString(fmt.Sprintf("WARNING: inconsistent linkage to parent: got (%p, %d), want (%p, %d)\n", child.parent, child.parentIndex, n, i))
  2000  			}
  2001  			child.writeDebugString(buf, fmt.Sprintf("%s- % 3d ", prefix, i))
  2002  		}
  2003  		buf.WriteString(prefix)
  2004  		if n.hasChildren {
  2005  			if trackGaps != 0 {
  2006  				buf.WriteString(fmt.Sprintf("- % 3d: %v => %v, maxGap: %d\n", i, n.keys[i], n.values[i], n.maxGap.Get()))
  2007  			} else {
  2008  				buf.WriteString(fmt.Sprintf("- % 3d: %v => %v\n", i, n.keys[i], n.values[i]))
  2009  			}
  2010  		} else {
  2011  			buf.WriteString(fmt.Sprintf("- % 3d: %v => %v\n", i, n.keys[i], n.values[i]))
  2012  		}
  2013  	}
  2014  	if child := n.children[n.nrSegments]; child != nil {
  2015  		child.writeDebugString(buf, fmt.Sprintf("%s- % 3d ", prefix, n.nrSegments))
  2016  	}
  2017  }
  2018  
  2019  // FlatSegment represents a segment as a single object. FlatSegment is used as
  2020  // an intermediate representation for save/restore and tests.
  2021  //
  2022  // +stateify savable
  2023  type FlatSegment struct {
  2024  	Start Key
  2025  	End   Key
  2026  	Value Value
  2027  }
  2028  
  2029  // ExportSlice returns a copy of all segments in the given set, in ascending
  2030  // key order.
  2031  func (s *Set) ExportSlice() []FlatSegment {
  2032  	var fs []FlatSegment
  2033  	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
  2034  		fs = append(fs, FlatSegment{
  2035  			Start: seg.Start(),
  2036  			End:   seg.End(),
  2037  			Value: seg.Value(),
  2038  		})
  2039  	}
  2040  	return fs
  2041  }
  2042  
  2043  // ImportSlice initializes the given set from the given slice.
  2044  //
  2045  // Preconditions:
  2046  //   - s must be empty.
  2047  //   - fs must represent a valid set (the segments in fs must have valid
  2048  //     lengths that do not overlap).
  2049  //   - The segments in fs must be sorted in ascending key order.
  2050  func (s *Set) ImportSlice(fs []FlatSegment) error {
  2051  	if !s.IsEmpty() {
  2052  		return fmt.Errorf("cannot import into non-empty set %v", s)
  2053  	}
  2054  	gap := s.FirstGap()
  2055  	for i := range fs {
  2056  		f := &fs[i]
  2057  		r := Range{f.Start, f.End}
  2058  		if !gap.Range().IsSupersetOf(r) {
  2059  			return fmt.Errorf("segment overlaps a preceding segment or is incorrectly sorted: %v => %v", r, f.Value)
  2060  		}
  2061  		gap = s.InsertWithoutMerging(gap, r, f.Value).NextGap()
  2062  	}
  2063  	return nil
  2064  }
  2065  
  2066  // segmentTestCheck returns an error if s is incorrectly sorted, does not
  2067  // contain exactly expectedSegments segments, or contains a segment which
  2068  // fails the passed check.
  2069  //
  2070  // This should be used only for testing, and has been added to this package for
  2071  // templating convenience.
  2072  func (s *Set) segmentTestCheck(expectedSegments int, segFunc func(int, Range, Value) error) error {
  2073  	havePrev := false
  2074  	prev := Key(0)
  2075  	nrSegments := 0
  2076  	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
  2077  		next := seg.Start()
  2078  		if havePrev && prev >= next {
  2079  			return fmt.Errorf("incorrect order: key %d (segment %d) >= key %d (segment %d)", prev, nrSegments-1, next, nrSegments)
  2080  		}
  2081  		if segFunc != nil {
  2082  			if err := segFunc(nrSegments, seg.Range(), seg.Value()); err != nil {
  2083  				return err
  2084  			}
  2085  		}
  2086  		prev = next
  2087  		havePrev = true
  2088  		nrSegments++
  2089  	}
  2090  	if nrSegments != expectedSegments {
  2091  		return fmt.Errorf("incorrect number of segments: got %d, wanted %d", nrSegments, expectedSegments)
  2092  	}
  2093  	return nil
  2094  }
  2095  
  2096  // countSegments counts the number of segments in the set.
  2097  //
  2098  // Similar to Check, this should only be used for testing.
  2099  func (s *Set) countSegments() (segments int) {
  2100  	for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() {
  2101  		segments++
  2102  	}
  2103  	return segments
  2104  }