Rewrite in Go

5 files changed, 163 insertions, 54 deletions

diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..8ae8579
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,5 @@
+module gitlab.com/rhogenson/words
+
+go 1.24.4
+
+require github.com/google/go-cmp v0.7.0
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..40e761a
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,2 @@
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
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)
+	}
+}
diff --git a/words.py b/words.py
deleted file mode 100644
index b3420ab..0000000
--- a/words.py
+++ /dev/null
@@ -1,54 +0,0 @@
-#!/usr/bin/env python
-
-import random
-
-wordList = open("words.txt")
-letterNumerators = [[0 for i in range(27)] for i in range(27)]
-letterDenominators = [0 for i in range(27)]
-wordBeginnings = [0 for i in range(27)]
-wordsNum = 0
-
-for line in wordList:
-    line = line.lower()
-    line = line[:-1]
-    acceptable = True
-    for letter in line:
-        if (ord(letter) < 97 or ord(letter) > 122):
-            acceptable = False
-
-    if len(line) < 0:
-        acceptable = False
-    if acceptable:
-        wordsNum += 1
-        wordBeginnings[ord(line[0]) - 96] += 1
-        for i in range(len(line) - 1):
-            letterDenominators[ord(line[i]) - 96] += 1
-            letterNumerators[ord(line[i]) - 96][ord(line[i + 1]) - 96] += 1
-        letterDenominators[ord(line[-1]) - 96] += 1
-        letterNumerators[ord(line[-1]) - 96][0] += 1
-
-
-while True:
-    newWord = ""
-    j = 0
-    number = random.SystemRandom().random() * wordsNum
-    for i in range(len(wordBeginnings)):
-        number -= wordBeginnings[i]
-        if number < 0:
-            newWord += (chr(i + 96))
-            break
-
-    newWordFinished = False
-    while not newWordFinished:
-        number = letterDenominators[ord(newWord[j]) - 96] * random.SystemRandom().random()
-        for i in range(len(letterNumerators[ord(newWord[j]) - 96])):
-            number -= letterNumerators[ord(newWord[j]) - 96][i]
-            if number < 0:
-                if i == 0:
-                    newWordFinished = True
-                    break
-                j = j + 1
-                newWord += (chr(i + 96))
-                break
-
-    print(newWord)
diff --git a/words_test.go b/words_test.go
new file mode 100644
index 0000000..9c79464
--- /dev/null
+++ b/words_test.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestBuildNGrams(t *testing.T) {
+	const testWordList = `abc
+abd
+def
+`
+
+	got, err := buildNGrams(strings.NewReader(testWordList))
+	if err != nil {
+		t.Fatalf("buildNGrams failed: %s", err)
+	}
+	want := map[biGram][]choice{
+		{-1, -1}:   {{'a', 2}, {'d', 3}},
+		{-1, 'a'}:  {{'b', 2}},
+		{-1, 'd'}:  {{'e', 1}},
+		{'a', 'b'}: {{'c', 1}, {'d', 2}},
+		{'b', 'c'}: {{-1, 1}},
+		{'b', 'd'}: {{-1, 1}},
+		{'d', 'e'}: {{'f', 1}},
+		{'e', 'f'}: {{-1, 1}},
+	}
+	if diff := cmp.Diff(want, got, cmp.AllowUnexported(biGram{}, choice{})); diff != "" {
+		t.Errorf("buildNGrams(%q) returned unexpected diff (-want +got):\n%s", testWordList, diff)
+	}
+}