96 lines
2.2 KiB
Go
96 lines
2.2 KiB
Go
package markov
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
)
|
|
|
|
// model is the on-disk representation. It is versioned so a future change
|
|
// to the format can be detected rather than silently misread.
|
|
type model struct {
|
|
Version int `json:"version"`
|
|
Order int `json:"order"`
|
|
Freq map[string]map[string]int `json:"freq"`
|
|
}
|
|
|
|
const modelVersion = 1
|
|
|
|
// Save writes the chain to w as JSON.
|
|
func (c *Chain) Save(w io.Writer) error {
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
enc := json.NewEncoder(w)
|
|
enc.SetIndent("", " ")
|
|
return enc.Encode(model{Version: modelVersion, Order: c.order, Freq: c.freq})
|
|
}
|
|
|
|
// SaveFile writes the chain to path atomically: it writes a temporary file
|
|
// alongside the target and renames it into place, so an interrupted save
|
|
// cannot leave a half-written model where a good one used to be.
|
|
func (c *Chain) SaveFile(path string) (err error) {
|
|
dir := filepath.Dir(path)
|
|
tmp, err := os.CreateTemp(dir, filepath.Base(path)+".tmp*")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tmpName := tmp.Name()
|
|
defer func() {
|
|
if err != nil {
|
|
tmp.Close()
|
|
os.Remove(tmpName)
|
|
}
|
|
}()
|
|
|
|
if err = c.Save(tmp); err != nil {
|
|
return err
|
|
}
|
|
if err = tmp.Sync(); err != nil {
|
|
return err
|
|
}
|
|
if err = tmp.Close(); err != nil {
|
|
return err
|
|
}
|
|
return os.Rename(tmpName, path)
|
|
}
|
|
|
|
// Load reads a chain previously written by Save.
|
|
func Load(r io.Reader) (*Chain, error) {
|
|
var m model
|
|
if err := json.NewDecoder(r).Decode(&m); err != nil {
|
|
return nil, fmt.Errorf("markov: decoding model: %w", err)
|
|
}
|
|
if m.Version != modelVersion {
|
|
return nil, fmt.Errorf("markov: unsupported model version %d (want %d)", m.Version, modelVersion)
|
|
}
|
|
c, err := New(m.Order)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("markov: invalid model: %w", err)
|
|
}
|
|
if m.Freq != nil {
|
|
c.freq = m.Freq
|
|
}
|
|
// Totals are derived rather than stored, so a hand-edited or truncated
|
|
// model cannot desynchronise counts from their sum.
|
|
for prefix, succ := range c.freq {
|
|
sum := 0
|
|
for _, n := range succ {
|
|
sum += n
|
|
}
|
|
c.total[prefix] = sum
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
// LoadFile reads a chain from path.
|
|
func LoadFile(path string) (*Chain, error) {
|
|
f, err := os.Open(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer f.Close()
|
|
return Load(f)
|
|
}
|