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// Tests coroutine switches.
6
7//go:build ignore
8
9package main
10
11import (
12	"iter"
13	"log"
14	"os"
15	"runtime/trace"
16	"sync"
17)
18
19func main() {
20	// Start tracing.
21	if err := trace.Start(os.Stdout); err != nil {
22		log.Fatalf("failed to start tracing: %v", err)
23	}
24
25	// Try simple pull iteration.
26	i := pullRange(100)
27	for {
28		_, ok := i.next()
29		if !ok {
30			break
31		}
32	}
33
34	// Try bouncing the pull iterator between two goroutines.
35	var wg sync.WaitGroup
36	var iterChans [2]chan intIter
37	wg.Add(2)
38	iterChans[0] = make(chan intIter)
39	iterChans[1] = make(chan intIter)
40	go func() {
41		defer wg.Done()
42
43		iter := pullRange(100)
44		iterChans[1] <- iter
45
46		for i := range iterChans[0] {
47			_, ok := i.next()
48			if !ok {
49				close(iterChans[1])
50				break
51			}
52			iterChans[1] <- i
53		}
54	}()
55	go func() {
56		defer wg.Done()
57
58		for i := range iterChans[1] {
59			_, ok := i.next()
60			if !ok {
61				close(iterChans[0])
62				break
63			}
64			iterChans[0] <- i
65		}
66	}()
67	wg.Wait()
68
69	// End of traced execution.
70	trace.Stop()
71}
72
73func pullRange(n int) intIter {
74	next, stop := iter.Pull(func(yield func(v int) bool) {
75		for i := range n {
76			yield(i)
77		}
78	})
79	return intIter{next: next, stop: stop}
80}
81
82type intIter struct {
83	next func() (int, bool)
84	stop func()
85}
86