sumprimes.go

package main

import (
    "fmt"
    "os"
    "strconv"
    "math"
    "runtime"
)

var workers = runtime.NumCPU()

func isPrime(n int) bool {
    if n == 1 || n == 2 {
        return true
    }

    if math.Mod(float64(n), 2) == 0 {
        return false
    }

    for i := 3.0; i <= math.Floor(math.Sqrt(float64(n))); i += 2.0 {
        if math.Mod(float64(n), i) == 0 {
            return false
        }
    }

    return true
}

type Job struct {
    n int
    result chan<-int
}

func (job *Job) Do() {
    if isPrime(job.n) {
        job.result <- job.n
    }
}

func sumPrimesUpto(n int) int {
    jobs := make(chan Job, workers)
    results := make(chan int, n)
    done := make(chan struct{}, workers)

    go addJobs(jobs, results, n)

    for i := 0; i < workers; i++ {
        go doJobs(done, jobs)
    }

    go wait(done, results)

    return tally(results)
}

func tally(results <-chan int) int {
    retval := 0
    for result := range results {
        retval += result
    }

    return retval
}

func addJobs(jobs chan<-Job, results chan<-int, n int) {
    for i := 1; i <= n; i++ {
        jobs <- Job{i, results}
    }
    close(jobs)
}

func doJobs(done chan<-struct{}, jobs <-chan Job) {
    for job := range jobs {
        job.Do()
    }
    done <- struct{}{}
}

func wait(done <-chan struct{}, results chan int) {
    for i := 0; i < workers; i++ {
        <-done
    }
    close(results)
}

func main() {
    runtime.GOMAXPROCS(workers)
    fmt.Printf("CPUS=%d\n", workers)
    upperBound, err := strconv.Atoi(os.Args[1])
    if err != nil {
        fmt.Println("Invalid argument.")
        os.Exit(1);
    }

    result := sumPrimesUpto(upperBound)
    fmt.Println(result)
}