word list -> BST JSON for a project I was working on

avl.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
)

type node struct {
	left, right  *node
	heightMemo   int
	childrenMemo int
	value        string
}

func (n *node) clearMemos() {
	n.heightMemo = 0
	n.childrenMemo = 0
}

func (n *node) height() int {
	if n == nil {
		return 0
	}
	if n.heightMemo != 0 {
		return n.heightMemo
	}
	l := n.left.height()
	r := n.right.height()
	if l > r {
		n.heightMemo = 1 + l
	} else {
		n.heightMemo = 1 + r
	}
	return n.heightMemo
}

func (n *node) children() int {
	if n == nil {
		return 0
	}
	return n.left.children() + n.right.children() + 1
}

func (n *node) balanceFactor() int {
	if n == nil {
		return 0
	}
	return n.left.height() - n.right.height()
}

func (n *node) inOrder(f func(string)) {
	if n == nil {
		return
	}
	n.left.inOrder(f)
	f(n.value)
	n.right.inOrder(f)
}

func (n *node) toMap() map[string]interface{} {
	m := map[string]interface{}{
		"value":    n.value,
		"height":   n.height(),
		"children": n.children(),
	}
	if n.left != nil {
		m["left"] = n.left.toMap()
	}
	if n.right != nil {
		m["right"] = n.right.toMap()
	}
	return m
}

func (n *node) toDot(w io.Writer) {
	if n == nil {
		return
	}
	// fmt.Fprintf(w, " %q [label=\"%v - bf=%v, height=%v\"];\n",
	// 	n.value, n.value, n.balanceFactor(), n.height())
	if n.left != nil {
		fmt.Fprintf(w, " %q -> %q [label=\"Left\"];\n", n.value, n.left.value)
		n.left.toDot(w)
	}
	if n.right != nil {
		fmt.Fprintf(w, " %q -> %q [label=\"Right\"];\n", n.value, n.right.value)
		n.right.toDot(w)
	}
}

type Tree struct {
	root      *node
	rotations int
}

func (t *Tree) Add(s string) {
	if t.root == nil {
		t.root = &node{value: s}
		return
	}

	rots := 0
	t.root, rots = addAt(s, t.root)
	t.rotations += rots
}

func addAt(s string, n *node) (*node, int) {
	rots := 0
	switch {
	case n.value == s:
		return n, 0
	case s < n.value:
		if n.left == nil {
			n.left = &node{value: s}
		} else {
			n.left, rots = addAt(s, n.left)
			n.left.clearMemos()
		}
	case s > n.value:
		if n.right == nil {
			n.right = &node{value: s}
		} else {
			n.right, rots = addAt(s, n.right)
			n.right.clearMemos()
		}
	}
	n.clearMemos()

	bf := n.balanceFactor()
	if bf < -1 || bf > 1 {
		r := 0
		n, r = rotate(n, bf)
		rots += r
	}

	return n, rots
}

func rotate(n *node, bf int) (*node, int) {
	rots := 0
	switch {
	case bf > 1:
		if n.left.balanceFactor() < -1 {
			n.left = rotateLeft(n.left)
			rots++
		}
		n = rotateRight(n)
		rots++
	case bf < -1:
		if n.right != nil && n.right.balanceFactor() > 1 {
			n.right = rotateRight(n.right)
			rots++
		}
		n = rotateLeft(n)
		rots++
	}
	return n, rots
}

func rotateRight(n *node) *node {
	pivot := n.left
	n.left = pivot.right
	pivot.right = n
	pivot.clearMemos()
	pivot.right.clearMemos()
	return pivot
}

func rotateLeft(n *node) *node {
	pivot := n.right
	n.right = pivot.left
	pivot.left = n
	pivot.clearMemos()
	pivot.left.clearMemos()
	return pivot
}

func (t Tree) Visit(f func(string)) {
	t.root.inOrder(f)
}

func (t Tree) ToDot(w io.Writer) error {
	fmt.Fprintf(w, "digraph \"g\" {\n")
	t.root.toDot(w)
	fmt.Fprintf(w, "}\n")
	return nil
}

func (t Tree) MarshalJSON() ([]byte, error) {
	return json.Marshal(t.root.toMap())
}

func main() {
	f, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatalf("Error opening file: %v", err)
	}
	defer f.Close()

	tree := Tree{}

	scanner := bufio.NewScanner(f)
	i := 0
	for scanner.Scan() {
		tree.Add(scanner.Text())
		i++
		if i%1000 == 0 {
			log.Printf("Read %v words, %v rotations, last was %q",
				i, tree.rotations, scanner.Text())
		}
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading stuff: %v", err)
	}

	dotf, err := os.Create("/tmp/alpha.dot")
	if err != nil {
		log.Fatalf("Error creating dot file: %v", err)
	}
	defer dotf.Close()
	tree.ToDot(dotf)

	jf, err := os.Create("/tmp/alpha.json")
	if err != nil {
		log.Fatalf("Error creating json file: %v", err)
	}
	defer jf.Close()
	e := json.NewEncoder(jf)
	err = e.Encode(tree)
	if err != nil {
		log.Fatalf("Error encoding json: %v", err)
	}
}

avl_test.go

package main

import "testing"

var testData = []string{
	"hopingly",
	"disentrancing",
	"attingency",
	"felsite",
	"untelevised",
	"floccose",
	"semiporcelain",
	"tided",
	"unaccountable",
	"unbumptious",
}

var exp = []string{
	"attingency",
	"disentrancing",
	"felsite",
	"floccose",
	"hopingly",
	"semiporcelain",
	"tided",
	"unaccountable",
	"unbumptious",
	"untelevised",
}

func TestInsertSome(t *testing.T) {
	tree := Tree{}
	for _, w := range testData {
		tree.Add(w)
	}

	found := []string{}
	tree.Visit(func(s string) {
		found = append(found, s)
	})
	if len(found) < len(testData) {
		t.Fatalf("Only visisted %v of %v items", len(found), len(testData))
	}
	for i := range exp {
		if exp[i] != found[i] {
			t.Errorf("Expected %v at %v, got %v", exp[i], i, found[i])
		}
	}

}