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

View as plain text