1// run
2
3// Copyright 2009 The Go Authors. All rights reserved.
4// Use of this source code is governed by a BSD-style
5// license that can be found in the LICENSE file.
6
7// Test concurrency primitives: prime sieve of Eratosthenes.
8
9// Generate primes up to 100 using channels, checking the results.
10// This sieve is Eratosthenesque and only considers odd candidates.
11// See discussion at <http://blog.onideas.ws/eratosthenes.go>.
12
13package main
14
15import (
16	"container/heap"
17	"container/ring"
18)
19
20// Return a chan of odd numbers, starting from 5.
21func odds() chan int {
22	out := make(chan int, 50)
23	go func() {
24		n := 5
25		for {
26			out <- n
27			n += 2
28		}
29	}()
30	return out
31}
32
33// Return a chan of odd multiples of the prime number p, starting from p*p.
34func multiples(p int) chan int {
35	out := make(chan int, 10)
36	go func() {
37		n := p * p
38		for {
39			out <- n
40			n += 2 * p
41		}
42	}()
43	return out
44}
45
46type PeekCh struct {
47	head int
48	ch   chan int
49}
50
51// Heap of PeekCh, sorting by head values, satisfies Heap interface.
52type PeekChHeap []*PeekCh
53
54func (h *PeekChHeap) Less(i, j int) bool {
55	return (*h)[i].head < (*h)[j].head
56}
57
58func (h *PeekChHeap) Swap(i, j int) {
59	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
60}
61
62func (h *PeekChHeap) Len() int {
63	return len(*h)
64}
65
66func (h *PeekChHeap) Pop() (v interface{}) {
67	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
68	return
69}
70
71func (h *PeekChHeap) Push(v interface{}) {
72	*h = append(*h, v.(*PeekCh))
73}
74
75// Return a channel to serve as a sending proxy to 'out'.
76// Use a goroutine to receive values from 'out' and store them
77// in an expanding buffer, so that sending to 'out' never blocks.
78func sendproxy(out chan<- int) chan<- int {
79	proxy := make(chan int, 10)
80	go func() {
81		n := 16 // the allocated size of the circular queue
82		first := ring.New(n)
83		last := first
84		var c chan<- int
85		var e int
86		for {
87			c = out
88			if first == last {
89				// buffer empty: disable output
90				c = nil
91			} else {
92				e = first.Value.(int)
93			}
94			select {
95			case e = <-proxy:
96				last.Value = e
97				if last.Next() == first {
98					// buffer full: expand it
99					last.Link(ring.New(n))
100					n *= 2
101				}
102				last = last.Next()
103			case c <- e:
104				first = first.Next()
105			}
106		}
107	}()
108	return proxy
109}
110
111// Return a chan int of primes.
112func Sieve() chan int {
113	// The output values.
114	out := make(chan int, 10)
115	out <- 2
116	out <- 3
117
118	// The channel of all composites to be eliminated in increasing order.
119	composites := make(chan int, 50)
120
121	// The feedback loop.
122	primes := make(chan int, 10)
123	primes <- 3
124
125	// Merge channels of multiples of 'primes' into 'composites'.
126	go func() {
127		var h PeekChHeap
128		min := 15
129		for {
130			m := multiples(<-primes)
131			head := <-m
132			for min < head {
133				composites <- min
134				minchan := heap.Pop(&h).(*PeekCh)
135				min = minchan.head
136				minchan.head = <-minchan.ch
137				heap.Push(&h, minchan)
138			}
139			for min == head {
140				minchan := heap.Pop(&h).(*PeekCh)
141				min = minchan.head
142				minchan.head = <-minchan.ch
143				heap.Push(&h, minchan)
144			}
145			composites <- head
146			heap.Push(&h, &PeekCh{<-m, m})
147		}
148	}()
149
150	// Sieve out 'composites' from 'candidates'.
151	go func() {
152		// In order to generate the nth prime we only need multiples of
153		// primes ≤ sqrt(nth prime).  Thus, the merging goroutine will
154		// receive from 'primes' much slower than this goroutine
155		// will send to it, making the buffer accumulate and block this
156		// goroutine from sending, causing a deadlock.  The solution is to
157		// use a proxy goroutine to do automatic buffering.
158		primes := sendproxy(primes)
159
160		candidates := odds()
161		p := <-candidates
162
163		for {
164			c := <-composites
165			for p < c {
166				primes <- p
167				out <- p
168				p = <-candidates
169			}
170			if p == c {
171				p = <-candidates
172			}
173		}
174	}()
175
176	return out
177}
178
179func main() {
180	primes := Sieve()
181	a := []int{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97}
182	for i := 0; i < len(a); i++ {
183		if x := <-primes; x != a[i] {
184			println(x, " != ", a[i])
185			panic("fail")
186		}
187	}
188}
189