github.com/anacrolix/torrent@v1.61.0/internal/indexed/btree-ajwerner.go (about)

     1  package indexed
     2  
     3  import (
     4  	"iter"
     5  
     6  	"github.com/anacrolix/btree"
     7  	"github.com/anacrolix/missinggo/v2/panicif"
     8  )
     9  
    10  type btreeIterator[R any] struct {
    11  	btree.MapIterator[R, struct{}]
    12  	version int
    13  }
    14  
    15  type ajwernerBtreeSet[R any] struct {
    16  	inner   btree.Set[R]
    17  	version int
    18  }
    19  
    20  func (me *ajwernerBtreeSet[R]) Delete(r R) (actual R, removed bool) {
    21  	actual, _, removed = me.inner.Map.Delete(r)
    22  	return
    23  }
    24  
    25  func (me *ajwernerBtreeSet[R]) Upsert(r R) (_ R, overwrote bool) {
    26  	return me.inner.Upsert(r)
    27  }
    28  
    29  func (me *ajwernerBtreeSet[R]) Contains(r R) bool {
    30  	_, ok := me.inner.Get(r)
    31  	return ok
    32  }
    33  
    34  func (me *ajwernerBtreeSet[R]) Len() int {
    35  	return me.inner.Len()
    36  }
    37  
    38  func (me *ajwernerBtreeSet[R]) getBtreeIterator() btreeIterator[R] {
    39  	return btreeIterator[R]{
    40  		MapIterator: me.inner.Iterator(),
    41  		version:     me.version,
    42  	}
    43  }
    44  
    45  func (me *ajwernerBtreeSet[R]) assertIteratorVersion(it btreeIterator[R]) {
    46  	panicif.NotEq(me.version, it.version)
    47  }
    48  
    49  func (me *ajwernerBtreeSet[R]) Iter(yield func(R) bool) {
    50  	it := me.getBtreeIterator()
    51  	for it.First(); it.Valid(); it.Next() {
    52  		if !yield(it.Cur()) {
    53  			return
    54  		}
    55  	}
    56  }
    57  
    58  func (me *ajwernerBtreeSet[R]) IterFrom(start R) iter.Seq[R] {
    59  	return func(yield func(R) bool) {
    60  		it := me.getBtreeIterator()
    61  		it.SeekGE(start)
    62  		for ; it.Valid(); it.Next() {
    63  			if !yield(it.Cur()) {
    64  				return
    65  			}
    66  			me.assertIteratorVersion(it)
    67  		}
    68  	}
    69  }
    70  
    71  func makeAjwernerSet[R any](cmp func(R, R) int) *ajwernerBtreeSet[R] {
    72  	return &ajwernerBtreeSet[R]{inner: btree.MakeSet(cmp)}
    73  }