Source file src/iter/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  /*
     6  Package iter provides basic definitions and operations related to
     7  iterators over sequences.
     8  
     9  # Iterators
    10  
    11  An iterator is a function that passes successive elements of a
    12  sequence to a callback function, conventionally named yield.
    13  The function stops either when the sequence is finished or
    14  when yield returns false, indicating to stop the iteration early.
    15  This package defines [Seq] and [Seq2]
    16  (pronounced like seek—the first syllable of sequence)
    17  as shorthands for iterators that pass 1 or 2 values per sequence element
    18  to yield:
    19  
    20  	type (
    21  		Seq[V any]     func(yield func(V) bool)
    22  		Seq2[K, V any] func(yield func(K, V) bool)
    23  	)
    24  
    25  Seq2 represents a sequence of paired values, conventionally key-value
    26  or index-value pairs.
    27  
    28  Yield returns true if the iterator should continue with the next
    29  element in the sequence, false if it should stop.
    30  
    31  For instance, [maps.Keys] returns an iterator that produces the sequence
    32  of keys of the map m, implemented as follows:
    33  
    34  	func Keys[Map ~map[K]V, K comparable, V any](m Map) iter.Seq[K] {
    35  		return func(yield func(K) bool) {
    36  			for k := range m {
    37  				if !yield(k) {
    38  					return
    39  				}
    40  			}
    41  		}
    42  	}
    43  
    44  Further examples can be found in [The Go Blog: Range Over Function Types].
    45  
    46  Iterator functions are most often called by a [range loop], as in:
    47  
    48  	func PrintAll[V any](seq iter.Seq[V]) {
    49  		for v := range seq {
    50  			fmt.Println(v)
    51  		}
    52  	}
    53  
    54  # Naming Conventions
    55  
    56  Iterator functions and methods are named for the sequence being walked:
    57  
    58  	// All returns an iterator over all elements in s.
    59  	func (s *Set[V]) All() iter.Seq[V]
    60  
    61  The iterator method on a collection type is conventionally named All,
    62  because it iterates a sequence of all the values in the collection.
    63  
    64  For a type containing multiple possible sequences, the iterator's name
    65  can indicate which sequence is being provided:
    66  
    67  	// Cities returns an iterator over the major cities in the country.
    68  	func (c *Country) Cities() iter.Seq[*City]
    69  
    70  	// Languages returns an iterator over the official spoken languages of the country.
    71  	func (c *Country) Languages() iter.Seq[string]
    72  
    73  If an iterator requires additional configuration, the constructor function
    74  can take additional configuration arguments:
    75  
    76  	// Scan returns an iterator over key-value pairs with min ≤ key ≤ max.
    77  	func (m *Map[K, V]) Scan(min, max K) iter.Seq2[K, V]
    78  
    79  	// Split returns an iterator over the (possibly-empty) substrings of s
    80  	// separated by sep.
    81  	func Split(s, sep string) iter.Seq[string]
    82  
    83  When there are multiple possible iteration orders, the method name may
    84  indicate that order:
    85  
    86  	// All returns an iterator over the list from head to tail.
    87  	func (l *List[V]) All() iter.Seq[V]
    88  
    89  	// Backward returns an iterator over the list from tail to head.
    90  	func (l *List[V]) Backward() iter.Seq[V]
    91  
    92  	// Preorder returns an iterator over all nodes of the syntax tree
    93  	// beneath (and including) the specified root, in depth-first preorder,
    94  	// visiting a parent node before its children.
    95  	func Preorder(root Node) iter.Seq[Node]
    96  
    97  # Single-Use Iterators
    98  
    99  Most iterators provide the ability to walk an entire sequence:
   100  when called, the iterator does any setup necessary to start the
   101  sequence, then calls yield on successive elements of the sequence,
   102  and then cleans up before returning. Calling the iterator again
   103  walks the sequence again.
   104  
   105  Some iterators break that convention, providing the ability to walk a
   106  sequence only once. These “single-use iterators” typically report values
   107  from a data stream that cannot be rewound to start over.
   108  Calling the iterator again after stopping early may continue the
   109  stream, but calling it again after the sequence is finished will yield
   110  no values at all. Doc comments for functions or methods that return
   111  single-use iterators should document this fact:
   112  
   113  	// Lines returns an iterator over lines read from r.
   114  	// It returns a single-use iterator.
   115  	func (r *Reader) Lines() iter.Seq[string]
   116  
   117  # Pulling Values
   118  
   119  Functions and methods that accept or return iterators
   120  should use the standard [Seq] or [Seq2] types, to ensure
   121  compatibility with range loops and other iterator adapters.
   122  The standard iterators can be thought of as “push iterators”, which
   123  push values to the yield function.
   124  
   125  Sometimes a range loop is not the most natural way to consume values
   126  of the sequence. In this case, [Pull] converts a standard push iterator
   127  to a “pull iterator”, which can be called to pull one value at a time
   128  from the sequence. [Pull] starts an iterator and returns a pair
   129  of functions—next and stop—which return the next value from the iterator
   130  and stop it, respectively.
   131  
   132  For example:
   133  
   134  	// Pairs returns an iterator over successive pairs of values from seq.
   135  	func Pairs[V any](seq iter.Seq[V]) iter.Seq2[V, V] {
   136  		return func(yield func(V, V) bool) {
   137  			next, stop := iter.Pull(seq)
   138  			defer stop()
   139  			for {
   140  				v1, ok1 := next()
   141  				if !ok1 {
   142  					return
   143  				}
   144  				v2, ok2 := next()
   145  				// If ok2 is false, v2 should be the
   146  				// zero value; yield one last pair.
   147  				if !yield(v1, v2) {
   148  					return
   149  				}
   150  				if !ok2 {
   151  					return
   152  				}
   153  			}
   154  		}
   155  	}
   156  
   157  If clients do not consume the sequence to completion, they must call stop,
   158  which allows the iterator function to finish and return. As shown in
   159  the example, the conventional way to ensure this is to use defer.
   160  
   161  # Standard Library Usage
   162  
   163  A few packages in the standard library provide iterator-based APIs,
   164  most notably the [maps] and [slices] packages.
   165  For example, [maps.Keys] returns an iterator over the keys of a map,
   166  while [slices.Sorted] collects the values of an iterator into a slice,
   167  sorts them, and returns the slice, so to iterate over the sorted keys of a map:
   168  
   169  	for _, key := range slices.Sorted(maps.Keys(m)) {
   170  		...
   171  	}
   172  
   173  # Mutation
   174  
   175  Iterators provide only the values of the sequence, not any direct way
   176  to modify it. If an iterator wishes to provide a mechanism for modifying
   177  a sequence during iteration, the usual approach is to define a position type
   178  with the extra operations and then provide an iterator over positions.
   179  
   180  For example, a tree implementation might provide:
   181  
   182  	// Positions returns an iterator over positions in the sequence.
   183  	func (t *Tree[V]) Positions() iter.Seq[*Pos]
   184  
   185  	// A Pos represents a position in the sequence.
   186  	// It is only valid during the yield call it is passed to.
   187  	type Pos[V any] struct { ... }
   188  
   189  	// Pos returns the value at the cursor.
   190  	func (p *Pos[V]) Value() V
   191  
   192  	// Delete deletes the value at this point in the iteration.
   193  	func (p *Pos[V]) Delete()
   194  
   195  	// Set changes the value v at the cursor.
   196  	func (p *Pos[V]) Set(v V)
   197  
   198  And then a client could delete boring values from the tree using:
   199  
   200  	for p := range t.Positions() {
   201  		if boring(p.Value()) {
   202  			p.Delete()
   203  		}
   204  	}
   205  
   206  [The Go Blog: Range Over Function Types]: https://go.dev/blog/range-functions
   207  [range loop]: https://go.dev/ref/spec#For_range
   208  */
   209  package iter
   210  
   211  import (
   212  	"internal/race"
   213  	"runtime"
   214  	"unsafe"
   215  )
   216  
   217  // Seq is an iterator over sequences of individual values.
   218  // When called as seq(yield), seq calls yield(v) for each value v in the sequence,
   219  // stopping early if yield returns false.
   220  // See the [iter] package documentation for more details.
   221  type Seq[V any] func(yield func(V) bool)
   222  
   223  // Seq2 is an iterator over sequences of pairs of values, most commonly key-value pairs.
   224  // When called as seq(yield), seq calls yield(k, v) for each pair (k, v) in the sequence,
   225  // stopping early if yield returns false.
   226  // See the [iter] package documentation for more details.
   227  type Seq2[K, V any] func(yield func(K, V) bool)
   228  
   229  type coro struct{}
   230  
   231  //go:linkname newcoro runtime.newcoro
   232  func newcoro(func(*coro)) *coro
   233  
   234  //go:linkname coroswitch runtime.coroswitch
   235  func coroswitch(*coro)
   236  
   237  // Pull converts the “push-style” iterator sequence seq
   238  // into a “pull-style” iterator accessed by the two functions
   239  // next and stop.
   240  //
   241  // Next returns the next value in the sequence
   242  // and a boolean indicating whether the value is valid.
   243  // When the sequence is over, next returns the zero V and false.
   244  // It is valid to call next after reaching the end of the sequence
   245  // or after calling stop. These calls will continue
   246  // to return the zero V and false.
   247  //
   248  // Stop ends the iteration. It must be called when the caller is
   249  // no longer interested in next values and next has not yet
   250  // signaled that the sequence is over (with a false boolean return).
   251  // It is valid to call stop multiple times and when next has
   252  // already returned false. Typically, callers should “defer stop()”.
   253  //
   254  // It is an error to call next or stop from multiple goroutines
   255  // simultaneously.
   256  //
   257  // If the iterator panics during a call to next (or stop),
   258  // then next (or stop) itself panics with the same value.
   259  func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) {
   260  	var (
   261  		v          V
   262  		ok         bool
   263  		done       bool
   264  		yieldNext  bool
   265  		racer      int
   266  		panicValue any
   267  		seqDone    bool // to detect Goexit
   268  	)
   269  	c := newcoro(func(c *coro) {
   270  		race.Acquire(unsafe.Pointer(&racer))
   271  		if done {
   272  			race.Release(unsafe.Pointer(&racer))
   273  			return
   274  		}
   275  		yield := func(v1 V) bool {
   276  			if done {
   277  				return false
   278  			}
   279  			if !yieldNext {
   280  				panic("iter.Pull: yield called again before next")
   281  			}
   282  			yieldNext = false
   283  			v, ok = v1, true
   284  			race.Release(unsafe.Pointer(&racer))
   285  			coroswitch(c)
   286  			race.Acquire(unsafe.Pointer(&racer))
   287  			return !done
   288  		}
   289  		// Recover and propagate panics from seq.
   290  		defer func() {
   291  			if p := recover(); p != nil {
   292  				panicValue = p
   293  			} else if !seqDone {
   294  				panicValue = goexitPanicValue
   295  			}
   296  			done = true // Invalidate iterator
   297  			race.Release(unsafe.Pointer(&racer))
   298  		}()
   299  		seq(yield)
   300  		var v0 V
   301  		v, ok = v0, false
   302  		seqDone = true
   303  	})
   304  	next = func() (v1 V, ok1 bool) {
   305  		race.Write(unsafe.Pointer(&racer)) // detect races
   306  
   307  		if done {
   308  			return
   309  		}
   310  		if yieldNext {
   311  			panic("iter.Pull: next called again before yield")
   312  		}
   313  		yieldNext = true
   314  		race.Release(unsafe.Pointer(&racer))
   315  		coroswitch(c)
   316  		race.Acquire(unsafe.Pointer(&racer))
   317  
   318  		// Propagate panics and goexits from seq.
   319  		if panicValue != nil {
   320  			if panicValue == goexitPanicValue {
   321  				// Propagate runtime.Goexit from seq.
   322  				runtime.Goexit()
   323  			} else {
   324  				panic(panicValue)
   325  			}
   326  		}
   327  		return v, ok
   328  	}
   329  	stop = func() {
   330  		race.Write(unsafe.Pointer(&racer)) // detect races
   331  
   332  		if !done {
   333  			done = true
   334  			race.Release(unsafe.Pointer(&racer))
   335  			coroswitch(c)
   336  			race.Acquire(unsafe.Pointer(&racer))
   337  
   338  			// Propagate panics and goexits from seq.
   339  			if panicValue != nil {
   340  				if panicValue == goexitPanicValue {
   341  					// Propagate runtime.Goexit from seq.
   342  					runtime.Goexit()
   343  				} else {
   344  					panic(panicValue)
   345  				}
   346  			}
   347  		}
   348  	}
   349  	return next, stop
   350  }
   351  
   352  // Pull2 converts the “push-style” iterator sequence seq
   353  // into a “pull-style” iterator accessed by the two functions
   354  // next and stop.
   355  //
   356  // Next returns the next pair in the sequence
   357  // and a boolean indicating whether the pair is valid.
   358  // When the sequence is over, next returns a pair of zero values and false.
   359  // It is valid to call next after reaching the end of the sequence
   360  // or after calling stop. These calls will continue
   361  // to return a pair of zero values and false.
   362  //
   363  // Stop ends the iteration. It must be called when the caller is
   364  // no longer interested in next values and next has not yet
   365  // signaled that the sequence is over (with a false boolean return).
   366  // It is valid to call stop multiple times and when next has
   367  // already returned false. Typically, callers should “defer stop()”.
   368  //
   369  // It is an error to call next or stop from multiple goroutines
   370  // simultaneously.
   371  //
   372  // If the iterator panics during a call to next (or stop),
   373  // then next (or stop) itself panics with the same value.
   374  func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func()) {
   375  	var (
   376  		k          K
   377  		v          V
   378  		ok         bool
   379  		done       bool
   380  		yieldNext  bool
   381  		racer      int
   382  		panicValue any
   383  		seqDone    bool
   384  	)
   385  	c := newcoro(func(c *coro) {
   386  		race.Acquire(unsafe.Pointer(&racer))
   387  		if done {
   388  			race.Release(unsafe.Pointer(&racer))
   389  			return
   390  		}
   391  		yield := func(k1 K, v1 V) bool {
   392  			if done {
   393  				return false
   394  			}
   395  			if !yieldNext {
   396  				panic("iter.Pull2: yield called again before next")
   397  			}
   398  			yieldNext = false
   399  			k, v, ok = k1, v1, true
   400  			race.Release(unsafe.Pointer(&racer))
   401  			coroswitch(c)
   402  			race.Acquire(unsafe.Pointer(&racer))
   403  			return !done
   404  		}
   405  		// Recover and propagate panics from seq.
   406  		defer func() {
   407  			if p := recover(); p != nil {
   408  				panicValue = p
   409  			} else if !seqDone {
   410  				panicValue = goexitPanicValue
   411  			}
   412  			done = true // Invalidate iterator.
   413  			race.Release(unsafe.Pointer(&racer))
   414  		}()
   415  		seq(yield)
   416  		var k0 K
   417  		var v0 V
   418  		k, v, ok = k0, v0, false
   419  		seqDone = true
   420  	})
   421  	next = func() (k1 K, v1 V, ok1 bool) {
   422  		race.Write(unsafe.Pointer(&racer)) // detect races
   423  
   424  		if done {
   425  			return
   426  		}
   427  		if yieldNext {
   428  			panic("iter.Pull2: next called again before yield")
   429  		}
   430  		yieldNext = true
   431  		race.Release(unsafe.Pointer(&racer))
   432  		coroswitch(c)
   433  		race.Acquire(unsafe.Pointer(&racer))
   434  
   435  		// Propagate panics and goexits from seq.
   436  		if panicValue != nil {
   437  			if panicValue == goexitPanicValue {
   438  				// Propagate runtime.Goexit from seq.
   439  				runtime.Goexit()
   440  			} else {
   441  				panic(panicValue)
   442  			}
   443  		}
   444  		return k, v, ok
   445  	}
   446  	stop = func() {
   447  		race.Write(unsafe.Pointer(&racer)) // detect races
   448  
   449  		if !done {
   450  			done = true
   451  			race.Release(unsafe.Pointer(&racer))
   452  			coroswitch(c)
   453  			race.Acquire(unsafe.Pointer(&racer))
   454  
   455  			// Propagate panics and goexits from seq.
   456  			if panicValue != nil {
   457  				if panicValue == goexitPanicValue {
   458  					// Propagate runtime.Goexit from seq.
   459  					runtime.Goexit()
   460  				} else {
   461  					panic(panicValue)
   462  				}
   463  			}
   464  		}
   465  	}
   466  	return next, stop
   467  }
   468  
   469  // goexitPanicValue is a sentinel value indicating that an iterator
   470  // exited via runtime.Goexit.
   471  var goexitPanicValue any = new(int)
   472  

View as plain text