1 // Copyright 2023 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 //go:build goexperiment.rangefunc 6 7 // Package iter provides basic definitions and operations 8 // related to iteration in Go. 9 // 10 // This package is experimental and can only be imported 11 // when building with GOEXPERIMENT=rangefunc. 12 package iter 13 14 import ( 15 "internal/race" 16 "unsafe" 17 ) 18 19 // Seq is an iterator over sequences of individual values. 20 // When called as seq(yield), seq calls yield(v) for each value v in the sequence, 21 // stopping early if yield returns false. 22 type Seq[V any] func(yield func(V) bool) 23 24 // Seq2 is an iterator over sequences of pairs of values, most commonly key-value pairs. 25 // When called as seq(yield), seq calls yield(k, v) for each pair (k, v) in the sequence, 26 // stopping early if yield returns false. 27 type Seq2[K, V any] func(yield func(K, V) bool) 28 29 type coro struct{} 30 31 //go:linkname newcoro runtime.newcoro 32 func newcoro(func(*coro)) *coro 33 34 //go:linkname coroswitch runtime.coroswitch 35 func coroswitch(*coro) 36 37 // Pull converts the “push-style” iterator sequence seq 38 // into a “pull-style” iterator accessed by the two functions 39 // next and stop. 40 // 41 // Next returns the next value in the sequence 42 // and a boolean indicating whether the value is valid. 43 // When the sequence is over, next returns the zero V and false. 44 // It is valid to call next after reaching the end of the sequence 45 // or after calling stop. These calls will continue 46 // to return the zero V and false. 47 // 48 // Stop ends the iteration. It must be called when the caller is 49 // no longer interested in next values and next has not yet 50 // signaled that the sequence is over (with a false boolean return). 51 // It is valid to call stop multiple times and when next has 52 // already returned false. 53 // 54 // It is an error to call next or stop from multiple goroutines 55 // simultaneously. 56 func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) { 57 var ( 58 v V 59 ok bool 60 done bool 61 racer int 62 ) 63 c := newcoro(func(c *coro) { 64 race.Acquire(unsafe.Pointer(&racer)) 65 yield := func(v1 V) bool { 66 if done { 67 return false 68 } 69 v, ok = v1, true 70 race.Release(unsafe.Pointer(&racer)) 71 coroswitch(c) 72 race.Acquire(unsafe.Pointer(&racer)) 73 return !done 74 } 75 seq(yield) 76 var v0 V 77 v, ok = v0, false 78 done = true 79 race.Release(unsafe.Pointer(&racer)) 80 }) 81 next = func() (v1 V, ok1 bool) { 82 race.Write(unsafe.Pointer(&racer)) // detect races 83 if done { 84 return 85 } 86 race.Release(unsafe.Pointer(&racer)) 87 coroswitch(c) 88 race.Acquire(unsafe.Pointer(&racer)) 89 return v, ok 90 } 91 stop = func() { 92 race.Write(unsafe.Pointer(&racer)) // detect races 93 if !done { 94 done = true 95 race.Release(unsafe.Pointer(&racer)) 96 coroswitch(c) 97 race.Acquire(unsafe.Pointer(&racer)) 98 } 99 } 100 return next, stop 101 } 102 103 // Pull2 converts the “push-style” iterator sequence seq 104 // into a “pull-style” iterator accessed by the two functions 105 // next and stop. 106 // 107 // Next returns the next pair in the sequence 108 // and a boolean indicating whether the pair is valid. 109 // When the sequence is over, next returns a pair of zero values and false. 110 // It is valid to call next after reaching the end of the sequence 111 // or after calling stop. These calls will continue 112 // to return a pair of zero values and false. 113 // 114 // Stop ends the iteration. It must be called when the caller is 115 // no longer interested in next values and next has not yet 116 // signaled that the sequence is over (with a false boolean return). 117 // It is valid to call stop multiple times and when next has 118 // already returned false. 119 // 120 // It is an error to call next or stop from multiple goroutines 121 // simultaneously. 122 func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func()) { 123 var ( 124 k K 125 v V 126 ok bool 127 done bool 128 racer int 129 ) 130 c := newcoro(func(c *coro) { 131 race.Acquire(unsafe.Pointer(&racer)) 132 yield := func(k1 K, v1 V) bool { 133 if done { 134 return false 135 } 136 k, v, ok = k1, v1, true 137 race.Release(unsafe.Pointer(&racer)) 138 coroswitch(c) 139 race.Acquire(unsafe.Pointer(&racer)) 140 return !done 141 } 142 seq(yield) 143 var k0 K 144 var v0 V 145 k, v, ok = k0, v0, false 146 done = true 147 race.Release(unsafe.Pointer(&racer)) 148 }) 149 next = func() (k1 K, v1 V, ok1 bool) { 150 race.Write(unsafe.Pointer(&racer)) // detect races 151 if done { 152 return 153 } 154 race.Release(unsafe.Pointer(&racer)) 155 coroswitch(c) 156 race.Acquire(unsafe.Pointer(&racer)) 157 return k, v, ok 158 } 159 stop = func() { 160 race.Write(unsafe.Pointer(&racer)) // detect races 161 if !done { 162 done = true 163 race.Release(unsafe.Pointer(&racer)) 164 coroswitch(c) 165 race.Acquire(unsafe.Pointer(&racer)) 166 } 167 } 168 return next, stop 169 } 170