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  //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
View as plain text