1 files changed, 123 insertions, 0 deletions
diff --git a/words.go b/words.go
new file mode 100644
index 0000000..cdd1306
--- /dev/null
+++ b/words.go
@@ -0,0 +1,123 @@
+package main
+
+import (
+	"bufio"
+	"cmp"
+	_ "embed"
+	"flag"
+	"fmt"
+	"io"
+	"math/rand/v2"
+	"os"
+	"slices"
+	"strings"
+)
+
+var (
+	wordList = flag.String("w", "", "use `wordlist` for N-grams")
+	numWords = flag.Int("n", 1, "generate `n` words")
+)
+
+type choice struct {
+	r         rune
+	cumWeight int
+}
+
+func pick(choices []choice) rune {
+	if len(choices) == 0 {
+		panic("pick: nothing to choose from")
+	}
+	n, _ := slices.BinarySearchFunc(choices, rand.IntN(choices[len(choices)-1].cumWeight)+1, func(c choice, n int) int {
+		return cmp.Compare(c.cumWeight, n)
+	})
+	return choices[n].r
+}
+
+type biGram struct {
+	a, b rune
+}
+
+func buildNGrams(r io.Reader) (map[biGram][]choice, error) {
+	type triGram struct {
+		a, b, c rune
+	}
+	triGrams := make(map[triGram]int)
+	total := 0
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		line := scanner.Text()
+		var a, b rune = -1, -1
+		for _, r := range line {
+			triGrams[triGram{a, b, r}]++
+			total++
+			a, b = b, r
+		}
+		triGrams[triGram{a, b, -1}]++
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+	triGramsKeys := make([]triGram, 0, len(triGrams))
+	for t := range triGrams {
+		triGramsKeys = append(triGramsKeys, t)
+	}
+	slices.SortFunc(triGramsKeys, func(x, y triGram) int {
+		return cmp.Or(
+			cmp.Compare(x.a, y.a),
+			cmp.Compare(x.b, y.b),
+			cmp.Compare(x.c, y.c))
+	})
+	biGrams := make(map[biGram][]choice)
+	for _, t := range triGramsKeys {
+		choices := biGrams[biGram{t.a, t.b}]
+		currentCumTotal := 0
+		if len(choices) > 0 {
+			currentCumTotal = choices[len(choices)-1].cumWeight
+		}
+		biGrams[biGram{t.a, t.b}] = append(choices, choice{t.c, currentCumTotal + triGrams[t]})
+	}
+	return biGrams, nil
+}
+
+//go:embed words.txt
+var defaultWordList string
+
+func run() error {
+	var wordListReader io.Reader
+	if *wordList != "" {
+		var err error
+		wordListReader, err = os.Open(*wordList)
+		if err != nil {
+			return err
+		}
+	} else {
+		wordListReader = strings.NewReader(defaultWordList)
+	}
+	nGrams, err := buildNGrams(wordListReader)
+	if err != nil {
+		return err
+	}
+	w := bufio.NewWriter(os.Stdout)
+	defer w.Flush()
+	for i := 0; *numWords < 0 || i < *numWords; i++ {
+		var a, b rune = -1, -1
+		for {
+			c := pick(nGrams[biGram{a, b}])
+			if c == -1 {
+				break
+			}
+			w.WriteRune(c)
+			a, b = b, c
+		}
+		w.WriteByte('\n')
+	}
+	return nil
+}
+
+func main() {
+	flag.Parse()
+	if err := run(); err != nil {
+		fmt.Fprintf(os.Stderr, "words: %s\n", err)
+		os.Exit(1)
+	}
+}