github.com/graybobo/golang.org-package-offline-cache@v0.0.0-20200626051047-6608995c132f/x/talks/2012/chat/markov/markov.go (about)

     1  // +build OMIT
     2  
     3  package main
     4  
     5  // This Markov chain code is taken from the "Generating arbitrary text"
     6  // codewalk: http://golang.org/doc/codewalk/markov/
     7  
     8  import (
     9  	"bytes"
    10  	"fmt"
    11  	"math/rand"
    12  	"strings"
    13  	"sync"
    14  )
    15  
    16  // Prefix is a Markov chain prefix of one or more words.
    17  type Prefix []string
    18  
    19  // String returns the Prefix as a string (for use as a map key).
    20  func (p Prefix) String() string {
    21  	return strings.Join(p, " ")
    22  }
    23  
    24  // Shift removes the first word from the Prefix and appends the given word.
    25  func (p Prefix) Shift(word string) {
    26  	copy(p, p[1:])
    27  	p[len(p)-1] = word
    28  }
    29  
    30  // Chain contains a map ("chain") of prefixes to a list of suffixes.
    31  // A prefix is a string of prefixLen words joined with spaces.
    32  // A suffix is a single word. A prefix can have multiple suffixes.
    33  type Chain struct {
    34  	chain     map[string][]string
    35  	prefixLen int
    36  	mu        sync.Mutex
    37  }
    38  
    39  // NewChain returns a new Chain with prefixes of prefixLen words.
    40  func NewChain(prefixLen int) *Chain {
    41  	return &Chain{
    42  		chain:     make(map[string][]string),
    43  		prefixLen: prefixLen,
    44  	}
    45  }
    46  
    47  // Write parses the bytes into prefixes and suffixes that are stored in Chain.
    48  func (c *Chain) Write(b []byte) (int, error) {
    49  	br := bytes.NewReader(b)
    50  	p := make(Prefix, c.prefixLen)
    51  	for {
    52  		var s string
    53  		if _, err := fmt.Fscan(br, &s); err != nil {
    54  			break
    55  		}
    56  		key := p.String()
    57  		c.mu.Lock()
    58  		c.chain[key] = append(c.chain[key], s)
    59  		c.mu.Unlock()
    60  		p.Shift(s)
    61  	}
    62  	return len(b), nil
    63  }
    64  
    65  // Generate returns a string of at most n words generated from Chain.
    66  func (c *Chain) Generate(n int) string {
    67  	c.mu.Lock()
    68  	defer c.mu.Unlock()
    69  	p := make(Prefix, c.prefixLen)
    70  	var words []string
    71  	for i := 0; i < n; i++ {
    72  		choices := c.chain[p.String()]
    73  		if len(choices) == 0 {
    74  			break
    75  		}
    76  		next := choices[rand.Intn(len(choices))]
    77  		words = append(words, next)
    78  		p.Shift(next)
    79  	}
    80  	return strings.Join(words, " ")
    81  }