github.com/anacrolix/torrent@v1.61.0/ordered-bitmap.go (about)

     1  package torrent
     2  
     3  import (
     4  	"iter"
     5  
     6  	g "github.com/anacrolix/generics"
     7  	list "github.com/bahlo/generic-list-go"
     8  
     9  	typedRoaring "github.com/anacrolix/torrent/typed-roaring"
    10  )
    11  
    12  type orderedBitmap[T typedRoaring.BitConstraint] struct {
    13  	bitmap typedRoaring.Bitmap[T]
    14  	// There should be way more efficient ways to do this.
    15  	order    list.List[T]
    16  	elements map[T]*list.Element[T]
    17  }
    18  
    19  func (o *orderedBitmap[T]) IterateSnapshot(f func(T) bool) {
    20  	o.bitmap.Clone().Iterate(f)
    21  }
    22  
    23  func (o *orderedBitmap[T]) IsEmpty() bool {
    24  	return o.bitmap.IsEmpty()
    25  }
    26  
    27  func (o *orderedBitmap[T]) GetCardinality() uint64 {
    28  	return uint64(o.order.Len())
    29  }
    30  
    31  func (o *orderedBitmap[T]) Contains(index T) bool {
    32  	return o.bitmap.Contains(index)
    33  }
    34  
    35  func (o *orderedBitmap[T]) Add(index T) {
    36  	o.bitmap.Add(index)
    37  	if _, ok := o.elements[index]; !ok {
    38  		g.MakeMapIfNilAndSet(&o.elements, index, o.order.PushBack(index))
    39  	}
    40  }
    41  
    42  func (o *orderedBitmap[T]) Rank(index T) uint64 {
    43  	return o.bitmap.Rank(index)
    44  }
    45  
    46  func (o *orderedBitmap[T]) Iterate(f func(T) bool) (all bool) {
    47  	for e := o.order.Front(); e != nil; e = e.Next() {
    48  		if !f(e.Value) {
    49  			return
    50  		}
    51  	}
    52  	all = true
    53  	return
    54  }
    55  
    56  func (o *orderedBitmap[T]) Iterator() iter.Seq[T] {
    57  	return func(yield func(T) bool) {
    58  		for e := o.order.Front(); e != nil; e = e.Next() {
    59  			if !yield(e.Value) {
    60  				return
    61  			}
    62  		}
    63  	}
    64  }
    65  
    66  func (o *orderedBitmap[T]) CheckedRemove(index T) bool {
    67  	if !o.bitmap.CheckedRemove(index) {
    68  		return false
    69  	}
    70  	o.order.Remove(o.elements[index])
    71  	delete(o.elements, index)
    72  	return true
    73  }