aboutsummaryrefslogtreecommitdiff
path: root/src/iter/iter.go
blob: 240df00f7f44d26daa8de8c73ad3f36182b9df08 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build goexperiment.rangefunc

// Package iter provides basic definitions and operations
// related to iteration in Go.
//
// This package is experimental and can only be imported
// when building with GOEXPERIMENT=rangefunc.
package iter

import (
	"internal/race"
	"unsafe"
	_ "unsafe"
) // for linkname

// Seq is an iterator over sequences of individual values.
// When called as seq(yield), seq calls yield(v) for each value v in the sequence,
// stopping early if yield returns false.
type Seq[V any] func(yield func(V) bool)

// Seq2 is an iterator over sequences of pairs of values, most commonly key-value pairs.
// When called as seq(yield), seq calls yield(k, v) for each pair (k, v) in the sequence,
// stopping early if yield returns false.
type Seq2[K, V any] func(yield func(K, V) bool)

type coro struct{}

//go:linkname newcoro runtime.newcoro
func newcoro(func(*coro)) *coro

//go:linkname coroswitch runtime.coroswitch
func coroswitch(*coro)

// Pull converts the “push-style” iterator sequence seq
// into a “pull-style” iterator accessed by the two functions
// next and stop.
//
// Next returns the next value in the sequence
// and a boolean indicating whether the value is valid.
// When the sequence is over, next returns the zero V and false.
// It is valid to call next after reaching the end of the sequence
// or after calling stop. These calls will continue
// to return the zero V and false.
//
// Stop ends the iteration. It must be called when the caller is
// no longer interested in next values and next has not yet
// signaled that the sequence is over (with a false boolean return).
// It is valid to call stop multiple times and when next has
// already returned false.
//
// It is an error to call next or stop from multiple goroutines
// simultaneously.
func Pull[V any](seq Seq[V]) (next func() (V, bool), stop func()) {
	var (
		v     V
		ok    bool
		done  bool
		racer int
	)
	c := newcoro(func(c *coro) {
		race.Acquire(unsafe.Pointer(&racer))
		yield := func(v1 V) bool {
			if done {
				return false
			}
			v, ok = v1, true
			race.Release(unsafe.Pointer(&racer))
			coroswitch(c)
			race.Acquire(unsafe.Pointer(&racer))
			return !done
		}
		seq(yield)
		var v0 V
		v, ok = v0, false
		done = true
		race.Release(unsafe.Pointer(&racer))
	})
	next = func() (v1 V, ok1 bool) {
		race.Write(unsafe.Pointer(&racer)) // detect races
		if done {
			return
		}
		race.Release(unsafe.Pointer(&racer))
		coroswitch(c)
		race.Acquire(unsafe.Pointer(&racer))
		return v, ok
	}
	stop = func() {
		race.Write(unsafe.Pointer(&racer)) // detect races
		if !done {
			done = true
			race.Release(unsafe.Pointer(&racer))
			coroswitch(c)
			race.Acquire(unsafe.Pointer(&racer))
		}
	}
	return next, stop
}

// Pull2 converts the “push-style” iterator sequence seq
// into a “pull-style” iterator accessed by the two functions
// next and stop.
//
// Next returns the next pair in the sequence
// and a boolean indicating whether the pair is valid.
// When the sequence is over, next returns a pair of zero values and false.
// It is valid to call next after reaching the end of the sequence
// or after calling stop. These calls will continue
// to return a pair of zero values and false.
//
// Stop ends the iteration. It must be called when the caller is
// no longer interested in next values and next has not yet
// signaled that the sequence is over (with a false boolean return).
// It is valid to call stop multiple times and when next has
// already returned false.
//
// It is an error to call next or stop from multiple goroutines
// simultaneously.
func Pull2[K, V any](seq Seq2[K, V]) (next func() (K, V, bool), stop func()) {
	var (
		k    K
		v    V
		ok   bool
		done bool
	)
	c := newcoro(func(c *coro) {
		yield := func(k1 K, v1 V) bool {
			if done {
				return false
			}
			k, v, ok = k1, v1, true
			coroswitch(c)
			return !done
		}
		seq(yield)
		var k0 K
		var v0 V
		k, v, ok = k0, v0, false
		done = true
	})
	next = func() (k1 K, v1 V, ok1 bool) {
		race.Write(unsafe.Pointer(&c)) // detect races
		if done {
			return
		}
		coroswitch(c)
		return k, v, ok
	}
	stop = func() {
		race.Write(unsafe.Pointer(&c)) // detect races
		if !done {
			done = true
			coroswitch(c)
		}
	}
	return next, stop
}