ordered-bitmap.go

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