Iterator functions

README.md

@@ -34,13 +34,24 @@ fmt.Println(base)
 [10]{1010000000}
 ```
 
-### Iteration
+### Iterator
+
+Go >=1.23 iterator is exposed by `.Bits()` method:
 
 ```go
 bm := bitmask.New(5)
 bm.Set(0)
 bm.Set(3)
 
+for idx, isSet := range bm.Bits() {
+    // use the value
+    fmt.Printf("%v) %v\n", index, value)
+}
+```
+
+There's also an old-style equivalent:
+
+```go
 it := bm.Iterator()
 for {
     ok, value, index := it.Next()
@@ -72,3 +83,16 @@ fmt.Println(bm)
 ```
 [513]{0000000000000000000000000000000000000000000000000000000000000000 0000000000000000000000000000000000000000000000000000000000000000 0000000000000000000000000000000000000000000000000000000000000000 0000000000000000000000000000000000000000000000000000000000000001 <more 64 bits> 1000000000000000000000000000000000000000000000000000000000000000 0000000000000000000000000000000000000000000000000000000000000000 0000000000000000000000000000000000000000000000000000000000000000 0}
 ```
+
+### Non-copying constructor to avoid heap-allocating bit buffer
+
+```go
+bm := bitmask.NewFromUintRawNocopy(1, 2, 3)
+```
+
+From `go build -gcflags=-m`
+```
+inlining call to bitmask.NewFromUintRawNocopy
+... argument does not escape
+&bitmask.BitMask{...} does not escape
+```

bitmask.go

@@ -2,6 +2,7 @@ package bitmask
 
 import (
 	"fmt"
+	"iter"
 	"reflect"
 	"strings"
 	"unsafe"
@@ -321,6 +322,7 @@ type BitIterator struct {
 // Creates stateful iterator to iterate through all the bits.
 // See BitIterator doc for an example.
 // It is an equivalent of just calling IsSet for each bit of a BitMask.
+// For Go >=1.23 also consider Bits(), which is equivalent.
 func (bm *BitMask) Iterator() BitIterator {
 	index := uint(0)
 	return BitIterator{
@@ -340,6 +342,19 @@ func (bm *BitMask) Iterator() BitIterator {
 	}
 }
 
+// Go >=1.23 iterator. Equivalent of Iterator().
+func (bm *BitMask) Bits() iter.Seq2[uint, bool] {
+	return func(yield func(uint, bool) bool) {
+		for i := uint(0); i < bm.len; i++ {
+			bref, m := bm.getBit(i)
+			isSet := (*bref & m) != 0
+			if !yield(i, isSet) {
+				return
+			}
+		}
+	}
+}
+
 // Returns string representation of a bitmask in the form "[length]{bits}".
 // For example: [4]{0100}
 // It is O(1) operation and it will skip bits after some amount of them.

bitmask_test.go

@@ -2,8 +2,11 @@ package bitmask
 
 import (
 	"fmt"
+	"iter"
 	"math/rand"
+	"slices"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -404,17 +407,99 @@ func TestSlice(t *testing.T) {
 }
 
 func TestIterator(t *testing.T) {
-	bm := NewFromUint(uintMax, 0).Slice(62, 66)
-	it := bm.Iterator()
+	// arrange
+	uintNumber := uint(15)
+	bitNumber := uintSize * uintNumber
+	workerCount := uint(40)
+
+	shouldBeSet := func(i uint) bool {
+		return i/7%5 == 0 || i/2%5 == 0 || i%9 == 0
+	}
+
+	bm := NewFromUint(make([]uint, uintNumber)...)
+	for i := uint(0); i < bitNumber; i++ {
+		if shouldBeSet(i) {
+			bm.Set(i)
+		}
+
+	}
+
+	// act
+	wg := sync.WaitGroup{}
+
+	// new iterators
+	wg.Add(int(workerCount))
+	for workerId := uint(0); workerId < workerCount; workerId++ {
+		go func() {
+			defer wg.Done()
+
+			workerIdCopy := workerId
+
+			expectedIdx := uint(0)
+			for idx, isSet := range bm.Bits() {
+				assert.Equal(t, expectedIdx, idx)
+				assert.Equal(t, shouldBeSet(idx), isSet)
+				expectedIdx++
+
+				if workerIdCopy%10 == 0 && workerIdCopy == idx {
+					break
+				}
+			}
+		}()
+	}
+
+	// old iterators
+	wg.Add(int(workerCount))
+	for workerId := uint(0); workerId < workerCount; workerId++ {
+		go func() {
+			defer wg.Done()
+
+			workerIdCopy := workerId
+
+			it := bm.Iterator()
+			expectedIdx := uint(0)
+			totalIdx := uint(0)
+			for {
+				ok, value, idx := it.Next()
+
+				assert.Equal(t, expectedIdx, idx)
+				assert.Equal(t, expectedIdx < bitNumber, ok)
+				if ok {
+					if shouldBeSet(idx) != value {
+						t.Fail()
+					}
+					assert.Equal(t, shouldBeSet(idx), value)
+				} else {
+					assert.Equal(t, false, value)
+				}
+
+				expectedIdx++
+				totalIdx++
+
+				if workerIdCopy%8 == 0 && (totalIdx+workerIdCopy)%10 == 0 {
+					it.Reset()
+					expectedIdx = 0
+					// so that we can continue
+					workerIdCopy++
+				}
+
+				if !ok {
+					if workerIdCopy%7 == 0 && totalIdx < bitNumber*4 {
+						it.Reset()
+						expectedIdx = 0
+					} else {
+						break
+					}
+				}
+			}
 
-	for i := uint(0); i < 4; i++ {
-		ok, value, index := it.Next()
-		assert.Equal(t, i, index)
-		assert.Equal(t, true, ok)
-		assert.Equal(t, i < 2, value)
+			// try iterate after end
+			ok, _, _ := it.Next()
+			assert.Equal(t, false, ok)
+		}()
 	}
-	ok, _, _ := it.Next()
-	assert.Equal(t, false, ok)
+
+	wg.Wait()
 }
 
 func TestSliceWithString(t *testing.T) {
@@ -702,6 +787,21 @@ func indexes(it BitIterator) []int {
 	}
 	return indexes
 }
+
+// returns indexes of all set bits in a bitmask
+func indexes2(bm *BitMask) iter.Seq[uint] {
+	return func(yield func(uint) bool) {
+		for idx, isSet := range bm.Bits() {
+			if isSet {
+				if !yield(idx) {
+					return
+				}
+			}
+		}
+
+	}
+}
+
 func TestIndexes(t *testing.T) {
 	bm := New(200)
 	bm.Set(1)
@@ -724,6 +824,24 @@ func TestIndexes(t *testing.T) {
 	assert.Equal(t, expected, i4)
 }
 
+func TestIndexes2(t *testing.T) {
+	bm := New(200)
+	bm.Set(1)
+	bm.Set(10)
+	bm.Set(100)
+
+	i1 := slices.Collect(indexes2(bm))
+	i2 := slices.Collect(indexes2(bm))
+	i3 := slices.Collect(indexes2(bm))
+	i4 := slices.Collect(indexes2(bm))
+
+	expected := []uint{1, 10, 100}
+	assert.Equal(t, expected, i1)
+	assert.Equal(t, expected, i2)
+	assert.Equal(t, expected, i3)
+	assert.Equal(t, expected, i4)
+}
+
 func TestUintRaw(t *testing.T) {
 	bm := NewFromUintRaw(1, 1, 0)
 	indexes := indexes(bm.Iterator())
@@ -744,11 +862,11 @@ func TestUintRaw(t *testing.T) {
 func TestUintRawNocopy(t *testing.T) {
 	buf := []uint{1, 1, 0}
 	bm := NewFromUintRawNocopy(buf...)
-	indexes := indexes(bm.Iterator())
+	indexes := slices.Collect(indexes2(bm))
 
 	assert.Equal(t, 3, bm.LenUint())
 
-	assert.Equal(t, []int{uintSize - 1, uintSize + uintSize - 1}, indexes)
+	assert.Equal(t, []uint{uintSize - 1, uintSize + uintSize - 1}, indexes)
 
 	assert.Equal(t, uint(1<<63), bm.Uint(0))
 	assert.Equal(t, uint(1<<63), bm.Uint(1))

go.mod

@@ -1,8 +1,8 @@
 module github.com/astef/bitmask
 
-go 1.22
+go 1.23
 
-require github.com/stretchr/testify v1.8.2
+require github.com/stretchr/testify v1.9.0
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect

go.sum

@@ -1,17 +1,10 @@
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
-github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
-github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8=
-github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=