github.com/anacrolix/torrent@v1.61.0/internal/indexed/map.go (about)

     1  package indexed
     2  
     3  import (
     4  	g "github.com/anacrolix/generics"
     5  	"github.com/anacrolix/generics/option"
     6  	"github.com/anacrolix/missinggo/v2/panicif"
     7  )
     8  
     9  type mapTriggerFunc[K, V any] func(key K, old, new g.Option[V])
    10  
    11  // A table where the record has a key and value. Currently maps onto Table2 using Pair, but perhaps
    12  // requiring the record implements a KeyValue interface would be better. K and V are Record to
    13  // propagate the comparable requirement for now. TODO: We could use an actual map, and implement
    14  // relation. Another thing we can do is actually use btree.Map's value, I'm not sure if btree is
    15  // smarter about that.
    16  type Map[K, V Record] struct {
    17  	Table2[K, V]
    18  	keyCmp func(a, b K) int
    19  }
    20  
    21  func (me *Map[K, V]) Init(cmp func(a, b K) int) {
    22  	me.Table2.Init(func(a, b Pair[K, V]) int {
    23  		return cmp(a.Left, b.Left)
    24  	})
    25  	me.keyCmp = cmp
    26  }
    27  
    28  type UpdateResult struct {
    29  	Exists  bool
    30  	Changed bool
    31  }
    32  
    33  func (me *Map[K, V]) Update(key K, updateFunc func(V) V) (res UpdateResult) {
    34  	start := Pair[K, V]{Left: key}
    35  	gte := me.GetGte(start)
    36  	if !gte.Ok {
    37  		return
    38  	}
    39  	old := gte.Value
    40  	if me.keyCmp(old.Left, key) != 0 {
    41  		return
    42  	}
    43  	res.Exists = true
    44  	new := old
    45  	new.Right = updateFunc(old.Right)
    46  	if new.Right == old.Right {
    47  		return
    48  	}
    49  	replaced, overwrote := me.set.Upsert(new)
    50  	panicif.False(overwrote)
    51  	panicif.NotEq(replaced, old)
    52  	panicif.NotZero(me.cmp(replaced, old))
    53  	// Is this lazy? Should the caller be triggering instead?
    54  	res.Changed = true
    55  	me.Changed(g.Some(old), g.Some(new))
    56  	return
    57  }
    58  
    59  func (me *Map[K, V]) Alter(key K, updateFunc func(V, bool) (V, bool)) {
    60  	oldV, oldOk := me.Get(key)
    61  	newV, newOk := updateFunc(oldV, oldOk)
    62  	me.Change(
    63  		g.OptionFromTuple(NewPair(key, oldV), oldOk),
    64  		g.OptionFromTuple(NewPair(key, newV), newOk))
    65  }
    66  
    67  func (me *Map[K, V]) Get(k K) (v V, ok bool) {
    68  	for r := range me.set.IterFrom(Pair[K, V]{Left: k}) {
    69  		if me.keyCmp(r.Left, k) == 0 {
    70  			v = r.Right
    71  			ok = true
    72  			break
    73  		}
    74  	}
    75  	return
    76  }
    77  
    78  func (me *Map[K, V]) ContainsKey(k K) bool {
    79  	_, ok := me.Get(k)
    80  	return ok
    81  }
    82  
    83  func (me *Map[K, V]) Delete(k K) (removed bool) {
    84  	return me.Table2.Delete(Pair[K, V]{Left: k})
    85  }
    86  
    87  // Update the function otherwise create it, in both cases using the update function provided.
    88  func (me *Map[K, V]) UpdateOrCreate(k K, updateFunc func(old V) V) (created bool) {
    89  	me.Alter(k, func(v V, existed bool) (V, bool) {
    90  		created = !existed
    91  		return updateFunc(v), true
    92  	})
    93  	return
    94  }
    95  
    96  func (me *Map[K, V]) OnValueChange(do mapTriggerFunc[K, V]) {
    97  	me.OnChange(func(old, new g.Option[Pair[K, V]]) {
    98  		if old.Ok && new.Ok {
    99  			panicif.NotZero(me.keyCmp(old.Value.Left, new.Value.Left))
   100  		}
   101  		// Key must be one or the other.
   102  		var key K
   103  		if old.Ok {
   104  			key = old.Value.Left
   105  		} else {
   106  			key = new.Unwrap().Left
   107  		}
   108  		do(key, option.Map(pairMapRight, old), option.Map(pairMapRight, new))
   109  	})
   110  }
   111  
   112  // Maps don't compare on the value, so we can leave them as zeroes.
   113  func (me *Map[K, V]) SetMinRecord(min K) {
   114  	me.table.SetMinRecord(Pair[K, V]{Left: min})
   115  }