diff --git a/src/iter/iter.go b/src/iter/iter.go index e765378ef25f45b0c0329c7697ba2bd33e2d0cea..4d408e5e77c65e54e71abdc4103055f6296fa3a2 100644 --- a/src/iter/iter.go +++ b/src/iter/iter.go @@ -257,91 +257,91 @@ func coroswitch(*coro) // If the iterator panics during a call to next (or stop), // then next (or stop) itself panics with the same value. func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) { - var ( + var pull struct { v V ok bool done bool yieldNext bool + seqDone bool // to detect Goexit racer int panicValue any - seqDone bool // to detect Goexit - ) + } c := newcoro(func(c *coro) { - race.Acquire(unsafe.Pointer(&racer)) - if done { - race.Release(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) + if pull.done { + race.Release(unsafe.Pointer(&pull.racer)) return } yield := func(v1 V) bool { - if done { + if pull.done { return false } - if !yieldNext { + if !pull.yieldNext { panic("iter.Pull: yield called again before next") } - yieldNext = false - v, ok = v1, true - race.Release(unsafe.Pointer(&racer)) + pull.yieldNext = false + pull.v, pull.ok = v1, true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) - return !done + race.Acquire(unsafe.Pointer(&pull.racer)) + return !pull.done } // Recover and propagate panics from seq. defer func() { if p := recover(); p != nil { - panicValue = p - } else if !seqDone { - panicValue = goexitPanicValue + pull.panicValue = p + } else if !pull.seqDone { + pull.panicValue = goexitPanicValue } - done = true // Invalidate iterator - race.Release(unsafe.Pointer(&racer)) + pull.done = true // Invalidate iterator + race.Release(unsafe.Pointer(&pull.racer)) }() seq(yield) var v0 V - v, ok = v0, false - seqDone = true + pull.v, pull.ok = v0, false + pull.seqDone = true }) next = func() (v1 V, ok1 bool) { - race.Write(unsafe.Pointer(&racer)) // detect races + race.Write(unsafe.Pointer(&pull.racer)) // detect races - if done { + if pull.done { return } - if yieldNext { + if pull.yieldNext { panic("iter.Pull: next called again before yield") } - yieldNext = true - race.Release(unsafe.Pointer(&racer)) + pull.yieldNext = true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) // Propagate panics and goexits from seq. - if panicValue != nil { - if panicValue == goexitPanicValue { + if pull.panicValue != nil { + if pull.panicValue == goexitPanicValue { // Propagate runtime.Goexit from seq. runtime.Goexit() } else { - panic(panicValue) + panic(pull.panicValue) } } - return v, ok + return pull.v, pull.ok } stop = func() { - race.Write(unsafe.Pointer(&racer)) // detect races + race.Write(unsafe.Pointer(&pull.racer)) // detect races - if !done { - done = true - race.Release(unsafe.Pointer(&racer)) + if !pull.done { + pull.done = true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) // Propagate panics and goexits from seq. - if panicValue != nil { - if panicValue == goexitPanicValue { + if pull.panicValue != nil { + if pull.panicValue == goexitPanicValue { // Propagate runtime.Goexit from seq. runtime.Goexit() } else { - panic(panicValue) + panic(pull.panicValue) } } } @@ -372,93 +372,93 @@ func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) { // If the iterator panics during a call to next (or stop), // then next (or stop) itself panics with the same value. func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func()) { - var ( + var pull struct { k K v V ok bool done bool yieldNext bool + seqDone bool racer int panicValue any - seqDone bool - ) + } c := newcoro(func(c *coro) { - race.Acquire(unsafe.Pointer(&racer)) - if done { - race.Release(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) + if pull.done { + race.Release(unsafe.Pointer(&pull.racer)) return } yield := func(k1 K, v1 V) bool { - if done { + if pull.done { return false } - if !yieldNext { + if !pull.yieldNext { panic("iter.Pull2: yield called again before next") } - yieldNext = false - k, v, ok = k1, v1, true - race.Release(unsafe.Pointer(&racer)) + pull.yieldNext = false + pull.k, pull.v, pull.ok = k1, v1, true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) - return !done + race.Acquire(unsafe.Pointer(&pull.racer)) + return !pull.done } // Recover and propagate panics from seq. defer func() { if p := recover(); p != nil { - panicValue = p - } else if !seqDone { - panicValue = goexitPanicValue + pull.panicValue = p + } else if !pull.seqDone { + pull.panicValue = goexitPanicValue } - done = true // Invalidate iterator. - race.Release(unsafe.Pointer(&racer)) + pull.done = true // Invalidate iterator. + race.Release(unsafe.Pointer(&pull.racer)) }() seq(yield) var k0 K var v0 V - k, v, ok = k0, v0, false - seqDone = true + pull.k, pull.v, pull.ok = k0, v0, false + pull.seqDone = true }) next = func() (k1 K, v1 V, ok1 bool) { - race.Write(unsafe.Pointer(&racer)) // detect races + race.Write(unsafe.Pointer(&pull.racer)) // detect races - if done { + if pull.done { return } - if yieldNext { + if pull.yieldNext { panic("iter.Pull2: next called again before yield") } - yieldNext = true - race.Release(unsafe.Pointer(&racer)) + pull.yieldNext = true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) // Propagate panics and goexits from seq. - if panicValue != nil { - if panicValue == goexitPanicValue { + if pull.panicValue != nil { + if pull.panicValue == goexitPanicValue { // Propagate runtime.Goexit from seq. runtime.Goexit() } else { - panic(panicValue) + panic(pull.panicValue) } } - return k, v, ok + return pull.k, pull.v, pull.ok } stop = func() { - race.Write(unsafe.Pointer(&racer)) // detect races + race.Write(unsafe.Pointer(&pull.racer)) // detect races - if !done { - done = true - race.Release(unsafe.Pointer(&racer)) + if !pull.done { + pull.done = true + race.Release(unsafe.Pointer(&pull.racer)) coroswitch(c) - race.Acquire(unsafe.Pointer(&racer)) + race.Acquire(unsafe.Pointer(&pull.racer)) // Propagate panics and goexits from seq. - if panicValue != nil { - if panicValue == goexitPanicValue { + if pull.panicValue != nil { + if pull.panicValue == goexitPanicValue { // Propagate runtime.Goexit from seq. runtime.Goexit() } else { - panic(panicValue) + panic(pull.panicValue) } } } diff --git a/src/iter/pull_test.go b/src/iter/pull_test.go index e9e3bdadca099354f57bb17a3110b4431dd3949a..c66e20897bd8af05596c08448e8214bad7c64407 100644 --- a/src/iter/pull_test.go +++ b/src/iter/pull_test.go @@ -491,3 +491,19 @@ func TestPull2ImmediateStop(t *testing.T) { t.Fatal("next returned true after iterator was stopped") } } + +func BenchmarkPull(b *testing.B) { + seq := count(1) + for range b.N { + _, stop := Pull(seq) + stop() + } +} + +func BenchmarkPull2(b *testing.B) { + seq := squares(1) + for range b.N { + _, stop := Pull2(seq) + stop() + } +}