sieve2.go

     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  
    13  package main
    14  
    15  import (
    16  	"container/heap"
    17  	"container/ring"
    18  )
    19  
    20  // Return a chan of odd numbers, starting from 5.
    21  func 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.
    34  func 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  
    46  type PeekCh struct {
    47  	head int
    48  	ch   chan int
    49  }
    50  
    51  // Heap of PeekCh, sorting by head values, satisfies Heap interface.
    52  type PeekChHeap []*PeekCh
    53  
    54  func (h *PeekChHeap) Less(i, j int) bool {
    55  	return (*h)[i].head < (*h)[j].head
    56  }
    57  
    58  func (h *PeekChHeap) Swap(i, j int) {
    59  	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
    60  }
    61  
    62  func (h *PeekChHeap) Len() int {
    63  	return len(*h)
    64  }
    65  
    66  func (h *PeekChHeap) Pop() (v interface{}) {
    67  	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
    68  	return
    69  }
    70  
    71  func (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.
    78  func 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.
   112  func 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  
   179  func 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
View as plain text