Grains in Go - exercism.io solution

main.go

package main

import (
	"fmt"
	"time"
)

/*
 * There once was a wise servant who saved the life of a prince. The king promised to pay whatever the servant could dream up.
 * Knowing that the king loved chess, the servant told the king he would like to have grains of wheat.
 * One grain on the first square of a chess board, with the number of grains doubling on each successive square.
 */

func naiveSquare(number int) uint64 {
	result := uint64(1)
	for i := 2; i <= number; i++ {
		result *= 2
	}
	return result
}

func naiveTotal() uint64 {
	var total uint64
	for i := 1; i <= 64; i++ {
		total += naiveSquare(i)
	}
	return total
}

func square(number int) uint64 {
	return 0b1 << (number - 1)
}

func total() uint64 {
	var total uint64
	for i := 1; i <= 64; i++ {
		total += square(i)
	}
	return total
}

func main() {
	// timeFuncSquare(naiveSquare, 64, "naive square")
	// timeFuncTotal(naiveTotal, "naive total")
	// timeFuncSquare(square, 64, "bit shift square")
	// timeFuncTotal(total, "bit shift total")
	fmt.Println("---- Average perf ----")
	runs := 50
	perfTest(func() time.Duration {
		return timeFuncSquareFoo(naiveSquare, 64)
	}, "naive square", runs)
	perfTest(func() time.Duration {
		return timeFuncSquareFoo(square, 64)
	}, "bit shift square", runs)
	perfTest(func() time.Duration {
		return timeFuncTotalFoo(naiveTotal)
	}, "naive total", runs)
	perfTest(func() time.Duration {
		return timeFuncTotalFoo(total)
	}, "bit shift total", runs)
}

/* Performance testing and timing utility functions */

func timer(start time.Time, name string) {
	elapsed := time.Since(start)
	fmt.Printf("%s took %s\n", name, elapsed)
}

func timeFuncSquare(fn func(int) uint64, arg int, name string) uint64 {
	defer timer(time.Now(), name)
	return fn(arg)
}

func timeFuncTotal(fn func() uint64, name string) uint64 {
	defer timer(time.Now(), name)
	return fn()
}

func perfTest(fn func() time.Duration, name string, n int) {
	results := make([]time.Duration, n)
	for i := 0; i < n; i++ {
		results[i] = fn()
	}
	var sum int64
	for _, result := range results {
		sum += int64(result)
	}
	var meanDuration time.Duration = time.Duration(sum / int64(n))
	fmt.Printf("%s took an average of %s after %d runs\n", name, meanDuration, n)
}

func timeFuncSquareFoo(fn func(int) uint64, arg int) (timeElapsed time.Duration) {
	start := time.Now()
	defer func() {
		timeElapsed = time.Since(start)
	}()
	fn(arg)
	return timeElapsed
}

func timeFuncTotalFoo(fn func() uint64) (timeElapsed time.Duration) {
	start := time.Now()
	defer func() {
		timeElapsed = time.Since(start)
	}()
	fn()
	return timeElapsed
}