diff --git a/builder.go b/builder.go index b9012fd..2aef364 100644 --- a/builder.go +++ b/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) diff --git a/decoder_v1.go b/decoder_v1.go index d56e61d..0366d72 100644 --- a/decoder_v1.go +++ b/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) diff --git a/encoding.go b/encoding.go index 988d486..fbff5fe 100644 --- a/encoding.go +++ b/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) } diff --git a/fst.go b/fst.go index 3140042..c500a9c 100644 --- a/fst.go +++ b/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 } diff --git a/fst_iterator.go b/fst_iterator.go index f5c374e..8b4c674 100644 --- a/fst_iterator.go +++ b/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 diff --git a/levenshtein/alphabet.go b/levenshtein/alphabet.go index ec28512..5411165 100644 --- a/levenshtein/alphabet.go +++ b/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 { diff --git a/levenshtein/benchmark_test.go b/levenshtein/benchmark_test.go index 1c292b7..7d2a0db 100644 --- a/levenshtein/benchmark_test.go +++ b/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) diff --git a/levenshtein/parametric_dfa.go b/levenshtein/parametric_dfa.go index 41d2fcf..8bc2ed2 100644 --- a/levenshtein/parametric_dfa.go +++ b/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 } diff --git a/optim_bench_test.go b/optim_bench_test.go new file mode 100644 index 0000000..f6f126a --- /dev/null +++ b/optim_bench_test.go @@ -0,0 +1,183 @@ +// Copyright (c) 2017 Couchbase, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vellum + +import ( + "bytes" + "sort" + "testing" + + "github.com/blevesearch/vellum/levenshtein" + "github.com/blevesearch/vellum/regexp" +) + +// genWideKeys produces n deterministic, sorted, unique keys with a wide byte +// distribution, approximating a term dictionary of ids / hashes / tokens where +// the root and near-root FST states have very high fanout. +func genWideKeys(n, keyLen int) [][]byte { + keys := make([][]byte, 0, n) + seen := make(map[string]struct{}, n) + var state uint64 = 0x9e3779b97f4a7c15 + next := func() uint64 { + state ^= state << 13 + state ^= state >> 7 + state ^= state << 17 + return state + } + for len(keys) < n { + k := make([]byte, keyLen) + for j := 0; j < keyLen; j++ { + k[j] = byte(next()%255) + 1 + } + if _, ok := seen[string(k)]; ok { + continue + } + seen[string(k)] = struct{}{} + keys = append(keys, k) + } + sort.Slice(keys, func(i, j int) bool { return bytes.Compare(keys[i], keys[j]) < 0 }) + return keys +} + +func benchWordKeys(tb testing.TB) [][]byte { + words, err := loadWords("data/words-1000.txt") + if err != nil { + tb.Fatal(err) + } + sort.Strings(words) + keys := make([][]byte, len(words)) + for i, w := range words { + keys[i] = []byte(w) + } + return keys +} + +// buildBenchFST builds an FST over the (sorted, unique) keys using the default +// builder options. +func buildBenchFST(tb testing.TB, keys [][]byte) []byte { + var buf bytes.Buffer + b, err := New(&buf, nil) + if err != nil { + tb.Fatal(err) + } + for i, k := range keys { + if err := b.Insert(k, uint64(i)); err != nil { + tb.Fatal(err) + } + } + if err := b.Close(); err != nil { + tb.Fatal(err) + } + return buf.Bytes() +} + +func benchGet(b *testing.B, keys [][]byte) { + fst, err := Load(buildBenchFST(b, keys)) + if err != nil { + b.Fatal(err) + } + defer fst.Close() + r, err := fst.Reader() + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + var sink uint64 + for i := 0; i < b.N; i++ { + for _, k := range keys { + v, _, _ := r.Get(k) + sink += v + } + } + _ = sink +} + +func benchScan(b *testing.B, keys [][]byte) { + fst, err := Load(buildBenchFST(b, keys)) + if err != nil { + b.Fatal(err) + } + defer fst.Close() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + itr, err := fst.Iterator(nil, nil) + for err == nil { + _, _ = itr.Current() + err = itr.Next() + } + } +} + +func benchFuzzy(b *testing.B, keys [][]byte, query string, dist uint8) { + fst, err := Load(buildBenchFST(b, keys)) + if err != nil { + b.Fatal(err) + } + defer fst.Close() + lb, err := levenshtein.NewLevenshteinAutomatonBuilder(dist, false) + if err != nil { + b.Fatal(err) + } + dfa, err := lb.BuildDfa(query, dist) + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + itr, err := fst.Search(dfa, nil, nil) + for err == nil { + _, _ = itr.Current() + err = itr.Next() + } + } +} + +func benchRegex(b *testing.B, keys [][]byte, expr string) { + fst, err := Load(buildBenchFST(b, keys)) + if err != nil { + b.Fatal(err) + } + defer fst.Close() + r, err := regexp.New(expr) + if err != nil { + b.Fatal(err) + } + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + itr, err := fst.Search(r, nil, nil) + for err == nil { + _, _ = itr.Current() + err = itr.Next() + } + } +} + +// Exact lookups (Reader.Get) - exercises the cached root state. +func BenchmarkOptGetWords(b *testing.B) { benchGet(b, benchWordKeys(b)) } +func BenchmarkOptGetWide(b *testing.B) { benchGet(b, genWideKeys(50000, 8)) } +func BenchmarkOptGetWide2(b *testing.B) { benchGet(b, genWideKeys(60000, 2)) } + +// Full range scan - exercises the iterator offset accessor. +func BenchmarkOptScanWords(b *testing.B) { benchScan(b, benchWordKeys(b)) } +func BenchmarkOptScanWide(b *testing.B) { benchScan(b, genWideKeys(50000, 8)) } + +// Automaton-guided search. +func BenchmarkOptFuzzyWords1(b *testing.B) { benchFuzzy(b, benchWordKeys(b), "the", 1) } +func BenchmarkOptFuzzyWords2(b *testing.B) { benchFuzzy(b, benchWordKeys(b), "American", 2) } +func BenchmarkOptRegexWords(b *testing.B) { benchRegex(b, benchWordKeys(b), ".*a.*e.*") } diff --git a/regexp/dfa.go b/regexp/dfa.go index 7e6fb29..5916ffb 100644 --- a/regexp/dfa.go +++ b/regexp/dfa.go @@ -43,7 +43,7 @@ func newDfaBuilder(insts prog) *dfaBuilder { } // add 0 state that is invalid d.dfa.states = append(d.dfa.states, state{ - next: make([]int, 256), + next: make([]uint32, 256), match: false, }) return d @@ -87,7 +87,7 @@ func (d *dfaBuilder) runState(cur, next *sparseSet, state int, b byte, instsReus d.dfa.run(cur, next, b) var nextState int nextState, instsReuse = d.cachedState(next, instsReuse) - d.dfa.states[state].next[b] = nextState + d.dfa.states[state].next[b] = uint32(nextState) return nextState, instsReuse } @@ -130,7 +130,7 @@ func (d *dfaBuilder) cachedState(set *sparseSet, } d.dfa.states = append(d.dfa.states, state{ insts: insts, - next: make([]int, 256), + next: make([]uint32, 256), match: isMatch, }) newV := len(d.dfa.states) - 1 @@ -177,7 +177,10 @@ func (d *dfa) run(from, to *sparseSet, b byte) bool { type state struct { insts []uint - next []int + // next holds the destination state index for each of the 256 input + // bytes. State indices are bounded by StateLimit, so uint32 is always + // sufficient and halves the table size versus int on 64-bit platforms. + next []uint32 match bool } diff --git a/regexp/footprint_test.go b/regexp/footprint_test.go new file mode 100644 index 0000000..3494fa6 --- /dev/null +++ b/regexp/footprint_test.go @@ -0,0 +1,26 @@ +package regexp + +import ( + "testing" +) + +// BenchmarkDFAFootprint tracks the heap footprint of compiling regexps of +// varying complexity. The transition tables (state.next) dominate the +// allocation, so this guards against regressions in their size. +func BenchmarkDFAFootprint(b *testing.B) { + exprs := []string{"my.*h", "[a-z]+@[a-z]+\\.(com|net|org)", "(abc|def|ghi)*[0-9]{2,4}foo.*bar"} + for _, e := range exprs { + b.Run(e, func(b *testing.B) { + b.ReportAllocs() + var states int + for i := 0; i < b.N; i++ { + r, err := New(e) + if err != nil { + b.Fatal(err) + } + states = len(r.dfa.states) + } + b.ReportMetric(float64(states), "states") + }) + } +} diff --git a/regexp/regexp.go b/regexp/regexp.go index 8d28b23..f358d8e 100644 --- a/regexp/regexp.go +++ b/regexp/regexp.go @@ -113,7 +113,7 @@ func (r *Regexp) WillAlwaysMatch(int) bool { // when currently in the state s. func (r *Regexp) Accept(s int, b byte) int { if s < len(r.dfa.states) { - return r.dfa.states[s].next[b] + return int(r.dfa.states[s].next[b]) } return 0 } diff --git a/transducer_bench_test.go b/transducer_bench_test.go new file mode 100644 index 0000000..d84caa8 --- /dev/null +++ b/transducer_bench_test.go @@ -0,0 +1,66 @@ +package vellum + +import ( + "bytes" + "testing" +) + +func buildSmallFST(tb testing.TB) *FST { + var buf bytes.Buffer + b, err := New(&buf, nil) + if err != nil { + tb.Fatal(err) + } + keys := [][]byte{ + []byte("apple"), []byte("application"), []byte("apply"), + []byte("banana"), []byte("band"), []byte("bandana"), + []byte("orange"), []byte("orchard"), []byte("ordinary"), + } + for i, k := range keys { + if err := b.Insert(k, uint64(i)); err != nil { + tb.Fatal(err) + } + } + if err := b.Close(); err != nil { + tb.Fatal(err) + } + fst, err := Load(buf.Bytes()) + if err != nil { + tb.Fatal(err) + } + return fst +} + +// generic Transducer path: AcceptWithVal per byte + IsMatchWithVal once +func BenchmarkTransducerGet(b *testing.B) { + fst := buildSmallFST(b) + key := []byte("application") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _ = TransducerGet(fst, key) + } +} + +// FST.Get (allocates one state per call) +func BenchmarkFSTGet(b *testing.B) { + fst := buildSmallFST(b) + key := []byte("application") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _, _ = fst.Get(key) + } +} + +// Reader.Get (reuses prealloc state) +func BenchmarkReaderGet(b *testing.B) { + fst := buildSmallFST(b) + r, _ := fst.Reader() + key := []byte("application") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, _, _ = r.Get(key) + } +} diff --git a/transducer_race_test.go b/transducer_race_test.go new file mode 100644 index 0000000..da5a214 --- /dev/null +++ b/transducer_race_test.go @@ -0,0 +1,40 @@ +package vellum + +import ( + "sync" + "testing" +) + +// Exercises the pooled Automaton/Transducer methods concurrently to ensure +// the statePool sharing is race-free. +func TestTransducerConcurrent(t *testing.T) { + fst := buildSmallFST(t) + keys := [][]byte{ + []byte("application"), []byte("banana"), []byte("orchard"), + []byte("apply"), []byte("bandana"), []byte("ordinary"), + []byte("missing"), []byte("apple"), + } + var wg sync.WaitGroup + for g := 0; g < 16; g++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 2000; i++ { + k := keys[i%len(keys)] + _, _ = TransducerGet(fst, k) + // also drive the bare Automaton methods + st := fst.Start() + for _, b := range k { + st = fst.Accept(st, b) + if st == noneAddr { + break + } + } + if st != noneAddr { + _ = fst.IsMatch(st) + } + } + }() + } + wg.Wait() +}