Property-based testing, in Go! [go] [golang]

gen.go

package main

import (
	"fmt"
)

type RNG interface {
	Next() (uint64, RNG)
}

type simpleRNG struct {
	seed uint64
}

func (s simpleRNG) Next() (uint64, RNG) {
	return s.seed, simpleRNG{
		seed: s.seed*0x5DEECE66D + 0xB,
	}
}

type Gen[A any] interface {
	Sample(RNG) (A, RNG)
}

type GenFunc[A any] func(RNG) (A, RNG)

func (f GenFunc[A]) Sample(r RNG) (A, RNG) {
	return f(r)
}

// Generator combinators

func Map[A, B any](g Gen[A], f func(A) B) Gen[B] {
	return GenFunc[B](func(r RNG) (B, RNG) {
		s, nr := g.Sample(r)
		return f(s), nr
	})
}

func Unit[A any](a A) Gen[A] {
	return GenFunc[A](func(r RNG) (A, RNG) {
		return a, r
	})
}

func Map2[A, B, C any](ga Gen[A], gb Gen[B], f func(A, B) C) Gen[C] {
	return GenFunc[C](func(r RNG) (C, RNG) {
		s1, nr1 := ga.Sample(r)
		s2, nr2 := gb.Sample(nr1)
		return f(s1, s2), nr2
	})
}

func Apply[A, B any](gf Gen[func(A) B], ga Gen[A]) Gen[B] {
	return Map2(gf, ga, func(f func(A) B, a A) B {
		return f(a)
	})
}

func FlatMap[A, B any](ga Gen[A], f func(A) Gen[B]) Gen[B] {
	return GenFunc[B](func(r RNG) (B, RNG) {
		s, nr := ga.Sample(r)
		return f(s).Sample(nr)
	})
}

func Traverse[A, B any](as []A, f func(A) Gen[B]) Gen[[]B] {
	return GenFunc[[]B](func(r RNG) ([]B, RNG) {
		res := make([]B, len(as))
		rng := r
		for i, a := range as {
			res[i], rng = f(a).Sample(rng)
		}
		return res, rng
	})
}

func Sequence[A any](gs []Gen[A]) Gen[[]A] {
	return Traverse(gs, func(ga Gen[A]) Gen[A] { return ga })
}

func Fill[A any](n int, f func() A) []A {
	res := make([]A, n)
	for i := 0; i < n; i++ {
		res[i] = f()
	}
	return res
}

func SliceOfN[A any](g Gen[A], n int) Gen[[]A] {
	return Sequence(Fill(n, func() Gen[A] { return g }))
}

// Generators

func Choose(min, max int) Gen[int] {
	return GenFunc[int](func(r RNG) (int, RNG) {
		s, nr := r.Next()
		return min + int(s&(1<<32-1))%(max-min), nr
	})
}

func OneOf[A any](gs ...Gen[A]) Gen[A] {
	return FlatMap(Choose(0, len(gs)), func(n int) Gen[A] {
		return gs[n%len(gs)]
	})
}

func Boolean() Gen[bool] {
	return GenFunc[bool](func(r RNG) (bool, RNG) {
		n, nr := r.Next()
		return n%2 == 0, nr
	})
}

func Byte() Gen[byte] {
	return GenFunc[byte](func(r RNG) (byte, RNG) {
		n, nr := r.Next()
		return byte(n), nr
	})
}

func StringN(n int) Gen[string] {
	return Map(SliceOfN(Byte(), n), func(bs []byte) string {
		return string(bs)
	})
}

func main() {
	fmt.Println(SliceOfN(OneOf(Unit("a"), Unit("b"), Unit("c")), 100).Sample(simpleRNG{0}))
}