FST read-path and search optimizations

builder.go

@@ -146,7 +146,7 @@ func (b *Builder) compileFrom(iState int) error {
 		var err error
 		addr, err = b.compile(node)
 		if err != nil {
-			return nil
+			return err
 		}
 	}
 	b.unfinished.topLastFreeze(addr)

decoder_v1.go

@@ -261,6 +261,29 @@ func (f *fstStateV1) TransitionFor(b byte) (int, int, uint64) {
 	return f.numTrans - pos - 1, dest, out
 }
 
+func (f *fstStateV1) TransitionDestForOffset(i int) (int, uint64) {
+	if f.isEncodedSingle() {
+		return int(f.singleTransAddr), f.singleTransOut
+	}
+	// transitions are stored reversed, so sorted offset i lives at reversed
+	// position numTrans-i-1 across the dest/output arrays.
+	pos := f.numTrans - i - 1
+
+	transDests := f.data[f.destBottom:f.destTop]
+	dest := int(readPackedUint(transDests[pos*f.transSize : pos*f.transSize+f.transSize]))
+	if dest > 0 {
+		// convert delta
+		dest = f.bottom - dest
+	}
+
+	var out uint64
+	if f.outSize > 0 {
+		transVals := f.data[f.outBottom:f.outTop]
+		out = readPackedUint(transVals[pos*f.outSize : pos*f.outSize+f.outSize])
+	}
+	return dest, out
+}
+
 func (f *fstStateV1) String() string {
 	rv := ""
 	rv += fmt.Sprintf("State: %d (%#x)", f.top, f.top)

encoding.go

@@ -84,4 +84,10 @@ type fstState interface {
 	NumTransitions() int
 	TransitionFor(b byte) (int, int, uint64)
 	TransitionAt(i int) byte
+	// TransitionDestForOffset returns the destination address and output for
+	// the transition at sorted offset i, without searching by key. Callers
+	// enumerating transitions in order (e.g. the automaton-guided iterator)
+	// already know the offset from TransitionAt, so this avoids the redundant
+	// key lookup that TransitionFor would perform.
+	TransitionDestForOffset(i int) (int, uint64)
 }

fst.go

@@ -16,6 +16,7 @@ package vellum
 
 import (
 	"io"
+	"sync"
 
 	"github.com/bits-and-blooms/bitset"
 )
@@ -31,6 +32,13 @@ type FST struct {
 	typ     int
 	data    []byte
 	decoder decoder
+
+	// statePool recycles transient fstState instances used by the
+	// Automaton/Transducer interface methods (Accept/AcceptWithVal,
+	// IsMatch/IsMatchWithVal), which would otherwise allocate a fresh
+	// state on every call. The pool is safe for concurrent use, so the
+	// FST remains usable from multiple goroutines.
+	statePool sync.Pool
 }
 
 func new(data []byte, f io.Closer) (rv *FST, err error) {
@@ -160,21 +168,26 @@ func (f *FST) Accept(addr int, b byte) int {
 // IsMatchWithVal returns if this state is a matching state in this Automaton
 // and also returns the final output value for this state
 func (f *FST) IsMatchWithVal(addr int) (bool, uint64) {
-	s, err := f.decoder.stateAt(addr, nil)
+	prealloc, _ := f.statePool.Get().(fstState)
+	s, err := f.decoder.stateAt(addr, prealloc)
 	if err != nil {
 		return false, 0
 	}
-	return s.Final(), s.FinalOutput()
+	final, out := s.Final(), s.FinalOutput()
+	f.statePool.Put(s)
+	return final, out
 }
 
 // AcceptWithVal returns the next state for this Automaton on input of byte b
 // and also returns the output value for the transition
 func (f *FST) AcceptWithVal(addr int, b byte) (int, uint64) {
-	s, err := f.decoder.stateAt(addr, nil)
+	prealloc, _ := f.statePool.Get().(fstState)
+	s, err := f.decoder.stateAt(addr, prealloc)
 	if err != nil {
 		return noneAddr, 0
 	}
 	_, next, output := s.TransitionFor(b)
+	f.statePool.Put(s)
 	return next, output
 }
 
@@ -240,7 +253,17 @@ func (a addrStack) Pop() (addrStack, int) {
 // Reader() returns a Reader instance that a single thread may use to
 // retrieve data from the FST
 func (f *FST) Reader() (*Reader, error) {
-	return &Reader{f: f}, nil
+	r := &Reader{f: f, rootAddr: f.decoder.getRoot()}
+	// Parse the root state once and cache it. The root is traversed on every
+	// Get, so caching the (immutable) parsed form lets each lookup skip
+	// re-decoding it - a per-call saving on top of avoiding the state alloc.
+	rs, err := f.decoder.stateAt(r.rootAddr, &r.root)
+	if err == nil {
+		if _, ok := rs.(*fstStateV1); ok {
+			r.hasRoot = true
+		}
+	}
+	return r, nil
 }
 
 func (f *FST) GetMinKey() ([]byte, error) {
@@ -293,8 +316,38 @@ func (f *FST) GetMaxKey() ([]byte, error) {
 type Reader struct {
 	f        *FST
 	prealloc fstStateV1
+
+	// root holds the pre-parsed root state; it is copied into prealloc at the
+	// start of each Get so the root never has to be re-decoded.
+	rootAddr int
+	hasRoot  bool
+	root     fstStateV1
 }
 
 func (r *Reader) Get(input []byte) (uint64, bool, error) {
-	return r.f.get(input, &r.prealloc)
+	if !r.hasRoot {
+		return r.f.get(input, &r.prealloc)
+	}
+
+	var total uint64
+	r.prealloc = r.root // reuse the cached, already-decoded root
+	state := &r.prealloc
+	for _, c := range input {
+		_, curr, output := state.TransitionFor(c)
+		if curr == noneAddr {
+			return 0, false, nil
+		}
+		next, err := r.f.decoder.stateAt(curr, state)
+		if err != nil {
+			return 0, false, err
+		}
+		state = next.(*fstStateV1)
+		total += output
+	}
+
+	if state.Final() {
+		total += state.FinalOutput()
+		return total, true, nil
+	}
+	return 0, false, nil
 }

fst_iterator.go

@@ -240,7 +240,12 @@ OUTER:
 				continue INNER
 			}
 
-			pos, nextAddr, v := curr.TransitionFor(t)
+			// only now that the automaton has accepted do we resolve the
+			// dest/output - by offset, so no key search is needed. The sorted
+			// offset is exactly the position recorded on the stack
+			// (TransitionFor returns numTrans-revpos-1 == offset).
+			nextAddr, v := curr.TransitionDestForOffset(nextOffset)
+			pos := nextOffset
 
 			// the next slot in the statesStack might have an
 			// fstState instance that we can reuse

levenshtein/alphabet.go

@@ -17,6 +17,7 @@ package levenshtein
 import (
 	"fmt"
 	"sort"
+	"strings"
 	"unicode/utf8"
 )
 
@@ -78,16 +79,17 @@ func (a *Alphabet) next() (rune, FullCharacteristicVector, error) {
 
 func dedupe(in string) string {
 	lookUp := make(map[rune]struct{}, len(in))
-	var rv string
+	var sb strings.Builder
+	sb.Grow(len(in))
 	for len(in) > 0 {
 		r, size := utf8.DecodeRuneInString(in)
 		in = in[size:]
 		if _, ok := lookUp[r]; !ok {
-			rv += string(r)
+			sb.WriteRune(r)
 			lookUp[r] = struct{}{}
 		}
 	}
-	return rv
+	return sb.String()
 }
 
 func queryChars(qChars string) Alphabet {

levenshtein/benchmark_test.go

@@ -18,6 +18,24 @@ import (
 	"testing"
 )
 
+func BenchmarkNewLevenshteinAutomatonBuilder1(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_, err := NewLevenshteinAutomatonBuilder(1, true)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkNewLevenshteinAutomatonBuilder2(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_, err := NewLevenshteinAutomatonBuilder(2, false)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
 func BenchmarkNewEvalEditDistance1(b *testing.B) {
 	lb, _ := NewLevenshteinAutomatonBuilder(1, true)
 

levenshtein/parametric_dfa.go

@@ -15,8 +15,7 @@
 package levenshtein
 
 import (
-	"crypto/sha256"
-	"encoding/json"
+	"encoding/binary"
 	"fmt"
 	"math"
 )
@@ -311,13 +310,14 @@ func fromNfa(nfa *LevenshteinNFA) (*ParametricDFA, error) {
 }
 
 type hash struct {
-	index map[[32]byte]int
-	items []MultiState
+	index  map[string]int
+	items  []MultiState
+	keyBuf []byte
 }
 
 func newHash() *hash {
 	return &hash{
-		index: make(map[[32]byte]int, 100),
+		index: make(map[string]int, 100),
 		items: make([]MultiState, 0, 100),
 	}
 }
@@ -326,9 +326,11 @@ func (h *hash) getOrAllocate(m MultiState) int {
 	size := len(h.items)
 	var exists bool
 	var pos int
-	sha := getHash(&m)
-	if pos, exists = h.index[sha]; !exists {
-		h.index[sha] = size
+	h.keyBuf = encodeMultiState(&m, h.keyBuf)
+	if pos, exists = h.index[string(h.keyBuf)]; !exists {
+		// string(h.keyBuf) allocates a fresh, immutable copy for the map key,
+		// so reusing h.keyBuf on the next call is safe.
+		h.index[string(h.keyBuf)] = size
 		pos = size
 		h.items = append(h.items, m)
 	}
@@ -339,11 +341,25 @@ func (h *hash) getFromID(id int) *MultiState {
 	return &h.items[id]
 }
 
-func getHash(ms *MultiState) [32]byte {
-	msBytes := []byte{}
-	for _, state := range ms.states {
-		jsonBytes, _ := json.Marshal(&state)
-		msBytes = append(msBytes, jsonBytes...)
+// encodeMultiState serializes ms into buf (reusing its capacity) as an exact,
+// collision-free key: 6 bytes per NFAState (4 offset, 1 distance, 1 transpose).
+func encodeMultiState(ms *MultiState, buf []byte) []byte {
+	const stateSize = 6
+	need := len(ms.states) * stateSize
+	if cap(buf) < need {
+		buf = make([]byte, need)
+	} else {
+		buf = buf[:need]
 	}
-	return sha256.Sum256(msBytes)
+	for i, state := range ms.states {
+		off := i * stateSize
+		binary.LittleEndian.PutUint32(buf[off:], state.Offset)
+		buf[off+4] = state.Distance
+		if state.InTranspose {
+			buf[off+5] = 1
+		} else {
+			buf[off+5] = 0
+		}
+	}
+	return buf
 }