iter.go

     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  // Package iter provides basic definitions and operations
     6  // related to iteration in Go.
     7  package iter
     8  
     9  import (
    10  	"internal/race"
    11  	"unsafe"
    12  )
    13  
    14  // Seq is an iterator over sequences of individual values.
    15  // When called as seq(yield), seq calls yield(v) for each value v in the sequence,
    16  // stopping early if yield returns false.
    17  type Seq[V any] func(yield func(V) bool)
    18  
    19  // Seq2 is an iterator over sequences of pairs of values, most commonly key-value pairs.
    20  // When called as seq(yield), seq calls yield(k, v) for each pair (k, v) in the sequence,
    21  // stopping early if yield returns false.
    22  type Seq2[K, V any] func(yield func(K, V) bool)
    23  
    24  type coro struct{}
    25  
    26  //go:linkname newcoro runtime.newcoro
    27  func newcoro(func(*coro)) *coro
    28  
    29  //go:linkname coroswitch runtime.coroswitch
    30  func coroswitch(*coro)
    31  
    32  // Pull converts the “push-style” iterator sequence seq
    33  // into a “pull-style” iterator accessed by the two functions
    34  // next and stop.
    35  //
    36  // Next returns the next value in the sequence
    37  // and a boolean indicating whether the value is valid.
    38  // When the sequence is over, next returns the zero V and false.
    39  // It is valid to call next after reaching the end of the sequence
    40  // or after calling stop. These calls will continue
    41  // to return the zero V and false.
    42  //
    43  // Stop ends the iteration. It must be called when the caller is
    44  // no longer interested in next values and next has not yet
    45  // signaled that the sequence is over (with a false boolean return).
    46  // It is valid to call stop multiple times and when next has
    47  // already returned false.
    48  //
    49  // It is an error to call next or stop from multiple goroutines
    50  // simultaneously.
    51  func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) {
    52  	var (
    53  		v         V
    54  		ok        bool
    55  		done      bool
    56  		yieldNext bool
    57  		racer     int
    58  	)
    59  	c := newcoro(func(c *coro) {
    60  		race.Acquire(unsafe.Pointer(&racer))
    61  		yield := func(v1 V) bool {
    62  			if done {
    63  				return false
    64  			}
    65  			if !yieldNext {
    66  				panic("iter.Pull: yield called again before next")
    67  			}
    68  			yieldNext = false
    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  		if yieldNext {
    87  			panic("iter.Pull: next called again before yield")
    88  		}
    89  		yieldNext = true
    90  		race.Release(unsafe.Pointer(&racer))
    91  		coroswitch(c)
    92  		race.Acquire(unsafe.Pointer(&racer))
    93  		return v, ok
    94  	}
    95  	stop = func() {
    96  		race.Write(unsafe.Pointer(&racer)) // detect races
    97  		if !done {
    98  			done = true
    99  			race.Release(unsafe.Pointer(&racer))
   100  			coroswitch(c)
   101  			race.Acquire(unsafe.Pointer(&racer))
   102  		}
   103  	}
   104  	return next, stop
   105  }
   106  
   107  // Pull2 converts the “push-style” iterator sequence seq
   108  // into a “pull-style” iterator accessed by the two functions
   109  // next and stop.
   110  //
   111  // Next returns the next pair in the sequence
   112  // and a boolean indicating whether the pair is valid.
   113  // When the sequence is over, next returns a pair of zero values and false.
   114  // It is valid to call next after reaching the end of the sequence
   115  // or after calling stop. These calls will continue
   116  // to return a pair of zero values and false.
   117  //
   118  // Stop ends the iteration. It must be called when the caller is
   119  // no longer interested in next values and next has not yet
   120  // signaled that the sequence is over (with a false boolean return).
   121  // It is valid to call stop multiple times and when next has
   122  // already returned false.
   123  //
   124  // It is an error to call next or stop from multiple goroutines
   125  // simultaneously.
   126  func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func()) {
   127  	var (
   128  		k         K
   129  		v         V
   130  		ok        bool
   131  		done      bool
   132  		yieldNext bool
   133  		racer     int
   134  	)
   135  	c := newcoro(func(c *coro) {
   136  		race.Acquire(unsafe.Pointer(&racer))
   137  		yield := func(k1 K, v1 V) bool {
   138  			if done {
   139  				return false
   140  			}
   141  			if !yieldNext {
   142  				panic("iter.Pull2: yield called again before next")
   143  			}
   144  			yieldNext = false
   145  			k, v, ok = k1, v1, true
   146  			race.Release(unsafe.Pointer(&racer))
   147  			coroswitch(c)
   148  			race.Acquire(unsafe.Pointer(&racer))
   149  			return !done
   150  		}
   151  		seq(yield)
   152  		var k0 K
   153  		var v0 V
   154  		k, v, ok = k0, v0, false
   155  		done = true
   156  		race.Release(unsafe.Pointer(&racer))
   157  	})
   158  	next = func() (k1 K, v1 V, ok1 bool) {
   159  		race.Write(unsafe.Pointer(&racer)) // detect races
   160  		if done {
   161  			return
   162  		}
   163  		if yieldNext {
   164  			panic("iter.Pull2: next called again before yield")
   165  		}
   166  		yieldNext = true
   167  		race.Release(unsafe.Pointer(&racer))
   168  		coroswitch(c)
   169  		race.Acquire(unsafe.Pointer(&racer))
   170  		return k, v, ok
   171  	}
   172  	stop = func() {
   173  		race.Write(unsafe.Pointer(&racer)) // detect races
   174  		if !done {
   175  			done = true
   176  			race.Release(unsafe.Pointer(&racer))
   177  			coroswitch(c)
   178  			race.Acquire(unsafe.Pointer(&racer))
   179  		}
   180  	}
   181  	return next, stop
   182  }
   183
View as plain text