ordered-bitmap.go

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