Markov chain generator in Go

markov.go

// Markov chain generator
// Builds chain from stdin, outputs n generated words to stdout
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"
)

const (
	NPREF   = 2
	NONWORD = "\n"
)

type Prefix [NPREF]string

func NewPrefix() *Prefix {
	var p Prefix
	for i := 0; i < NPREF; i++ {
		p.PushBack(NONWORD)
	}
	return &p
}

func (p *Prefix) PushBack(s string) {
	for i := 0; i < len(p)-1; i++ {
		p[i] = p[i+1]
	}
	p[len(p)-1] = s
}

type Sufflist []string

type State map[Prefix]Sufflist

type Chain struct {
	Pref *Prefix
	St   State
}

func NewChain() *Chain {
	c := new(Chain)
	c.Pref = NewPrefix()
	c.St = make(State)
	return c
}

func (c *Chain) Add(s string) {
	sufflist := c.St[*c.Pref]
	c.St[*c.Pref] = append(sufflist, s)
	c.Pref.PushBack(s)
}

func (c *Chain) Build(r io.Reader) {
	ch := make(chan string)
	go ReadWords(r, ch)
	for word := range ch {
		c.Add(word)
	}
	c.Add(NONWORD)
}

func (c *Chain) Generate(n int, ch chan string) {
	pref := NewPrefix()
	for i := 0; i < n; i++ {
		sufflist := c.St[*pref]
		i := rand.Intn(len(sufflist))
		word := sufflist[i]
		if word == NONWORD {
			break
		}
		ch <- word
		pref.PushBack(word)
	}
	close(ch)
}

func ReadWords(r io.Reader, ch chan string) {
	var (
		line     []byte
		isPrefix bool
		err      error
	)
	br := bufio.NewReader(r)
	for {
		if line, isPrefix, err = br.ReadLine(); err != nil {
			if err == io.EOF {
				break
			} else {
				log.Fatal(err)
			}
		}
		if !isPrefix {
			for _, word := range strings.Split(string(line), " ") {
				word = strings.Trim(word, "\n\r ")
				if word != "" {
					ch <- word
				}
			}
		}
	}
	close(ch)
}

var nwords = flag.Int("n", 200, "number of words to generate")

func main() {
	flag.Parse()
	c := NewChain()
	c.Build(os.Stdin)
	ch := make(chan string)
	rand.Seed(time.Now().Unix())
	go c.Generate(*nwords, ch)
	for word := range ch {
		fmt.Println(word)
	}
}

Usage:

$ cat novel.txt | go run markov.go