From daf014d650ef87c82141b8d8208dcffee1c48432 Mon Sep 17 00:00:00 2001 From: Engineering Team Date: Thu, 18 Jun 2026 14:12:37 +0800 Subject: [PATCH] Add ReusableState for allocation-free incremental FST walks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AcceptWithVal/IsMatchWithVal decode each source state with a nil prealloc, so they allocate a fresh fstStateV1 on every transition. Callers that walk the FST one byte at a time over many overlapping prefixes — e.g. building a CJK segmentation DAG that restarts from every input position — therefore allocate once per transition. In a real FST-backed CJK tokenizer this was ~69% of all allocations (and ~2x the wall-clock once removed). Add an opaque ReusableState plus AcceptWithValState/IsMatchWithValState that thread a reusable *fstStateV1 through decoder.stateAt. fstStateV1 only references the FST data (no owned slices), so a single ReusableState backs an entire walk with zero per-transition allocation. Results are identical to the stock methods; the existing API is untouched. Tests: TestReusableStateMatchesStock (parity across match/prefix/absent keys) and TestReusableStateNoAlloc (AllocsPerRun == 0). Signed-off-by: Engineering Team Co-Authored-By: Claude Opus 4.8 (1M context) --- fst_reuse.go | 43 ++++++++++++++++++++++++ fst_reuse_test.go | 86 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+) create mode 100644 fst_reuse.go create mode 100644 fst_reuse_test.go diff --git a/fst_reuse.go b/fst_reuse.go new file mode 100644 index 0000000..8430ab9 --- /dev/null +++ b/fst_reuse.go @@ -0,0 +1,43 @@ +package vellum + +// ReusableState is an opaque, reusable decode buffer for the incremental +// AcceptWithValState / IsMatchWithValState walk. +// +// The stock AcceptWithVal / IsMatchWithVal decode each source state with a nil +// prealloc, so they allocate a fresh internal state on every transition. Callers +// that walk the FST one byte at a time over many overlapping prefixes — e.g. +// building a CJK segmentation DAG, where the walk restarts from every input +// position — pay one allocation per transition. Holding a single ReusableState +// for the whole walk eliminates that allocation (the internal state only +// references the FST data, so it can be reused in place). +// +// A ReusableState is NOT safe for concurrent use; give each goroutine its own. +type ReusableState struct { + s fstStateV1 +} + +// NewReusableState returns a fresh, reusable decode buffer. +func NewReusableState() *ReusableState { + return &ReusableState{} +} + +// AcceptWithValState behaves exactly like AcceptWithVal but decodes the source +// state into the caller-owned rs instead of allocating a new internal state. +func (f *FST) AcceptWithValState(addr int, b byte, rs *ReusableState) (int, uint64) { + s, err := f.decoder.stateAt(addr, &rs.s) + if err != nil { + return noneAddr, 0 + } + _, next, output := s.TransitionFor(b) + return next, output +} + +// IsMatchWithValState behaves exactly like IsMatchWithVal but decodes the state +// into the caller-owned rs instead of allocating a new internal state. +func (f *FST) IsMatchWithValState(addr int, rs *ReusableState) (bool, uint64) { + s, err := f.decoder.stateAt(addr, &rs.s) + if err != nil { + return false, 0 + } + return s.Final(), s.FinalOutput() +} diff --git a/fst_reuse_test.go b/fst_reuse_test.go new file mode 100644 index 0000000..5e7752f --- /dev/null +++ b/fst_reuse_test.go @@ -0,0 +1,86 @@ +package vellum + +import ( + "bytes" + "testing" +) + +// buildTestFST builds a small FST from sorted keys with val = len(key). +func buildTestFST(t *testing.T, keys []string) *FST { + t.Helper() + var buf bytes.Buffer + b, err := New(&buf, nil) + if err != nil { + t.Fatal(err) + } + for _, k := range keys { + if err := b.Insert([]byte(k), uint64(len(k))); err != nil { + t.Fatal(err) + } + } + if err := b.Close(); err != nil { + t.Fatal(err) + } + fst, err := Load(buf.Bytes()) + if err != nil { + t.Fatal(err) + } + return fst +} + +// walk consumes key one byte at a time, returning the accumulated output and +// whether the final state matched. When rs != nil it uses the reusable-state +// variants; otherwise the stock ones. +func walk(fst *FST, key string, rs *ReusableState) (uint64, bool) { + addr := fst.Start() + var sum uint64 + for i := 0; i < len(key); i++ { + var next int + var out uint64 + if rs != nil { + next, out = fst.AcceptWithValState(addr, key[i], rs) + } else { + next, out = fst.AcceptWithVal(addr, key[i]) + } + if next == noneAddr { + return 0, false + } + addr = next + sum += out + } + var final bool + var fout uint64 + if rs != nil { + final, fout = fst.IsMatchWithValState(addr, rs) + } else { + final, fout = fst.IsMatchWithVal(addr) + } + return sum + fout, final +} + +// TestReusableStateMatchesStock verifies the *State variants return identical +// results to AcceptWithVal/IsMatchWithVal for matching, prefix, and absent keys. +func TestReusableStateMatchesStock(t *testing.T) { + fst := buildTestFST(t, []string{"cat", "cats", "dog", "doge"}) + rs := NewReusableState() + for _, k := range []string{"cat", "cats", "dog", "doge", "ca", "do", "zzz", "catx"} { + wantVal, wantFinal := walk(fst, k, nil) + gotVal, gotFinal := walk(fst, k, rs) // reuse rs across keys on purpose + if gotVal != wantVal || gotFinal != wantFinal { + t.Errorf("%q: state walk (%d,%v) != stock (%d,%v)", k, gotVal, gotFinal, wantVal, wantFinal) + } + } +} + +// TestReusableStateNoAlloc asserts the reusable-state walk does not allocate per +// transition (the whole point of the API). +func TestReusableStateNoAlloc(t *testing.T) { + fst := buildTestFST(t, []string{"cat", "cats", "dog", "doge"}) + rs := NewReusableState() + allocs := testing.AllocsPerRun(200, func() { + walk(fst, "cats", rs) + }) + if allocs != 0 { + t.Errorf("reusable-state walk allocated %v objects/run, want 0", allocs) + } +}