private/intset: add intset package

Change-Id: I801275eb48a923ffca349973d7c14ccdfa37c46a

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])
+	}
+}

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)
+		}
+	})
+}