diff --git a/iter.go b/iter.go index 0f5221d..66b1997 100644 --- a/iter.go +++ b/iter.go @@ -12,20 +12,112 @@ import "iter" // All returns an iterator over key-value pairs from m. func (m *Map[K, E]) All() iter.Seq2[K, E] { - return func(yield func(K, E) bool) { - for it := m.Iter(); it.Next(); { - if !yield(it.Key(), it.Elem()) { + return m.iterFunc +} + +func (m *Map[K, E]) iterFunc(yield func(k K, e E) bool) { + if m == nil || m.count == 0 { + return + } + // Remember we have an iterator. + // Can run concurrently with another m.Iter(). + atomicOr(&m.flags, iterator|oldIterator) + + var ( + r = rand64() + buckets = m.buckets + startBucket = int(r & m.bucketMask()) + nextBucket = startBucket + b *bucket[K, E] + checkBucket int + offset = uint8(r >> (64 - bucketCntBits)) + wrapped bool + ) + + for { + if b == nil { + if nextBucket == startBucket && wrapped { return } + + if m.growing() && len(buckets) == len(m.buckets) { + // Iterator was started in the middle of a grow, and the grow isn't done yet. + // If the bucket we're looking at hasn't been filled in yet (i.e. the old + // bucket hasn't been evacuated) then we need to iterate through the old + // bucket and only return the ones that will be migrated to this bucket. + oldbucket := uint64(nextBucket) & m.oldbucketmask() + b = &(*m.oldbuckets)[oldbucket] + if !evacuated(b) { + checkBucket = nextBucket + } else { + b = &buckets[nextBucket] + checkBucket = noCheck + } + } else { + b = &buckets[nextBucket] + checkBucket = noCheck + } + nextBucket++ + if nextBucket == len(buckets) { + nextBucket = 0 + wrapped = true + } + } + + for i := uint8(0); i < bucketCnt; i++ { + offi := (i + offset) & (bucketCnt - 1) + if isEmpty(b.tophash[offi]) || b.tophash[offi] == evacuatedEmpty { + // TODO: emptyRest is hard to use here, as we start iterating + // in the middle of a bucket. It's feasible, just tricky. + continue + } + k := b.keys[offi] + if checkBucket != noCheck && !m.sameSizeGrow() { + // Special case: iterator was started during a grow to a larger size + // and the grow is not done yet. We're working on a bucket whose + // oldbucket has not been evacuated yet. Or at least, it wasn't + // evacuated when we started the bucket. So we're iterating + // through the oldbucket, skipping any keys that will go + // to the other new bucket (each oldbucket expands to two + // buckets during a grow). + // If the item in the oldbucket is not destined for + // the current new bucket in the iteration, skip it. + hash := m.hash(m.seed, k) + if int(hash&m.bucketMask()) != checkBucket { + continue + } + } + if b.tophash[offi] != evacuatedX && b.tophash[offi] != evacuatedY { + // This is the golden data, we can return it. + if !yield(k, b.elems[offi]) { + return + } + } else { + // The hash table has grown since the iterator was started. + // The golden data for this key is now somewhere else. + // Check the current hash table for the data. + // This code handles the case where the key + // has been deleted, updated, or deleted and reinserted. + // NOTE: we need to regrab the key as it has potentially been + // updated to an equal() but not identical key (e.g. +0.0 vs -0.0). + rk, re := m.mapaccessK(k) + if rk == nil { + continue // key has been deleted + } + if !yield(*rk, *re) { + return + } + } } + b = b.overflow } } // Keys returns an iterator over keys in m. func (m *Map[K, E]) Keys() iter.Seq[K] { return func(yield func(K) bool) { - for it := m.Iter(); it.Next(); { - if !yield(it.Key()) { + for k := range m.All() { + if !yield(k) { return } } @@ -35,8 +127,8 @@ func (m *Map[K, E]) Keys() iter.Seq[K] { // Values returns an iterator over values in m. func (m *Map[K, E]) Values() iter.Seq[E] { return func(yield func(E) bool) { - for it := m.Iter(); it.Next(); { - if !yield(it.Elem()) { + for _, v := range m.All() { + if !yield(v) { return } } diff --git a/iter_test.go b/iter_test.go index b4bbef6..d2d8e1f 100644 --- a/iter_test.go +++ b/iter_test.go @@ -92,3 +92,137 @@ func TestRangeFuncs(t *testing.T) { } }) } + +func TestAll(t *testing.T) { + m := New[uint64, uint64]( + func(a, b uint64) bool { return a == b }, + badIntHash, + ) + expected := make(map[uint64]uint64, 9) + for i := uint64(0); i < 9; i++ { + expected[i] = i + m.Set(i, i) + } + for k, v := range m.All() { + e, ok := expected[k] + if !ok { + t.Errorf("unexpected value in m: [%d: %d]", k, v) + continue + } + if e != v { + t.Errorf("wrong value for key %d. Expected: %d Got: %d", k, e, v) + continue + } + delete(expected, k) + } + if len(expected) > 0 { + t.Errorf("Values not found in m: %v", expected) + } +} + +func TestAllDuringResize(t *testing.T) { + m := New[uint64, uint64]( + func(a, b uint64) bool { return a == b }, + badIntHash, + ) + + // insert numbers that initially hash to the same bucket, but will + // be split into different buckets on resize. Evens will end up in + // bucket[0], odds end up in bucket[1] thanks to badIntHash. + initial := map[uint64]uint64{0: 0, 1: 1, 2: 2, 3: 3} + for k, e := range initial { + m.Set(k, e) + } + additional := map[uint64]uint64{100: 100, 101: 101, 102: 102, 103: 103, 104: 104} + + first := true + // start the iter + for k, v := range m.All() { + if first { + first = false + // Add some additional data to cause a resize + for k, e := range additional { + m.Set(k, e) + } + // Remove 1 value that in each of the initial and split + // buckets that we haven't seen yet + if k == 0 { + m.Delete(2) + delete(initial, 2) + } else { + m.Delete(0) + delete(initial, 0) + } + if k == 1 { + m.Delete(3) + delete(initial, 3) + } else { + m.Delete(1) + delete(initial, 1) + } + } + if k != v { + t.Errorf("expected key == elem, but got: %d != %d", k, v) + t.Error(m.debugString()) + } + if _, ok := initial[k]; ok { + delete(initial, v) + continue + } + if _, ok := additional[k]; ok { + t.Logf("Saw key from additional: %d", k) + continue + } + t.Errorf("Unexpected value from iter: %d", k) + } + for k := range initial { + t.Errorf("iter missing key: %d", k) + } +} + +func TestAllDuringGrow(t *testing.T) { + m := New[uint64, uint64]( + func(a, b uint64) bool { return a == b }, + badIntHash, + ) + + // Insert exactly 27 numbers so we end up in the middle of a grow. + expected := make(map[uint64]uint64, 27) + for i := uint64(0); i < 27; i++ { + expected[i] = i + m.Set(i, i) + } + + for k, v := range m.All() { + t.Logf("Key: %d", k) + if k != v { + t.Errorf("expected key == elem, but got: %d != %d", k, v) + t.Error(m.debugString()) + } + + if _, ok := expected[k]; ok { + delete(expected, v) + continue + } + t.Errorf("Unexpected value from iter: %d", k) + } + for k := range expected { + t.Errorf("iter missing key: %d", k) + } +} + +func BenchmarkAll(b *testing.B) { + m := New[string, int]( + func(a, b string) bool { return a == b }, + maphash.String, + KeyElem[string, int]{"one", 1}, + KeyElem[string, int]{"two", 2}, + KeyElem[string, int]{"three", 3}, + ) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + for range m.All() { + } + } +} diff --git a/map.go b/map.go index 3595f61..08ef460 100644 --- a/map.go +++ b/map.go @@ -649,6 +649,9 @@ search: // Iter instantiates an Iterator to explore the elements of the Map. // Ordering is undefined and is intentionally randomized. +// +// Prefer [Map.All] over Iter when using go1.23 or later as it works +// with for-range loops and has less overhead. func (m *Map[K, E]) Iter() *Iterator[K, E] { // Iter() is a small function to encourage the compiler to inline // it into its caller and let `it` be kept on the stack.