go-primes-counter/main.go

package main

import (
	"flag"
	"fmt"
	"sync"
)

func main() {
	start, max, step, printCount := parseArgs()
	if max < start {
		start, max = max, start
	}
	totalLength := max - start
	if totalLength < step {
		step = totalLength
	}

	// Start jobs
	var wg sync.WaitGroup
	results := make(chan int, 100)
	for i := start; i < max; i += step {
		wg.Add(1)
		go worker(i, i+step, &wg, results)
	}

	// Collect results
	go waitForWorkers(&wg, results)
	var primes []int
	for prime := range results {
		primes = append(primes, prime)
	}

	// Print results
	fmt.Println("There are", len(primes), "primes between", start, "and", max)
	if printCount > len(primes) {
		printCount = -1
	}
	if printCount < 0 {
		fmt.Println("Here they are:")
		fmt.Println(primes)
	} else {
		fmt.Println("Here's a free sample:")
		fmt.Println(primes[:printCount])
	}
}

func worker(start, end int, wg *sync.WaitGroup, results chan<- int) {
	defer wg.Done()
	for i := start; i < end; i++ {
		if IsPrime(i) {
			results <- i
		}
	}
}

func waitForWorkers(wg *sync.WaitGroup, results chan int) {
	wg.Wait()
	close(results)
}

// IsPrime determines if the given int is a prime number or not.
// The method currently implemented by this method is not smart and not
// optimized for large numbers.
func IsPrime(n int) bool {
	if n <= 3 {
		return n > 1
	}
	if n%2 == 0 || n%3 == 0 {
		return false
	}
	i := 5
	for i*i <= n {
		if n%i == 0 || n%(i+2) == 0 {
			return false
		}
		i += 6
	}
	return true
}

func parseArgs() (start, max, step, printCount int) {
	flag.IntVar(&start, "start", 0, "The start number")
	flag.IntVar(&max, "max", 10_000, "The end value")
	flag.IntVar(&step, "step-size", 1_000, "The job size")
	printCountDesc := `Number of primes to print. -1 to print all.
Primes are not guaranteed to be sorted.`
	flag.IntVar(&printCount, "print-count", 20, printCountDesc)
	flag.Parse()
	return
}