diff --git a/private/intset/set.go b/private/intset/set.go new file mode 100644 index 000000000000..2111919e9396 --- /dev/null +++ b/private/intset/set.go @@ -0,0 +1,86 @@ +// Copyright (C) 2023 Storj Labs, Inc. +// See LICENSE for copying information. + +package intset + +import "math/bits" + +const ( + bucketBits = 32 +) + +// Set set of int values. +type Set struct { + size int + bits []uint32 + count int +} + +// NewSet creates new int set. +func NewSet(size int) Set { + return Set{ + size: size, + bits: make([]uint32, (size+(bucketBits-1))/bucketBits), + } +} + +// Contains returns true if set includes int value. +// Ignores negative values and above set length. +func (s Set) Contains(value int) bool { + if value < 0 || value >= s.size { + return false + } + bucket, bit := offset(value) + return (s.bits[bucket] & (uint32(1) << bit)) != 0 +} + +// Include includes int value into set. +// Ignores negative values and above set length. +func (s *Set) Include(value int) { + if value < 0 || value >= s.size { + return + } + + if !s.Contains(value) { + bucket, bit := offset(value) + s.bits[bucket] |= (uint32(1) << bit) + s.count++ + } +} + +// Exclude removes int value from set. +// Ignores negative values and above set length. +func (s *Set) Exclude(value int) { + if value < 0 || value >= s.size { + return + } + if s.Contains(value) { + bucket, bit := offset(value) + s.bits[bucket] &= ^(uint32(1) << bit) + s.count-- + } +} + +// Count returns number of int values in a set. +func (s Set) Count() int { + return s.count +} + +func offset(value int) (bucket, bit int) { + return value / bucketBits, value % bucketBits +} + +// Add adds to the set content of other sets. +// Sources sets with different size than target will be ignored. +func (s *Set) Add(sources ...Set) { + s.count = 0 + for k, b := range s.bits { + for srci := range sources { + if s.size == sources[srci].size { + b |= sources[srci].bits[k] + } + } + s.bits[k] = b + s.count += bits.OnesCount32(s.bits[k]) + } +} diff --git a/private/intset/set_test.go b/private/intset/set_test.go new file mode 100644 index 000000000000..8e84436571df --- /dev/null +++ b/private/intset/set_test.go @@ -0,0 +1,79 @@ +// Copyright (C) 2023 Storj Labs, Inc. +// See LICENSE for copying information. + +package intset_test + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "storj.io/storj/private/intset" +) + +func TestSet(t *testing.T) { + + set := intset.NewSet(5) + + require.Zero(t, set.Count()) + + for i := -1; i < 10; i++ { + set.Include(i) + } + + require.Equal(t, 5, set.Count()) + + for i := 0; i < 5; i++ { + require.True(t, set.Contains(i)) + + set.Exclude(i) + } + + for i := -1; i < 10; i++ { + require.False(t, set.Contains(i), "#%d", i) + } +} + +func TestCopySet(t *testing.T) { + setA := intset.NewSet(10) + setB := intset.NewSet(10) + + for i := 0; i < 10; i++ { + if i%2 == 0 { + setA.Include(i) + } else { + setB.Include(i) + } + } + require.Equal(t, 5, setA.Count()) + require.Equal(t, 5, setB.Count()) + + setC := intset.NewSet(10) + setC.Add(setA, setB) + for i := 0; i < 10; i++ { + require.True(t, setC.Contains(i)) + } + require.Equal(t, 10, setC.Count()) + + setD := intset.NewSet(3) + setE := intset.NewSet(3) + setE.Include(0) + setE.Include(2) + // set with different initial size will be ignored while adding + setF := intset.NewSet(5) + setF.Include(1) + setD.Add(setE, setF) + + require.Equal(t, 2, setD.Count()) + for i, contains := range []bool{true, false, true, false, false} { + require.Equal(t, contains, setD.Contains(i), "#%d", i) + } +} + +func BenchmarkIntSet(b *testing.B) { + b.Run("create", func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = intset.NewSet(1000) + } + }) +}