github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/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 interface{} 35 36 // Value is a required type parameter. 37 type Value interface{} 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:".(*SegmentDataSlices)"` 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 // Add inserts the given segment into the set and returns true. If the new 296 // segment can be merged with adjacent segments, Add will do so. If the new 297 // segment would overlap an existing segment, Add returns false. If Add 298 // succeeds, all existing iterators are invalidated. 299 func (s *Set) Add(r Range, val Value) bool { 300 if r.Length() <= 0 { 301 panic(fmt.Sprintf("invalid segment range %v", r)) 302 } 303 gap := s.FindGap(r.Start) 304 if !gap.Ok() { 305 return false 306 } 307 if r.End > gap.End() { 308 return false 309 } 310 s.Insert(gap, r, val) 311 return true 312 } 313 314 // AddWithoutMerging inserts the given segment into the set and returns true. 315 // If it would overlap an existing segment, AddWithoutMerging does nothing and 316 // returns false. If AddWithoutMerging succeeds, all existing iterators are 317 // invalidated. 318 func (s *Set) AddWithoutMerging(r Range, val Value) bool { 319 if r.Length() <= 0 { 320 panic(fmt.Sprintf("invalid segment range %v", r)) 321 } 322 gap := s.FindGap(r.Start) 323 if !gap.Ok() { 324 return false 325 } 326 if r.End > gap.End() { 327 return false 328 } 329 s.InsertWithoutMergingUnchecked(gap, r, val) 330 return true 331 } 332 333 // Insert inserts the given segment into the given gap. If the new segment can 334 // be merged with adjacent segments, Insert will do so. Insert returns an 335 // iterator to the segment containing the inserted value (which may have been 336 // merged with other values). All existing iterators (including gap, but not 337 // including the returned iterator) are invalidated. 338 // 339 // If the gap cannot accommodate the segment, or if r is invalid, Insert panics. 340 // 341 // Insert is semantically equivalent to a InsertWithoutMerging followed by a 342 // Merge, but may be more efficient. Note that there is no unchecked variant of 343 // Insert since Insert must retrieve and inspect gap's predecessor and 344 // successor segments regardless. 345 func (s *Set) Insert(gap GapIterator, r Range, val Value) Iterator { 346 if r.Length() <= 0 { 347 panic(fmt.Sprintf("invalid segment range %v", r)) 348 } 349 prev, next := gap.PrevSegment(), gap.NextSegment() 350 if prev.Ok() && prev.End() > r.Start { 351 panic(fmt.Sprintf("new segment %v overlaps predecessor %v", r, prev.Range())) 352 } 353 if next.Ok() && next.Start() < r.End { 354 panic(fmt.Sprintf("new segment %v overlaps successor %v", r, next.Range())) 355 } 356 if prev.Ok() && prev.End() == r.Start { 357 if mval, ok := (Functions{}).Merge(prev.Range(), prev.Value(), r, val); ok { 358 shrinkMaxGap := trackGaps != 0 && gap.Range().Length() == gap.node.maxGap.Get() 359 prev.SetEndUnchecked(r.End) 360 prev.SetValue(mval) 361 if shrinkMaxGap { 362 gap.node.updateMaxGapLeaf() 363 } 364 if next.Ok() && next.Start() == r.End { 365 val = mval 366 if mval, ok := (Functions{}).Merge(prev.Range(), val, next.Range(), next.Value()); ok { 367 prev.SetEndUnchecked(next.End()) 368 prev.SetValue(mval) 369 return s.Remove(next).PrevSegment() 370 } 371 } 372 return prev 373 } 374 } 375 if next.Ok() && next.Start() == r.End { 376 if mval, ok := (Functions{}).Merge(r, val, next.Range(), next.Value()); ok { 377 shrinkMaxGap := trackGaps != 0 && gap.Range().Length() == gap.node.maxGap.Get() 378 next.SetStartUnchecked(r.Start) 379 next.SetValue(mval) 380 if shrinkMaxGap { 381 gap.node.updateMaxGapLeaf() 382 } 383 return next 384 } 385 } 386 // InsertWithoutMergingUnchecked will maintain maxGap if necessary. 387 return s.InsertWithoutMergingUnchecked(gap, r, val) 388 } 389 390 // InsertWithoutMerging inserts the given segment into the given gap and 391 // returns an iterator to the inserted segment. All existing iterators 392 // (including gap, but not including the returned iterator) are invalidated. 393 // 394 // If the gap cannot accommodate the segment, or if r is invalid, 395 // InsertWithoutMerging panics. 396 func (s *Set) InsertWithoutMerging(gap GapIterator, r Range, val Value) Iterator { 397 if r.Length() <= 0 { 398 panic(fmt.Sprintf("invalid segment range %v", r)) 399 } 400 if gr := gap.Range(); !gr.IsSupersetOf(r) { 401 panic(fmt.Sprintf("cannot insert segment range %v into gap range %v", r, gr)) 402 } 403 return s.InsertWithoutMergingUnchecked(gap, r, val) 404 } 405 406 // InsertWithoutMergingUnchecked inserts the given segment into the given gap 407 // and returns an iterator to the inserted segment. All existing iterators 408 // (including gap, but not including the returned iterator) are invalidated. 409 // 410 // Preconditions: 411 // * r.Start >= gap.Start(). 412 // * r.End <= gap.End(). 413 func (s *Set) InsertWithoutMergingUnchecked(gap GapIterator, r Range, val Value) Iterator { 414 gap = gap.node.rebalanceBeforeInsert(gap) 415 splitMaxGap := trackGaps != 0 && (gap.node.nrSegments == 0 || gap.Range().Length() == gap.node.maxGap.Get()) 416 copy(gap.node.keys[gap.index+1:], gap.node.keys[gap.index:gap.node.nrSegments]) 417 copy(gap.node.values[gap.index+1:], gap.node.values[gap.index:gap.node.nrSegments]) 418 gap.node.keys[gap.index] = r 419 gap.node.values[gap.index] = val 420 gap.node.nrSegments++ 421 if splitMaxGap { 422 gap.node.updateMaxGapLeaf() 423 } 424 return Iterator{gap.node, gap.index} 425 } 426 427 // Remove removes the given segment and returns an iterator to the vacated gap. 428 // All existing iterators (including seg, but not including the returned 429 // iterator) are invalidated. 430 func (s *Set) Remove(seg Iterator) GapIterator { 431 // We only want to remove directly from a leaf node. 432 if seg.node.hasChildren { 433 // Since seg.node has children, the removed segment must have a 434 // predecessor (at the end of the rightmost leaf of its left child 435 // subtree). Move the contents of that predecessor into the removed 436 // segment's position, and remove that predecessor instead. (We choose 437 // to steal the predecessor rather than the successor because removing 438 // from the end of a leaf node doesn't involve any copying unless 439 // merging is required.) 440 victim := seg.PrevSegment() 441 // This must be unchecked since until victim is removed, seg and victim 442 // overlap. 443 seg.SetRangeUnchecked(victim.Range()) 444 seg.SetValue(victim.Value()) 445 // Need to update the nextAdjacentNode's maxGap because the gap in between 446 // must have been modified by updating seg.Range() to victim.Range(). 447 // seg.NextSegment() must exist since the last segment can't be in a 448 // non-leaf node. 449 nextAdjacentNode := seg.NextSegment().node 450 if trackGaps != 0 { 451 nextAdjacentNode.updateMaxGapLeaf() 452 } 453 return s.Remove(victim).NextGap() 454 } 455 copy(seg.node.keys[seg.index:], seg.node.keys[seg.index+1:seg.node.nrSegments]) 456 copy(seg.node.values[seg.index:], seg.node.values[seg.index+1:seg.node.nrSegments]) 457 Functions{}.ClearValue(&seg.node.values[seg.node.nrSegments-1]) 458 seg.node.nrSegments-- 459 if trackGaps != 0 { 460 seg.node.updateMaxGapLeaf() 461 } 462 return seg.node.rebalanceAfterRemove(GapIterator{seg.node, seg.index}) 463 } 464 465 // RemoveAll removes all segments from the set. All existing iterators are 466 // invalidated. 467 func (s *Set) RemoveAll() { 468 s.root = node{} 469 } 470 471 // RemoveRange removes all segments in the given range. An iterator to the 472 // newly formed gap is returned, and all existing iterators are invalidated. 473 func (s *Set) RemoveRange(r Range) GapIterator { 474 seg, gap := s.Find(r.Start) 475 if seg.Ok() { 476 seg = s.Isolate(seg, r) 477 gap = s.Remove(seg) 478 } 479 for seg = gap.NextSegment(); seg.Ok() && seg.Start() < r.End; seg = gap.NextSegment() { 480 seg = s.Isolate(seg, r) 481 gap = s.Remove(seg) 482 } 483 return gap 484 } 485 486 // Merge attempts to merge two neighboring segments. If successful, Merge 487 // returns an iterator to the merged segment, and all existing iterators are 488 // invalidated. Otherwise, Merge returns a terminal iterator. 489 // 490 // If first is not the predecessor of second, Merge panics. 491 func (s *Set) Merge(first, second Iterator) Iterator { 492 if first.NextSegment() != second { 493 panic(fmt.Sprintf("attempt to merge non-neighboring segments %v, %v", first.Range(), second.Range())) 494 } 495 return s.MergeUnchecked(first, second) 496 } 497 498 // MergeUnchecked attempts to merge two neighboring segments. If successful, 499 // MergeUnchecked returns an iterator to the merged segment, and all existing 500 // iterators are invalidated. Otherwise, MergeUnchecked returns a terminal 501 // iterator. 502 // 503 // Precondition: first is the predecessor of second: first.NextSegment() == 504 // second, first == second.PrevSegment(). 505 func (s *Set) MergeUnchecked(first, second Iterator) Iterator { 506 if first.End() == second.Start() { 507 if mval, ok := (Functions{}).Merge(first.Range(), first.Value(), second.Range(), second.Value()); ok { 508 // N.B. This must be unchecked because until s.Remove(second), first 509 // overlaps second. 510 first.SetEndUnchecked(second.End()) 511 first.SetValue(mval) 512 // Remove will handle the maxGap update if necessary. 513 return s.Remove(second).PrevSegment() 514 } 515 } 516 return Iterator{} 517 } 518 519 // MergeAll attempts to merge all adjacent segments in the set. All existing 520 // iterators are invalidated. 521 func (s *Set) MergeAll() { 522 seg := s.FirstSegment() 523 if !seg.Ok() { 524 return 525 } 526 next := seg.NextSegment() 527 for next.Ok() { 528 if mseg := s.MergeUnchecked(seg, next); mseg.Ok() { 529 seg, next = mseg, mseg.NextSegment() 530 } else { 531 seg, next = next, next.NextSegment() 532 } 533 } 534 } 535 536 // MergeRange attempts to merge all adjacent segments that contain a key in the 537 // specific range. All existing iterators are invalidated. 538 func (s *Set) MergeRange(r Range) { 539 seg := s.LowerBoundSegment(r.Start) 540 if !seg.Ok() { 541 return 542 } 543 next := seg.NextSegment() 544 for next.Ok() && next.Range().Start < r.End { 545 if mseg := s.MergeUnchecked(seg, next); mseg.Ok() { 546 seg, next = mseg, mseg.NextSegment() 547 } else { 548 seg, next = next, next.NextSegment() 549 } 550 } 551 } 552 553 // MergeAdjacent attempts to merge the segment containing r.Start with its 554 // predecessor, and the segment containing r.End-1 with its successor. 555 func (s *Set) MergeAdjacent(r Range) { 556 first := s.FindSegment(r.Start) 557 if first.Ok() { 558 if prev := first.PrevSegment(); prev.Ok() { 559 s.Merge(prev, first) 560 } 561 } 562 last := s.FindSegment(r.End - 1) 563 if last.Ok() { 564 if next := last.NextSegment(); next.Ok() { 565 s.Merge(last, next) 566 } 567 } 568 } 569 570 // Split splits the given segment at the given key and returns iterators to the 571 // two resulting segments. All existing iterators (including seg, but not 572 // including the returned iterators) are invalidated. 573 // 574 // If the segment cannot be split at split (because split is at the start or 575 // end of the segment's range, so splitting would produce a segment with zero 576 // length, or because split falls outside the segment's range altogether), 577 // Split panics. 578 func (s *Set) Split(seg Iterator, split Key) (Iterator, Iterator) { 579 if !seg.Range().CanSplitAt(split) { 580 panic(fmt.Sprintf("can't split %v at %v", seg.Range(), split)) 581 } 582 return s.SplitUnchecked(seg, split) 583 } 584 585 // SplitUnchecked splits the given segment at the given key and returns 586 // iterators to the two resulting segments. All existing iterators (including 587 // seg, but not including the returned iterators) are invalidated. 588 // 589 // Preconditions: seg.Start() < key < seg.End(). 590 func (s *Set) SplitUnchecked(seg Iterator, split Key) (Iterator, Iterator) { 591 val1, val2 := (Functions{}).Split(seg.Range(), seg.Value(), split) 592 end2 := seg.End() 593 seg.SetEndUnchecked(split) 594 seg.SetValue(val1) 595 seg2 := s.InsertWithoutMergingUnchecked(seg.NextGap(), Range{split, end2}, val2) 596 // seg may now be invalid due to the Insert. 597 return seg2.PrevSegment(), seg2 598 } 599 600 // SplitAt splits the segment straddling split, if one exists. SplitAt returns 601 // true if a segment was split and false otherwise. If SplitAt splits a 602 // segment, all existing iterators are invalidated. 603 func (s *Set) SplitAt(split Key) bool { 604 if seg := s.FindSegment(split); seg.Ok() && seg.Range().CanSplitAt(split) { 605 s.SplitUnchecked(seg, split) 606 return true 607 } 608 return false 609 } 610 611 // Isolate ensures that the given segment's range does not escape r by 612 // splitting at r.Start and r.End if necessary, and returns an updated iterator 613 // to the bounded segment. All existing iterators (including seg, but not 614 // including the returned iterators) are invalidated. 615 func (s *Set) Isolate(seg Iterator, r Range) Iterator { 616 if seg.Range().CanSplitAt(r.Start) { 617 _, seg = s.SplitUnchecked(seg, r.Start) 618 } 619 if seg.Range().CanSplitAt(r.End) { 620 seg, _ = s.SplitUnchecked(seg, r.End) 621 } 622 return seg 623 } 624 625 // ApplyContiguous applies a function to a contiguous range of segments, 626 // splitting if necessary. The function is applied until the first gap is 627 // encountered, at which point the gap is returned. If the function is applied 628 // across the entire range, a terminal gap is returned. All existing iterators 629 // are invalidated. 630 // 631 // N.B. The Iterator must not be invalidated by the function. 632 func (s *Set) ApplyContiguous(r Range, fn func(seg Iterator)) GapIterator { 633 seg, gap := s.Find(r.Start) 634 if !seg.Ok() { 635 return gap 636 } 637 for { 638 seg = s.Isolate(seg, r) 639 fn(seg) 640 if seg.End() >= r.End { 641 return GapIterator{} 642 } 643 gap = seg.NextGap() 644 if !gap.IsEmpty() { 645 return gap 646 } 647 seg = gap.NextSegment() 648 if !seg.Ok() { 649 // This implies that the last segment extended all the 650 // way to the maximum value, since the gap was empty. 651 return GapIterator{} 652 } 653 } 654 } 655 656 // +stateify savable 657 type node struct { 658 // An internal binary tree node looks like: 659 // 660 // K 661 // / \ 662 // Cl Cr 663 // 664 // where all keys in the subtree rooted by Cl (the left subtree) are less 665 // than K (the key of the parent node), and all keys in the subtree rooted 666 // by Cr (the right subtree) are greater than K. 667 // 668 // An internal B-tree node's indexes work out to look like: 669 // 670 // K0 K1 K2 ... Kn-1 671 // / \/ \/ \ ... / \ 672 // C0 C1 C2 C3 ... Cn-1 Cn 673 // 674 // where n is nrSegments. 675 nrSegments int 676 677 // parent is a pointer to this node's parent. If this node is root, parent 678 // is nil. 679 parent *node 680 681 // parentIndex is the index of this node in parent.children. 682 parentIndex int 683 684 // Flag for internal nodes that is technically redundant with "children[0] 685 // != nil", but is stored in the first cache line. "hasChildren" rather 686 // than "isLeaf" because false must be the correct value for an empty root. 687 hasChildren bool 688 689 // The longest gap within this node. If the node is a leaf, it's simply the 690 // maximum gap among all the (nrSegments+1) gaps formed by its nrSegments keys 691 // including the 0th and nrSegments-th gap possibly shared with its upper-level 692 // nodes; if it's a non-leaf node, it's the max of all children's maxGap. 693 maxGap dynamicGap 694 695 // Nodes store keys and values in separate arrays to maximize locality in 696 // the common case (scanning keys for lookup). 697 keys [maxDegree - 1]Range 698 values [maxDegree - 1]Value 699 children [maxDegree]*node 700 } 701 702 // firstSegment returns the first segment in the subtree rooted by n. 703 // 704 // Preconditions: n.nrSegments != 0. 705 func (n *node) firstSegment() Iterator { 706 for n.hasChildren { 707 n = n.children[0] 708 } 709 return Iterator{n, 0} 710 } 711 712 // lastSegment returns the last segment in the subtree rooted by n. 713 // 714 // Preconditions: n.nrSegments != 0. 715 func (n *node) lastSegment() Iterator { 716 for n.hasChildren { 717 n = n.children[n.nrSegments] 718 } 719 return Iterator{n, n.nrSegments - 1} 720 } 721 722 func (n *node) prevSibling() *node { 723 if n.parent == nil || n.parentIndex == 0 { 724 return nil 725 } 726 return n.parent.children[n.parentIndex-1] 727 } 728 729 func (n *node) nextSibling() *node { 730 if n.parent == nil || n.parentIndex == n.parent.nrSegments { 731 return nil 732 } 733 return n.parent.children[n.parentIndex+1] 734 } 735 736 // rebalanceBeforeInsert splits n and its ancestors if they are full, as 737 // required for insertion, and returns an updated iterator to the position 738 // represented by gap. 739 func (n *node) rebalanceBeforeInsert(gap GapIterator) GapIterator { 740 if n.nrSegments < maxDegree-1 { 741 return gap 742 } 743 if n.parent != nil { 744 gap = n.parent.rebalanceBeforeInsert(gap) 745 } 746 if n.parent == nil { 747 // n is root. Move all segments before and after n's median segment 748 // into new child nodes adjacent to the median segment, which is now 749 // the only segment in root. 750 left := &node{ 751 nrSegments: minDegree - 1, 752 parent: n, 753 parentIndex: 0, 754 hasChildren: n.hasChildren, 755 } 756 right := &node{ 757 nrSegments: minDegree - 1, 758 parent: n, 759 parentIndex: 1, 760 hasChildren: n.hasChildren, 761 } 762 copy(left.keys[:minDegree-1], n.keys[:minDegree-1]) 763 copy(left.values[:minDegree-1], n.values[:minDegree-1]) 764 copy(right.keys[:minDegree-1], n.keys[minDegree:]) 765 copy(right.values[:minDegree-1], n.values[minDegree:]) 766 n.keys[0], n.values[0] = n.keys[minDegree-1], n.values[minDegree-1] 767 zeroValueSlice(n.values[1:]) 768 if n.hasChildren { 769 copy(left.children[:minDegree], n.children[:minDegree]) 770 copy(right.children[:minDegree], n.children[minDegree:]) 771 zeroNodeSlice(n.children[2:]) 772 for i := 0; i < minDegree; i++ { 773 left.children[i].parent = left 774 left.children[i].parentIndex = i 775 right.children[i].parent = right 776 right.children[i].parentIndex = i 777 } 778 } 779 n.nrSegments = 1 780 n.hasChildren = true 781 n.children[0] = left 782 n.children[1] = right 783 // In this case, n's maxGap won't violated as it's still the root, 784 // but the left and right children should be updated locally as they 785 // are newly split from n. 786 if trackGaps != 0 { 787 left.updateMaxGapLocal() 788 right.updateMaxGapLocal() 789 } 790 if gap.node != n { 791 return gap 792 } 793 if gap.index < minDegree { 794 return GapIterator{left, gap.index} 795 } 796 return GapIterator{right, gap.index - minDegree} 797 } 798 // n is non-root. Move n's median segment into its parent node (which can't 799 // be full because we've already invoked n.parent.rebalanceBeforeInsert) 800 // and move all segments after n's median into a new sibling node (the 801 // median segment's right child subtree). 802 copy(n.parent.keys[n.parentIndex+1:], n.parent.keys[n.parentIndex:n.parent.nrSegments]) 803 copy(n.parent.values[n.parentIndex+1:], n.parent.values[n.parentIndex:n.parent.nrSegments]) 804 n.parent.keys[n.parentIndex], n.parent.values[n.parentIndex] = n.keys[minDegree-1], n.values[minDegree-1] 805 copy(n.parent.children[n.parentIndex+2:], n.parent.children[n.parentIndex+1:n.parent.nrSegments+1]) 806 for i := n.parentIndex + 2; i < n.parent.nrSegments+2; i++ { 807 n.parent.children[i].parentIndex = i 808 } 809 sibling := &node{ 810 nrSegments: minDegree - 1, 811 parent: n.parent, 812 parentIndex: n.parentIndex + 1, 813 hasChildren: n.hasChildren, 814 } 815 n.parent.children[n.parentIndex+1] = sibling 816 n.parent.nrSegments++ 817 copy(sibling.keys[:minDegree-1], n.keys[minDegree:]) 818 copy(sibling.values[:minDegree-1], n.values[minDegree:]) 819 zeroValueSlice(n.values[minDegree-1:]) 820 if n.hasChildren { 821 copy(sibling.children[:minDegree], n.children[minDegree:]) 822 zeroNodeSlice(n.children[minDegree:]) 823 for i := 0; i < minDegree; i++ { 824 sibling.children[i].parent = sibling 825 sibling.children[i].parentIndex = i 826 } 827 } 828 n.nrSegments = minDegree - 1 829 // MaxGap of n's parent is not violated because the segments within is not changed. 830 // n and its sibling's maxGap need to be updated locally as they are two new nodes split from old n. 831 if trackGaps != 0 { 832 n.updateMaxGapLocal() 833 sibling.updateMaxGapLocal() 834 } 835 // gap.node can't be n.parent because gaps are always in leaf nodes. 836 if gap.node != n { 837 return gap 838 } 839 if gap.index < minDegree { 840 return gap 841 } 842 return GapIterator{sibling, gap.index - minDegree} 843 } 844 845 // rebalanceAfterRemove "unsplits" n and its ancestors if they are deficient 846 // (contain fewer segments than required by B-tree invariants), as required for 847 // removal, and returns an updated iterator to the position represented by gap. 848 // 849 // Precondition: n is the only node in the tree that may currently violate a 850 // B-tree invariant. 851 func (n *node) rebalanceAfterRemove(gap GapIterator) GapIterator { 852 for { 853 if n.nrSegments >= minDegree-1 { 854 return gap 855 } 856 if n.parent == nil { 857 // Root is allowed to be deficient. 858 return gap 859 } 860 // There's one other thing we can do before resorting to unsplitting. 861 // If either sibling node has at least minDegree segments, rotate that 862 // sibling's closest segment through the segment in the parent that 863 // separates us. That is, given: 864 // 865 // ... D ... 866 // / \ 867 // ... B C] [E ... 868 // 869 // where the node containing E is deficient, end up with: 870 // 871 // ... C ... 872 // / \ 873 // ... B] [D E ... 874 // 875 // As in Set.Remove, prefer rotating from the end of the sibling to the 876 // left: by precondition, n.node has fewer segments (to memcpy) than 877 // the sibling does. 878 if sibling := n.prevSibling(); sibling != nil && sibling.nrSegments >= minDegree { 879 copy(n.keys[1:], n.keys[:n.nrSegments]) 880 copy(n.values[1:], n.values[:n.nrSegments]) 881 n.keys[0] = n.parent.keys[n.parentIndex-1] 882 n.values[0] = n.parent.values[n.parentIndex-1] 883 n.parent.keys[n.parentIndex-1] = sibling.keys[sibling.nrSegments-1] 884 n.parent.values[n.parentIndex-1] = sibling.values[sibling.nrSegments-1] 885 Functions{}.ClearValue(&sibling.values[sibling.nrSegments-1]) 886 if n.hasChildren { 887 copy(n.children[1:], n.children[:n.nrSegments+1]) 888 n.children[0] = sibling.children[sibling.nrSegments] 889 sibling.children[sibling.nrSegments] = nil 890 n.children[0].parent = n 891 n.children[0].parentIndex = 0 892 for i := 1; i < n.nrSegments+2; i++ { 893 n.children[i].parentIndex = i 894 } 895 } 896 n.nrSegments++ 897 sibling.nrSegments-- 898 // n's parent's maxGap does not need to be updated as its content is unmodified. 899 // n and its sibling must be updated with (new) maxGap because of the shift of keys. 900 if trackGaps != 0 { 901 n.updateMaxGapLocal() 902 sibling.updateMaxGapLocal() 903 } 904 if gap.node == sibling && gap.index == sibling.nrSegments { 905 return GapIterator{n, 0} 906 } 907 if gap.node == n { 908 return GapIterator{n, gap.index + 1} 909 } 910 return gap 911 } 912 if sibling := n.nextSibling(); sibling != nil && sibling.nrSegments >= minDegree { 913 n.keys[n.nrSegments] = n.parent.keys[n.parentIndex] 914 n.values[n.nrSegments] = n.parent.values[n.parentIndex] 915 n.parent.keys[n.parentIndex] = sibling.keys[0] 916 n.parent.values[n.parentIndex] = sibling.values[0] 917 copy(sibling.keys[:sibling.nrSegments-1], sibling.keys[1:]) 918 copy(sibling.values[:sibling.nrSegments-1], sibling.values[1:]) 919 Functions{}.ClearValue(&sibling.values[sibling.nrSegments-1]) 920 if n.hasChildren { 921 n.children[n.nrSegments+1] = sibling.children[0] 922 copy(sibling.children[:sibling.nrSegments], sibling.children[1:]) 923 sibling.children[sibling.nrSegments] = nil 924 n.children[n.nrSegments+1].parent = n 925 n.children[n.nrSegments+1].parentIndex = n.nrSegments + 1 926 for i := 0; i < sibling.nrSegments; i++ { 927 sibling.children[i].parentIndex = i 928 } 929 } 930 n.nrSegments++ 931 sibling.nrSegments-- 932 // n's parent's maxGap does not need to be updated as its content is unmodified. 933 // n and its sibling must be updated with (new) maxGap because of the shift of keys. 934 if trackGaps != 0 { 935 n.updateMaxGapLocal() 936 sibling.updateMaxGapLocal() 937 } 938 if gap.node == sibling { 939 if gap.index == 0 { 940 return GapIterator{n, n.nrSegments} 941 } 942 return GapIterator{sibling, gap.index - 1} 943 } 944 return gap 945 } 946 // Otherwise, we must unsplit. 947 p := n.parent 948 if p.nrSegments == 1 { 949 // Merge all segments in both n and its sibling back into n.parent. 950 // This is the reverse of the root splitting case in 951 // node.rebalanceBeforeInsert. (Because we require minDegree >= 3, 952 // only root can have 1 segment in this path, so this reduces the 953 // height of the tree by 1, without violating the constraint that 954 // all leaf nodes remain at the same depth.) 955 left, right := p.children[0], p.children[1] 956 p.nrSegments = left.nrSegments + right.nrSegments + 1 957 p.hasChildren = left.hasChildren 958 p.keys[left.nrSegments] = p.keys[0] 959 p.values[left.nrSegments] = p.values[0] 960 copy(p.keys[:left.nrSegments], left.keys[:left.nrSegments]) 961 copy(p.values[:left.nrSegments], left.values[:left.nrSegments]) 962 copy(p.keys[left.nrSegments+1:], right.keys[:right.nrSegments]) 963 copy(p.values[left.nrSegments+1:], right.values[:right.nrSegments]) 964 if left.hasChildren { 965 copy(p.children[:left.nrSegments+1], left.children[:left.nrSegments+1]) 966 copy(p.children[left.nrSegments+1:], right.children[:right.nrSegments+1]) 967 for i := 0; i < p.nrSegments+1; i++ { 968 p.children[i].parent = p 969 p.children[i].parentIndex = i 970 } 971 } else { 972 p.children[0] = nil 973 p.children[1] = nil 974 } 975 // No need to update maxGap of p as its content is not changed. 976 if gap.node == left { 977 return GapIterator{p, gap.index} 978 } 979 if gap.node == right { 980 return GapIterator{p, gap.index + left.nrSegments + 1} 981 } 982 return gap 983 } 984 // Merge n and either sibling, along with the segment separating the 985 // two, into whichever of the two nodes comes first. This is the 986 // reverse of the non-root splitting case in 987 // node.rebalanceBeforeInsert. 988 var left, right *node 989 if n.parentIndex > 0 { 990 left = n.prevSibling() 991 right = n 992 } else { 993 left = n 994 right = n.nextSibling() 995 } 996 // Fix up gap first since we need the old left.nrSegments, which 997 // merging will change. 998 if gap.node == right { 999 gap = GapIterator{left, gap.index + left.nrSegments + 1} 1000 } 1001 left.keys[left.nrSegments] = p.keys[left.parentIndex] 1002 left.values[left.nrSegments] = p.values[left.parentIndex] 1003 copy(left.keys[left.nrSegments+1:], right.keys[:right.nrSegments]) 1004 copy(left.values[left.nrSegments+1:], right.values[:right.nrSegments]) 1005 if left.hasChildren { 1006 copy(left.children[left.nrSegments+1:], right.children[:right.nrSegments+1]) 1007 for i := left.nrSegments + 1; i < left.nrSegments+right.nrSegments+2; i++ { 1008 left.children[i].parent = left 1009 left.children[i].parentIndex = i 1010 } 1011 } 1012 left.nrSegments += right.nrSegments + 1 1013 copy(p.keys[left.parentIndex:], p.keys[left.parentIndex+1:p.nrSegments]) 1014 copy(p.values[left.parentIndex:], p.values[left.parentIndex+1:p.nrSegments]) 1015 Functions{}.ClearValue(&p.values[p.nrSegments-1]) 1016 copy(p.children[left.parentIndex+1:], p.children[left.parentIndex+2:p.nrSegments+1]) 1017 for i := 0; i < p.nrSegments; i++ { 1018 p.children[i].parentIndex = i 1019 } 1020 p.children[p.nrSegments] = nil 1021 p.nrSegments-- 1022 // Update maxGap of left locally, no need to change p and right because 1023 // p's contents is not changed and right is already invalid. 1024 if trackGaps != 0 { 1025 left.updateMaxGapLocal() 1026 } 1027 // This process robs p of one segment, so recurse into rebalancing p. 1028 n = p 1029 } 1030 } 1031 1032 // updateMaxGapLeaf updates maxGap bottom-up from the calling leaf until no 1033 // necessary update. 1034 // 1035 // Preconditions: n must be a leaf node, trackGaps must be 1. 1036 func (n *node) updateMaxGapLeaf() { 1037 if n.hasChildren { 1038 panic(fmt.Sprintf("updateMaxGapLeaf should always be called on leaf node: %v", n)) 1039 } 1040 max := n.calculateMaxGapLeaf() 1041 if max == n.maxGap.Get() { 1042 // If new max equals the old maxGap, no update is needed. 1043 return 1044 } 1045 oldMax := n.maxGap.Get() 1046 n.maxGap.Set(max) 1047 if max > oldMax { 1048 // Grow ancestor maxGaps. 1049 for p := n.parent; p != nil; p = p.parent { 1050 if p.maxGap.Get() >= max { 1051 // p and its ancestors already contain an equal or larger gap. 1052 break 1053 } 1054 // Only if new maxGap is larger than parent's 1055 // old maxGap, propagate this update to parent. 1056 p.maxGap.Set(max) 1057 } 1058 return 1059 } 1060 // Shrink ancestor maxGaps. 1061 for p := n.parent; p != nil; p = p.parent { 1062 if p.maxGap.Get() > oldMax { 1063 // p and its ancestors still contain a larger gap. 1064 break 1065 } 1066 // If new max is smaller than the old maxGap, and this gap used 1067 // to be the maxGap of its parent, iterate parent's children 1068 // and calculate parent's new maxGap.(It's probable that parent 1069 // has two children with the old maxGap, but we need to check it anyway.) 1070 parentNewMax := p.calculateMaxGapInternal() 1071 if p.maxGap.Get() == parentNewMax { 1072 // p and its ancestors still contain a gap of at least equal size. 1073 break 1074 } 1075 // If p's new maxGap differs from the old one, propagate this update. 1076 p.maxGap.Set(parentNewMax) 1077 } 1078 } 1079 1080 // updateMaxGapLocal updates maxGap of the calling node solely with no 1081 // propagation to ancestor nodes. 1082 // 1083 // Precondition: trackGaps must be 1. 1084 func (n *node) updateMaxGapLocal() { 1085 if !n.hasChildren { 1086 // Leaf node iterates its gaps. 1087 n.maxGap.Set(n.calculateMaxGapLeaf()) 1088 } else { 1089 // Non-leaf node iterates its children. 1090 n.maxGap.Set(n.calculateMaxGapInternal()) 1091 } 1092 } 1093 1094 // calculateMaxGapLeaf iterates the gaps within a leaf node and calculate the 1095 // max. 1096 // 1097 // Preconditions: n must be a leaf node. 1098 func (n *node) calculateMaxGapLeaf() Key { 1099 max := GapIterator{n, 0}.Range().Length() 1100 for i := 1; i <= n.nrSegments; i++ { 1101 if current := (GapIterator{n, i}).Range().Length(); current > max { 1102 max = current 1103 } 1104 } 1105 return max 1106 } 1107 1108 // calculateMaxGapInternal iterates children's maxGap within an internal node n 1109 // and calculate the max. 1110 // 1111 // Preconditions: n must be a non-leaf node. 1112 func (n *node) calculateMaxGapInternal() Key { 1113 max := n.children[0].maxGap.Get() 1114 for i := 1; i <= n.nrSegments; i++ { 1115 if current := n.children[i].maxGap.Get(); current > max { 1116 max = current 1117 } 1118 } 1119 return max 1120 } 1121 1122 // searchFirstLargeEnoughGap returns the first gap having at least minSize length 1123 // in the subtree rooted by n. If not found, return a terminal gap iterator. 1124 func (n *node) searchFirstLargeEnoughGap(minSize Key) GapIterator { 1125 if n.maxGap.Get() < minSize { 1126 return GapIterator{} 1127 } 1128 if n.hasChildren { 1129 for i := 0; i <= n.nrSegments; i++ { 1130 if largeEnoughGap := n.children[i].searchFirstLargeEnoughGap(minSize); largeEnoughGap.Ok() { 1131 return largeEnoughGap 1132 } 1133 } 1134 } else { 1135 for i := 0; i <= n.nrSegments; i++ { 1136 currentGap := GapIterator{n, i} 1137 if currentGap.Range().Length() >= minSize { 1138 return currentGap 1139 } 1140 } 1141 } 1142 panic(fmt.Sprintf("invalid maxGap in %v", n)) 1143 } 1144 1145 // searchLastLargeEnoughGap returns the last gap having at least minSize length 1146 // in the subtree rooted by n. If not found, return a terminal gap iterator. 1147 func (n *node) searchLastLargeEnoughGap(minSize Key) GapIterator { 1148 if n.maxGap.Get() < minSize { 1149 return GapIterator{} 1150 } 1151 if n.hasChildren { 1152 for i := n.nrSegments; i >= 0; i-- { 1153 if largeEnoughGap := n.children[i].searchLastLargeEnoughGap(minSize); largeEnoughGap.Ok() { 1154 return largeEnoughGap 1155 } 1156 } 1157 } else { 1158 for i := n.nrSegments; i >= 0; i-- { 1159 currentGap := GapIterator{n, i} 1160 if currentGap.Range().Length() >= minSize { 1161 return currentGap 1162 } 1163 } 1164 } 1165 panic(fmt.Sprintf("invalid maxGap in %v", n)) 1166 } 1167 1168 // A Iterator is conceptually one of: 1169 // 1170 // - A pointer to a segment in a set; or 1171 // 1172 // - A terminal iterator, which is a sentinel indicating that the end of 1173 // iteration has been reached. 1174 // 1175 // Iterators are copyable values and are meaningfully equality-comparable. The 1176 // zero value of Iterator is a terminal iterator. 1177 // 1178 // Unless otherwise specified, any mutation of a set invalidates all existing 1179 // iterators into the set. 1180 type Iterator struct { 1181 // node is the node containing the iterated segment. If the iterator is 1182 // terminal, node is nil. 1183 node *node 1184 1185 // index is the index of the segment in node.keys/values. 1186 index int 1187 } 1188 1189 // Ok returns true if the iterator is not terminal. All other methods are only 1190 // valid for non-terminal iterators. 1191 func (seg Iterator) Ok() bool { 1192 return seg.node != nil 1193 } 1194 1195 // Range returns the iterated segment's range key. 1196 func (seg Iterator) Range() Range { 1197 return seg.node.keys[seg.index] 1198 } 1199 1200 // Start is equivalent to Range().Start, but should be preferred if only the 1201 // start of the range is needed. 1202 func (seg Iterator) Start() Key { 1203 return seg.node.keys[seg.index].Start 1204 } 1205 1206 // End is equivalent to Range().End, but should be preferred if only the end of 1207 // the range is needed. 1208 func (seg Iterator) End() Key { 1209 return seg.node.keys[seg.index].End 1210 } 1211 1212 // SetRangeUnchecked mutates the iterated segment's range key. This operation 1213 // does not invalidate any iterators. 1214 // 1215 // Preconditions: 1216 // * r.Length() > 0. 1217 // * The new range must not overlap an existing one: 1218 // * If seg.NextSegment().Ok(), then r.end <= seg.NextSegment().Start(). 1219 // * If seg.PrevSegment().Ok(), then r.start >= seg.PrevSegment().End(). 1220 func (seg Iterator) SetRangeUnchecked(r Range) { 1221 seg.node.keys[seg.index] = r 1222 } 1223 1224 // SetRange mutates the iterated segment's range key. If the new range would 1225 // cause the iterated segment to overlap another segment, or if the new range 1226 // is invalid, SetRange panics. This operation does not invalidate any 1227 // iterators. 1228 func (seg Iterator) SetRange(r Range) { 1229 if r.Length() <= 0 { 1230 panic(fmt.Sprintf("invalid segment range %v", r)) 1231 } 1232 if prev := seg.PrevSegment(); prev.Ok() && r.Start < prev.End() { 1233 panic(fmt.Sprintf("new segment range %v overlaps segment range %v", r, prev.Range())) 1234 } 1235 if next := seg.NextSegment(); next.Ok() && r.End > next.Start() { 1236 panic(fmt.Sprintf("new segment range %v overlaps segment range %v", r, next.Range())) 1237 } 1238 seg.SetRangeUnchecked(r) 1239 } 1240 1241 // SetStartUnchecked mutates the iterated segment's start. This operation does 1242 // not invalidate any iterators. 1243 // 1244 // Preconditions: The new start must be valid: 1245 // * start < seg.End() 1246 // * If seg.PrevSegment().Ok(), then start >= seg.PrevSegment().End(). 1247 func (seg Iterator) SetStartUnchecked(start Key) { 1248 seg.node.keys[seg.index].Start = start 1249 } 1250 1251 // SetStart mutates the iterated segment's start. If the new start value would 1252 // cause the iterated segment to overlap another segment, or would result in an 1253 // invalid range, SetStart panics. This operation does not invalidate any 1254 // iterators. 1255 func (seg Iterator) SetStart(start Key) { 1256 if start >= seg.End() { 1257 panic(fmt.Sprintf("new start %v would invalidate segment range %v", start, seg.Range())) 1258 } 1259 if prev := seg.PrevSegment(); prev.Ok() && start < prev.End() { 1260 panic(fmt.Sprintf("new start %v would cause segment range %v to overlap segment range %v", start, seg.Range(), prev.Range())) 1261 } 1262 seg.SetStartUnchecked(start) 1263 } 1264 1265 // SetEndUnchecked mutates the iterated segment's end. This operation does not 1266 // invalidate any iterators. 1267 // 1268 // Preconditions: The new end must be valid: 1269 // * end > seg.Start(). 1270 // * If seg.NextSegment().Ok(), then end <= seg.NextSegment().Start(). 1271 func (seg Iterator) SetEndUnchecked(end Key) { 1272 seg.node.keys[seg.index].End = end 1273 } 1274 1275 // SetEnd mutates the iterated segment's end. If the new end value would cause 1276 // the iterated segment to overlap another segment, or would result in an 1277 // invalid range, SetEnd panics. This operation does not invalidate any 1278 // iterators. 1279 func (seg Iterator) SetEnd(end Key) { 1280 if end <= seg.Start() { 1281 panic(fmt.Sprintf("new end %v would invalidate segment range %v", end, seg.Range())) 1282 } 1283 if next := seg.NextSegment(); next.Ok() && end > next.Start() { 1284 panic(fmt.Sprintf("new end %v would cause segment range %v to overlap segment range %v", end, seg.Range(), next.Range())) 1285 } 1286 seg.SetEndUnchecked(end) 1287 } 1288 1289 // Value returns a copy of the iterated segment's value. 1290 func (seg Iterator) Value() Value { 1291 return seg.node.values[seg.index] 1292 } 1293 1294 // ValuePtr returns a pointer to the iterated segment's value. The pointer is 1295 // invalidated if the iterator is invalidated. This operation does not 1296 // invalidate any iterators. 1297 func (seg Iterator) ValuePtr() *Value { 1298 return &seg.node.values[seg.index] 1299 } 1300 1301 // SetValue mutates the iterated segment's value. This operation does not 1302 // invalidate any iterators. 1303 func (seg Iterator) SetValue(val Value) { 1304 seg.node.values[seg.index] = val 1305 } 1306 1307 // PrevSegment returns the iterated segment's predecessor. If there is no 1308 // preceding segment, PrevSegment returns a terminal iterator. 1309 func (seg Iterator) PrevSegment() Iterator { 1310 if seg.node.hasChildren { 1311 return seg.node.children[seg.index].lastSegment() 1312 } 1313 if seg.index > 0 { 1314 return Iterator{seg.node, seg.index - 1} 1315 } 1316 if seg.node.parent == nil { 1317 return Iterator{} 1318 } 1319 return segmentBeforePosition(seg.node.parent, seg.node.parentIndex) 1320 } 1321 1322 // NextSegment returns the iterated segment's successor. If there is no 1323 // succeeding segment, NextSegment returns a terminal iterator. 1324 func (seg Iterator) NextSegment() Iterator { 1325 if seg.node.hasChildren { 1326 return seg.node.children[seg.index+1].firstSegment() 1327 } 1328 if seg.index < seg.node.nrSegments-1 { 1329 return Iterator{seg.node, seg.index + 1} 1330 } 1331 if seg.node.parent == nil { 1332 return Iterator{} 1333 } 1334 return segmentAfterPosition(seg.node.parent, seg.node.parentIndex) 1335 } 1336 1337 // PrevGap returns the gap immediately before the iterated segment. 1338 func (seg Iterator) PrevGap() GapIterator { 1339 if seg.node.hasChildren { 1340 // Note that this isn't recursive because the last segment in a subtree 1341 // must be in a leaf node. 1342 return seg.node.children[seg.index].lastSegment().NextGap() 1343 } 1344 return GapIterator{seg.node, seg.index} 1345 } 1346 1347 // NextGap returns the gap immediately after the iterated segment. 1348 func (seg Iterator) NextGap() GapIterator { 1349 if seg.node.hasChildren { 1350 return seg.node.children[seg.index+1].firstSegment().PrevGap() 1351 } 1352 return GapIterator{seg.node, seg.index + 1} 1353 } 1354 1355 // PrevNonEmpty returns the iterated segment's predecessor if it is adjacent, 1356 // or the gap before the iterated segment otherwise. If seg.Start() == 1357 // Functions.MinKey(), PrevNonEmpty will return two terminal iterators. 1358 // Otherwise, exactly one of the iterators returned by PrevNonEmpty will be 1359 // non-terminal. 1360 func (seg Iterator) PrevNonEmpty() (Iterator, GapIterator) { 1361 gap := seg.PrevGap() 1362 if gap.Range().Length() != 0 { 1363 return Iterator{}, gap 1364 } 1365 return gap.PrevSegment(), GapIterator{} 1366 } 1367 1368 // NextNonEmpty returns the iterated segment's successor if it is adjacent, or 1369 // the gap after the iterated segment otherwise. If seg.End() == 1370 // Functions.MaxKey(), NextNonEmpty will return two terminal iterators. 1371 // Otherwise, exactly one of the iterators returned by NextNonEmpty will be 1372 // non-terminal. 1373 func (seg Iterator) NextNonEmpty() (Iterator, GapIterator) { 1374 gap := seg.NextGap() 1375 if gap.Range().Length() != 0 { 1376 return Iterator{}, gap 1377 } 1378 return gap.NextSegment(), GapIterator{} 1379 } 1380 1381 // A GapIterator is conceptually one of: 1382 // 1383 // - A pointer to a position between two segments, before the first segment, or 1384 // after the last segment in a set, called a *gap*; or 1385 // 1386 // - A terminal iterator, which is a sentinel indicating that the end of 1387 // iteration has been reached. 1388 // 1389 // Note that the gap between two adjacent segments exists (iterators to it are 1390 // non-terminal), but has a length of zero. GapIterator.IsEmpty returns true 1391 // for such gaps. An empty set contains a single gap, spanning the entire range 1392 // of the set's keys. 1393 // 1394 // GapIterators are copyable values and are meaningfully equality-comparable. 1395 // The zero value of GapIterator is a terminal iterator. 1396 // 1397 // Unless otherwise specified, any mutation of a set invalidates all existing 1398 // iterators into the set. 1399 type GapIterator struct { 1400 // The representation of a GapIterator is identical to that of an Iterator, 1401 // except that index corresponds to positions between segments in the same 1402 // way as for node.children (see comment for node.nrSegments). 1403 node *node 1404 index int 1405 } 1406 1407 // Ok returns true if the iterator is not terminal. All other methods are only 1408 // valid for non-terminal iterators. 1409 func (gap GapIterator) Ok() bool { 1410 return gap.node != nil 1411 } 1412 1413 // Range returns the range spanned by the iterated gap. 1414 func (gap GapIterator) Range() Range { 1415 return Range{gap.Start(), gap.End()} 1416 } 1417 1418 // Start is equivalent to Range().Start, but should be preferred if only the 1419 // start of the range is needed. 1420 func (gap GapIterator) Start() Key { 1421 if ps := gap.PrevSegment(); ps.Ok() { 1422 return ps.End() 1423 } 1424 return Functions{}.MinKey() 1425 } 1426 1427 // End is equivalent to Range().End, but should be preferred if only the end of 1428 // the range is needed. 1429 func (gap GapIterator) End() Key { 1430 if ns := gap.NextSegment(); ns.Ok() { 1431 return ns.Start() 1432 } 1433 return Functions{}.MaxKey() 1434 } 1435 1436 // IsEmpty returns true if the iterated gap is empty (that is, the "gap" is 1437 // between two adjacent segments.) 1438 func (gap GapIterator) IsEmpty() bool { 1439 return gap.Range().Length() == 0 1440 } 1441 1442 // PrevSegment returns the segment immediately before the iterated gap. If no 1443 // such segment exists, PrevSegment returns a terminal iterator. 1444 func (gap GapIterator) PrevSegment() Iterator { 1445 return segmentBeforePosition(gap.node, gap.index) 1446 } 1447 1448 // NextSegment returns the segment immediately after the iterated gap. If no 1449 // such segment exists, NextSegment returns a terminal iterator. 1450 func (gap GapIterator) NextSegment() Iterator { 1451 return segmentAfterPosition(gap.node, gap.index) 1452 } 1453 1454 // PrevGap returns the iterated gap's predecessor. If no such gap exists, 1455 // PrevGap returns a terminal iterator. 1456 func (gap GapIterator) PrevGap() GapIterator { 1457 seg := gap.PrevSegment() 1458 if !seg.Ok() { 1459 return GapIterator{} 1460 } 1461 return seg.PrevGap() 1462 } 1463 1464 // NextGap returns the iterated gap's successor. If no such gap exists, NextGap 1465 // returns a terminal iterator. 1466 func (gap GapIterator) NextGap() GapIterator { 1467 seg := gap.NextSegment() 1468 if !seg.Ok() { 1469 return GapIterator{} 1470 } 1471 return seg.NextGap() 1472 } 1473 1474 // NextLargeEnoughGap returns the iterated gap's first next gap with larger 1475 // length than minSize. If not found, return a terminal gap iterator (does NOT 1476 // include this gap itself). 1477 // 1478 // Precondition: trackGaps must be 1. 1479 func (gap GapIterator) NextLargeEnoughGap(minSize Key) GapIterator { 1480 if trackGaps != 1 { 1481 panic("set is not tracking gaps") 1482 } 1483 if gap.node != nil && gap.node.hasChildren && gap.index == gap.node.nrSegments { 1484 // If gap is the trailing gap of an non-leaf node, 1485 // translate it to the equivalent gap on leaf level. 1486 gap.node = gap.NextSegment().node 1487 gap.index = 0 1488 return gap.nextLargeEnoughGapHelper(minSize) 1489 } 1490 return gap.nextLargeEnoughGapHelper(minSize) 1491 } 1492 1493 // nextLargeEnoughGapHelper is the helper function used by NextLargeEnoughGap 1494 // to do the real recursions. 1495 // 1496 // Preconditions: gap is NOT the trailing gap of a non-leaf node. 1497 func (gap GapIterator) nextLargeEnoughGapHelper(minSize Key) GapIterator { 1498 // Crawl up the tree if no large enough gap in current node or the 1499 // current gap is the trailing one on leaf level. 1500 for gap.node != nil && 1501 (gap.node.maxGap.Get() < minSize || (!gap.node.hasChildren && gap.index == gap.node.nrSegments)) { 1502 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1503 } 1504 // If no large enough gap throughout the whole set, return a terminal 1505 // gap iterator. 1506 if gap.node == nil { 1507 return GapIterator{} 1508 } 1509 // Iterate subsequent gaps. 1510 gap.index++ 1511 for gap.index <= gap.node.nrSegments { 1512 if gap.node.hasChildren { 1513 if largeEnoughGap := gap.node.children[gap.index].searchFirstLargeEnoughGap(minSize); largeEnoughGap.Ok() { 1514 return largeEnoughGap 1515 } 1516 } else { 1517 if gap.Range().Length() >= minSize { 1518 return gap 1519 } 1520 } 1521 gap.index++ 1522 } 1523 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1524 if gap.node != nil && gap.index == gap.node.nrSegments { 1525 // If gap is the trailing gap of a non-leaf node, crawl up to 1526 // parent again and do recursion. 1527 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1528 } 1529 return gap.nextLargeEnoughGapHelper(minSize) 1530 } 1531 1532 // PrevLargeEnoughGap returns the iterated gap's first prev gap with larger or 1533 // equal length than minSize. If not found, return a terminal gap iterator 1534 // (does NOT include this gap itself). 1535 // 1536 // Precondition: trackGaps must be 1. 1537 func (gap GapIterator) PrevLargeEnoughGap(minSize Key) GapIterator { 1538 if trackGaps != 1 { 1539 panic("set is not tracking gaps") 1540 } 1541 if gap.node != nil && gap.node.hasChildren && gap.index == 0 { 1542 // If gap is the first gap of an non-leaf node, 1543 // translate it to the equivalent gap on leaf level. 1544 gap.node = gap.PrevSegment().node 1545 gap.index = gap.node.nrSegments 1546 return gap.prevLargeEnoughGapHelper(minSize) 1547 } 1548 return gap.prevLargeEnoughGapHelper(minSize) 1549 } 1550 1551 // prevLargeEnoughGapHelper is the helper function used by PrevLargeEnoughGap 1552 // to do the real recursions. 1553 // 1554 // Preconditions: gap is NOT the first gap of a non-leaf node. 1555 func (gap GapIterator) prevLargeEnoughGapHelper(minSize Key) GapIterator { 1556 // Crawl up the tree if no large enough gap in current node or the 1557 // current gap is the first one on leaf level. 1558 for gap.node != nil && 1559 (gap.node.maxGap.Get() < minSize || (!gap.node.hasChildren && gap.index == 0)) { 1560 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1561 } 1562 // If no large enough gap throughout the whole set, return a terminal 1563 // gap iterator. 1564 if gap.node == nil { 1565 return GapIterator{} 1566 } 1567 // Iterate previous gaps. 1568 gap.index-- 1569 for gap.index >= 0 { 1570 if gap.node.hasChildren { 1571 if largeEnoughGap := gap.node.children[gap.index].searchLastLargeEnoughGap(minSize); largeEnoughGap.Ok() { 1572 return largeEnoughGap 1573 } 1574 } else { 1575 if gap.Range().Length() >= minSize { 1576 return gap 1577 } 1578 } 1579 gap.index-- 1580 } 1581 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1582 if gap.node != nil && gap.index == 0 { 1583 // If gap is the first gap of a non-leaf node, crawl up to 1584 // parent again and do recursion. 1585 gap.node, gap.index = gap.node.parent, gap.node.parentIndex 1586 } 1587 return gap.prevLargeEnoughGapHelper(minSize) 1588 } 1589 1590 // segmentBeforePosition returns the predecessor segment of the position given 1591 // by n.children[i], which may or may not contain a child. If no such segment 1592 // exists, segmentBeforePosition returns a terminal iterator. 1593 func segmentBeforePosition(n *node, i int) Iterator { 1594 for i == 0 { 1595 if n.parent == nil { 1596 return Iterator{} 1597 } 1598 n, i = n.parent, n.parentIndex 1599 } 1600 return Iterator{n, i - 1} 1601 } 1602 1603 // segmentAfterPosition returns the successor segment of the position given by 1604 // n.children[i], which may or may not contain a child. If no such segment 1605 // exists, segmentAfterPosition returns a terminal iterator. 1606 func segmentAfterPosition(n *node, i int) Iterator { 1607 for i == n.nrSegments { 1608 if n.parent == nil { 1609 return Iterator{} 1610 } 1611 n, i = n.parent, n.parentIndex 1612 } 1613 return Iterator{n, i} 1614 } 1615 1616 func zeroValueSlice(slice []Value) { 1617 // TODO(jamieliu): check if Go is actually smart enough to optimize a 1618 // ClearValue that assigns nil to a memset here. 1619 for i := range slice { 1620 Functions{}.ClearValue(&slice[i]) 1621 } 1622 } 1623 1624 func zeroNodeSlice(slice []*node) { 1625 for i := range slice { 1626 slice[i] = nil 1627 } 1628 } 1629 1630 // String stringifies a Set for debugging. 1631 func (s *Set) String() string { 1632 return s.root.String() 1633 } 1634 1635 // String stringifies a node (and all of its children) for debugging. 1636 func (n *node) String() string { 1637 var buf bytes.Buffer 1638 n.writeDebugString(&buf, "") 1639 return buf.String() 1640 } 1641 1642 func (n *node) writeDebugString(buf *bytes.Buffer, prefix string) { 1643 if n.hasChildren != (n.nrSegments > 0 && n.children[0] != nil) { 1644 buf.WriteString(prefix) 1645 buf.WriteString(fmt.Sprintf("WARNING: inconsistent value of hasChildren: got %v, want %v\n", n.hasChildren, !n.hasChildren)) 1646 } 1647 for i := 0; i < n.nrSegments; i++ { 1648 if child := n.children[i]; child != nil { 1649 cprefix := fmt.Sprintf("%s- % 3d ", prefix, i) 1650 if child.parent != n || child.parentIndex != i { 1651 buf.WriteString(cprefix) 1652 buf.WriteString(fmt.Sprintf("WARNING: inconsistent linkage to parent: got (%p, %d), want (%p, %d)\n", child.parent, child.parentIndex, n, i)) 1653 } 1654 child.writeDebugString(buf, fmt.Sprintf("%s- % 3d ", prefix, i)) 1655 } 1656 buf.WriteString(prefix) 1657 if n.hasChildren { 1658 if trackGaps != 0 { 1659 buf.WriteString(fmt.Sprintf("- % 3d: %v => %v, maxGap: %d\n", i, n.keys[i], n.values[i], n.maxGap.Get())) 1660 } else { 1661 buf.WriteString(fmt.Sprintf("- % 3d: %v => %v\n", i, n.keys[i], n.values[i])) 1662 } 1663 } else { 1664 buf.WriteString(fmt.Sprintf("- % 3d: %v => %v\n", i, n.keys[i], n.values[i])) 1665 } 1666 } 1667 if child := n.children[n.nrSegments]; child != nil { 1668 child.writeDebugString(buf, fmt.Sprintf("%s- % 3d ", prefix, n.nrSegments)) 1669 } 1670 } 1671 1672 // SegmentDataSlices represents segments from a set as slices of start, end, and 1673 // values. SegmentDataSlices is primarily used as an intermediate representation 1674 // for save/restore and the layout here is optimized for that. 1675 // 1676 // +stateify savable 1677 type SegmentDataSlices struct { 1678 Start []Key 1679 End []Key 1680 Values []Value 1681 } 1682 1683 // ExportSortedSlices returns a copy of all segments in the given set, in 1684 // ascending key order. 1685 func (s *Set) ExportSortedSlices() *SegmentDataSlices { 1686 var sds SegmentDataSlices 1687 for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() { 1688 sds.Start = append(sds.Start, seg.Start()) 1689 sds.End = append(sds.End, seg.End()) 1690 sds.Values = append(sds.Values, seg.Value()) 1691 } 1692 sds.Start = sds.Start[:len(sds.Start):len(sds.Start)] 1693 sds.End = sds.End[:len(sds.End):len(sds.End)] 1694 sds.Values = sds.Values[:len(sds.Values):len(sds.Values)] 1695 return &sds 1696 } 1697 1698 // ImportSortedSlices initializes the given set from the given slice. 1699 // 1700 // Preconditions: 1701 // * s must be empty. 1702 // * sds must represent a valid set (the segments in sds must have valid 1703 // lengths that do not overlap). 1704 // * The segments in sds must be sorted in ascending key order. 1705 func (s *Set) ImportSortedSlices(sds *SegmentDataSlices) error { 1706 if !s.IsEmpty() { 1707 return fmt.Errorf("cannot import into non-empty set %v", s) 1708 } 1709 gap := s.FirstGap() 1710 for i := range sds.Start { 1711 r := Range{sds.Start[i], sds.End[i]} 1712 if !gap.Range().IsSupersetOf(r) { 1713 return fmt.Errorf("segment overlaps a preceding segment or is incorrectly sorted: [%d, %d) => %v", sds.Start[i], sds.End[i], sds.Values[i]) 1714 } 1715 gap = s.InsertWithoutMerging(gap, r, sds.Values[i]).NextGap() 1716 } 1717 return nil 1718 } 1719 1720 // segmentTestCheck returns an error if s is incorrectly sorted, does not 1721 // contain exactly expectedSegments segments, or contains a segment which 1722 // fails the passed check. 1723 // 1724 // This should be used only for testing, and has been added to this package for 1725 // templating convenience. 1726 func (s *Set) segmentTestCheck(expectedSegments int, segFunc func(int, Range, Value) error) error { 1727 havePrev := false 1728 prev := Key(0) 1729 nrSegments := 0 1730 for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() { 1731 next := seg.Start() 1732 if havePrev && prev >= next { 1733 return fmt.Errorf("incorrect order: key %d (segment %d) >= key %d (segment %d)", prev, nrSegments-1, next, nrSegments) 1734 } 1735 if segFunc != nil { 1736 if err := segFunc(nrSegments, seg.Range(), seg.Value()); err != nil { 1737 return err 1738 } 1739 } 1740 prev = next 1741 havePrev = true 1742 nrSegments++ 1743 } 1744 if nrSegments != expectedSegments { 1745 return fmt.Errorf("incorrect number of segments: got %d, wanted %d", nrSegments, expectedSegments) 1746 } 1747 return nil 1748 } 1749 1750 // countSegments counts the number of segments in the set. 1751 // 1752 // Similar to Check, this should only be used for testing. 1753 func (s *Set) countSegments() (segments int) { 1754 for seg := s.FirstSegment(); seg.Ok(); seg = seg.NextSegment() { 1755 segments++ 1756 } 1757 return segments 1758 }