Skip to content

Commit

Permalink
add method GetAndDelete
Browse files Browse the repository at this point in the history
  • Loading branch information
gaissmai committed Sep 29, 2024
1 parent 7c47de5 commit 3abcb43
Show file tree
Hide file tree
Showing 6 changed files with 63 additions and 11 deletions.
4 changes: 3 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,11 @@ The API has changed in v0.4.2, 0.5.3, v0.6.3, v0.10.1, v0.11.0, v0.12.0, v0.12.6
and/or writers.

func (t *Table[V]) Insert(pfx netip.Prefix, val V)
func (t *Table[V]) Update(pfx netip.Prefix, cb func(val V, ok bool) V) (newVal V)
func (t *Table[V]) Delete(pfx netip.Prefix)

func (t *Table[V]) Get(pfx netip.Prefix) (val V, ok bool)
func (t *Table[V]) Update(pfx netip.Prefix, cb func(val V, ok bool) V) (newVal V)
func (t *Table[V]) GetAndDelete(pfx netip.Prefix) (val V, ok bool)

func (t *Table[V]) Union(o *Table[V])
func (t *Table[V]) Clone() *Table[V]
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@ module github.com/gaissmai/bart

go 1.22.0

require github.com/bits-and-blooms/bitset v1.14.2
require github.com/bits-and-blooms/bitset v1.14.3
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
github.com/bits-and-blooms/bitset v1.14.2 h1:YXVoyPndbdvcEVcseEovVfp0qjJp7S+i5+xgp/Nfbdc=
github.com/bits-and-blooms/bitset v1.14.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
github.com/bits-and-blooms/bitset v1.14.3 h1:Gd2c8lSNf9pKXom5JtD7AaKO8o7fGQ2LtFj1436qilA=
github.com/bits-and-blooms/bitset v1.14.3/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8=
11 changes: 6 additions & 5 deletions node.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,17 +98,18 @@ func (n *node[V]) insertPrefix(idx uint, val V) (ok bool) {
return true
}

// deletePrefix removes the route octet/prefixLen.
// Returns false if there was no prefix to delete.
func (n *node[V]) deletePrefix(octet byte, prefixLen int) (ok bool) {
// deletePrefix removes the route octet/prefixLen and returns the associated value and true
// or false if there was no prefix to delete (and no value to return).
func (n *node[V]) deletePrefix(octet byte, prefixLen int) (val V, ok bool) {
idx := pfxToIdx(octet, prefixLen)

// no route entry
if !n.prefixesBitset.Test(idx) {
return false
return val, false
}

rnk := n.prefixRank(idx)
val = n.prefixes[rnk]

// delete from slice
n.prefixes = slices.Delete(n.prefixes, rnk, rnk+1)
Expand All @@ -117,7 +118,7 @@ func (n *node[V]) deletePrefix(octet byte, prefixLen int) (ok bool) {
n.prefixesBitset.Clear(idx)
n.prefixesBitset.Compact()

return true
return val, true
}

// updatePrefix, update or set the value at prefix via callback. The new value returned
Expand Down
15 changes: 13 additions & 2 deletions table.go
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,16 @@ func (t *Table[V]) Get(pfx netip.Prefix) (val V, ok bool) {

// Delete removes pfx from the tree, pfx does not have to be present.
func (t *Table[V]) Delete(pfx netip.Prefix) {
_, _ = t.getAndDelete(pfx)
}

// GetAndDelete deletes the prefix and returns the associated payload for prefix and true,
// or the zero vlaue and false if prefix is not set in the routing table.
func (t *Table[V]) GetAndDelete(pfx netip.Prefix) (val V, ok bool) {
return t.getAndDelete(pfx)
}

func (t *Table[V]) getAndDelete(pfx netip.Prefix) (val V, ok bool) {
if !pfx.IsValid() || !t.isInit() {
return
}
Expand Down Expand Up @@ -305,8 +315,7 @@ func (t *Table[V]) Delete(pfx netip.Prefix) {
}

// try to delete prefix in trie node
if !n.deletePrefix(lastOctet, lastOctetBits) {
// nothing deleted
if val, ok = n.deletePrefix(lastOctet, lastOctetBits); !ok {
return
}

Expand All @@ -324,6 +333,8 @@ func (t *Table[V]) Delete(pfx netip.Prefix) {
i--
n = stack[i]
}

return val, ok
}

// Lookup does a route lookup (longest prefix match) for IP and
Expand Down
38 changes: 38 additions & 0 deletions table_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1049,6 +1049,44 @@ func TestDeleteIsReverseOfInsert(t *testing.T) {
}
}

func TestGetAndDelete(t *testing.T) {
t.Parallel()
// Insert N prefixes, then delete those same prefixes in shuffled
// order.
const N = 10_000

tbl := new(Table[int])
prefixes := randomPrefixes(N)

// insert the prefixes
for _, p := range prefixes {
tbl.Insert(p.pfx, p.val)
}

// shuffle the prefixes
rand.Shuffle(N, func(i, j int) {
prefixes[i], prefixes[j] = prefixes[j], prefixes[i]
})

for _, p := range prefixes {
want, _ := tbl.Get(p.pfx)
val, ok := tbl.GetAndDelete(p.pfx)

if !ok {
t.Errorf("GetAndDelete, expected true, got %v", ok)
}

if val != want {
t.Errorf("GetAndDelete, expected %v, got %v", want, val)
}

val, ok = tbl.GetAndDelete(p.pfx)
if ok {
t.Errorf("GetAndDelete, expected false, got (%v, %v)", val, ok)
}
}
}

func TestGet(t *testing.T) {
t.Parallel()

Expand Down

0 comments on commit 3abcb43

Please sign in to comment.