Deprecate ScoreFunc, replace it with FilterFunc to provide more control over the result composition

Author: cyradin

README.md

@@ -102,11 +102,13 @@ See [options.go](./options.go) for the list of available options.
 You can provide a custom scoring function if needed:
 
 ```go
-	var scoreFunc spellchecker.ScoreFunc = func(src, candidate []rune, distance, cnt int) float64 {
-		return 1.0 // constant score
+	var fn spellchecker.FilterFunc = func(src, candidate []rune, cnt int) (float64, bool) {
+		// you can calculate Levenshtein distance here (see defaultFilterFunc in options.go for example)
+
+		return 1.0, true // constant score
 	}
 
-	sc, err := spellchecker.New("abc", spellchecker.WithScoreFunc(scoreFunc))
+	sc, err := spellchecker.New("abc", spellchecker.WithFilterFunc(fn))
 	if err != nil {
 		// handle err
 	}
@@ -118,7 +120,7 @@ You can provide a custom scoring function if needed:
 		// handle err
 	}
 
-	err = sc.WithOpts(spellchecker.WithScoreFunc(scoreFunc))
+	err = sc.WithOpts(spellchecker.WithFilterFunc(fn))
 	if err != nil {
 		// handle err
 	}

dictionary.go

@@ -7,12 +7,9 @@ import (
 	"sort"
 	"sync/atomic"
 
-	"github.com/agext/levenshtein"
 	"github.com/f1monkey/bitmap"
 )
 
-type scoreFunc func(src []rune, candidate []rune, distance int, cnt uint) float64
-
 type dictionary struct {
 	maxErrors int
 	alphabet  alphabet
@@ -24,24 +21,24 @@ type dictionary struct {
 
 	index map[uint64][]uint32
 
-	scoreFunc scoreFunc
+	filterFunc FilterFunc
 }
 
-func newDictionary(ab string, scoreFunc scoreFunc, maxErrors int) (*dictionary, error) {
+func newDictionary(ab string, filterFunc FilterFunc, maxErrors int) (*dictionary, error) {
 	alphabet, err := newAlphabet(ab)
 	if err != nil {
 		return nil, err
 	}
 
 	return &dictionary{
-		maxErrors: maxErrors,
-		alphabet:  alphabet,
-		nextID:    idSeq(0),
-		ids:       make(map[string]uint32),
-		words:     make(map[uint32]string),
-		counts:    make(map[uint32]uint),
-		index:     make(map[uint64][]uint32),
-		scoreFunc: scoreFunc,
+		maxErrors:  maxErrors,
+		alphabet:   alphabet,
+		nextID:     idSeq(0),
+		ids:        make(map[string]uint32),
+		words:      make(map[uint32]string),
+		counts:     make(map[uint32]uint),
+		index:      make(map[uint64][]uint32),
+		filterFunc: filterFunc,
 	}, nil
 }
 
@@ -60,10 +57,9 @@ func (d *dictionary) add(word string, n uint) (uint32, error) {
 	id := d.nextID()
 	d.ids[word] = id
 
-	runes := []rune(word)
 	d.counts[id] = n
 	d.words[id] = word
-	key := sum(d.alphabet.encode(runes))
+	key := sum(d.alphabet.encode([]rune(word)))
 	d.index[key] = append(d.index[key], id)
 
 	return id, nil
@@ -148,14 +144,14 @@ func (d *dictionary) fillWithCandidates(result *priorityQueue, wordRunes []rune,
 			continue
 		}
 
-		distance, _, _ := levenshtein.Calculate(wordRunes, []rune(docWord), 0, 1, 1, 1)
-		if distance > d.maxErrors {
+		score, ok := d.filterFunc(wordRunes, []rune(docWord), d.counts[id])
+		if !ok {
 			continue
 		}
 
 		result.Push(Match{
 			Value: docWord,
-			Score: d.scoreFunc(wordRunes, []rune(docWord), distance, d.counts[id]),
+			Score: score,
 		})
 	}
 }
@@ -206,7 +202,7 @@ func (d *dictionary) UnmarshalBinary(data []byte) error {
 	d.words = dictData.Words
 	d.index = dictData.Index
 	d.maxErrors = dictData.MaxErrors
-	d.scoreFunc = defaultScorefunc
+	d.filterFunc = defaultFilterFunc(dictData.MaxErrors)
 
 	var max uint32
 	for _, id := range d.ids {

dictionary_test.go

@@ -7,7 +7,7 @@ import (
 )
 
 func Test_dictionary_id(t *testing.T) {
-	dict, err := newDictionary(DefaultAlphabet, defaultScorefunc, DefaultMaxErrors)
+	dict, err := newDictionary(DefaultAlphabet, nil, DefaultMaxErrors)
 	require.NoError(t, err)
 
 	t.Run("must return 0 for unexisting word", func(t *testing.T) {
@@ -24,7 +24,7 @@ func Test_dictionary_id(t *testing.T) {
 
 func Test_dictionary_add(t *testing.T) {
 	t.Run("must add word to dictionary index", func(t *testing.T) {
-		dict, err := newDictionary(DefaultAlphabet, defaultScorefunc, DefaultMaxErrors)
+		dict, err := newDictionary(DefaultAlphabet, nil, DefaultMaxErrors)
 		require.NoError(t, err)
 
 		id, err := dict.add("qwe", 1)
@@ -49,7 +49,7 @@ func Test_dictionary_add(t *testing.T) {
 
 func Test_Dictionary_Inc(t *testing.T) {
 	t.Run("must increase counter value", func(t *testing.T) {
-		dict, err := newDictionary(DefaultAlphabet, defaultScorefunc, DefaultMaxErrors)
+		dict, err := newDictionary(DefaultAlphabet, nil, DefaultMaxErrors)
 		dict.counts[1] = 0
 		require.NoError(t, err)
 

options.go

@@ -3,8 +3,30 @@ package spellchecker
 import (
 	"bufio"
 	"math"
+
+	"github.com/agext/levenshtein"
 )
 
+// WithOpt set spellchecker options
+func (s *Spellchecker) WithOpts(opts ...OptionFunc) error {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	for _, o := range opts {
+		if err := o(s); err != nil {
+			return err
+		}
+	}
+
+	if s.scoreFunc != nil {
+		s.dict.filterFunc = wrapScoreFunc(s.scoreFunc, s.maxErrors)
+	} else {
+		s.dict.filterFunc = s.filterFunc
+	}
+
+	return nil
+}
+
 // WithSplitter set splitter func for AddFrom() reader
 func WithSplitter(f bufio.SplitFunc) OptionFunc {
 	return func(s *Spellchecker) error {
@@ -15,28 +37,82 @@ func WithSplitter(f bufio.SplitFunc) OptionFunc {
 
 // WithMaxErrors sets maxErrors — the maximum allowed difference in bits
 // between the "search word" and a "dictionary word".
-// For example, replacing a single character (problam => problem)
-// is treated as a two-bit difference.
-// It is not recommended to set a value greater than 2,
-// as it can significantly impact performance.
+// - deletion is a 1-bit change (proble → problem)
+// - insertion is a 1-bit change (problemm → problem)
+// - substitution is a 2-bit change (problam → problem)
+// - transposition is a 0-bit change (problme → problem)
+//
+// It is not recommended to set this value greater than 2,
+// as it can significantly affect performance.
 func WithMaxErrors(maxErrors int) OptionFunc {
 	return func(s *Spellchecker) error {
 		s.maxErrors = maxErrors
+
 		return nil
 	}
 }
 
-type ScoreFunc = scoreFunc
+// FilterFunc compares the source word with a candidate word.
+// It returns the candidate's score and a boolean flag.
+// If the flag is false, the candidate will be completely filtered out.
+type FilterFunc func(src, candidate []rune, count uint) (float64, bool)
+
+// WithFilterFunc set custom scoring function
+func WithFilterFunc(f FilterFunc) OptionFunc {
+	return func(s *Spellchecker) error {
+		s.filterFunc = f
+		return nil
+	}
+}
+
+// ScoreFunc custom scoring function type
+//
+// Deprecated: use FilterFunc instead
+type ScoreFunc func(src []rune, candidate []rune, distance int, cnt uint) float64
 
 // WithScoreFunc specify a function that will be used for scoring
+//
+// Deprecated: use WithFilterFunc instead
 func WithScoreFunc(f ScoreFunc) OptionFunc {
 	return func(s *Spellchecker) error {
-		s.dict.scoreFunc = f
+		s.scoreFunc = f
 		return nil
 	}
 }
 
-var defaultScorefunc scoreFunc = func(src, candidate []rune, distance int, cnt uint) float64 {
+func defaultFilterFunc(maxErrors int) FilterFunc {
+	return func(src, candidate []rune, count uint) (float64, bool) {
+		distance, _, _ := levenshtein.Calculate(src, candidate, 0, 1, 1, 1)
+		if distance > maxErrors {
+			return 0, false
+		}
+
+		mult := math.Log1p(float64(count))
+		// if first letters are the same, increase score
+		if src[0] == candidate[0] {
+			mult *= 1.5
+			// if second letters are the same too, increase score  even more
+			if len(src) > 1 && len(candidate) > 1 && src[1] == candidate[1] {
+				mult *= 1.5
+			}
+		}
+
+		return 1 / (1 + float64(distance*distance)) * mult, true
+	}
+}
+
+func wrapScoreFunc(f ScoreFunc, maxErrors int) FilterFunc {
+	return func(src, candidate []rune, count uint) (float64, bool) {
+		distance, _, _ := levenshtein.Calculate(src, candidate, 0, 1, 1, 1)
+		if distance > maxErrors {
+			return 0, false
+		}
+
+		return f(src, candidate, distance, count), true
+	}
+}
+
+var defaultScoreFunc ScoreFunc = func(src, candidate []rune, distance int, cnt uint) float64 {
 	mult := math.Log1p(float64(cnt))
 	// if first letters are the same, increase score
 	if src[0] == candidate[0] {

spellchecker.go

@@ -15,29 +15,36 @@ type OptionFunc func(s *Spellchecker) error
 type Spellchecker struct {
 	mtx sync.RWMutex
 
-	dict      *dictionary
-	splitter  bufio.SplitFunc
-	scoreFunc scoreFunc
-	maxErrors int
+	dict       *dictionary
+	splitter   bufio.SplitFunc
+	filterFunc FilterFunc
+	scoreFunc  ScoreFunc
+	maxErrors  int
 }
 
 func New(alphabet string, opts ...OptionFunc) (*Spellchecker, error) {
 	result := &Spellchecker{
-		maxErrors: DefaultMaxErrors,
-		scoreFunc: defaultScorefunc,
+		maxErrors:  DefaultMaxErrors,
+		filterFunc: defaultFilterFunc(DefaultMaxErrors),
 	}
-	dict, err := newDictionary(alphabet, result.scoreFunc, result.maxErrors)
-	if err != nil {
-		return nil, err
-	}
-	result.dict = dict
 
 	for _, o := range opts {
 		if err := o(result); err != nil {
 			return nil, err
 		}
 	}
 
+	if result.scoreFunc != nil {
+		result.filterFunc = wrapScoreFunc(result.scoreFunc, result.maxErrors)
+	}
+
+	dict, err := newDictionary(alphabet, result.filterFunc, result.maxErrors)
+	if err != nil {
+		return nil, err
+	}
+
+	result.dict = dict
+
 	return result, nil
 }
 
@@ -162,17 +169,3 @@ func (s *Spellchecker) SuggestScore(word string, n int) SuggestionResult {
 		Suggestions: s.dict.find(word, n),
 	}
 }
-
-// WithOpt set spellchecker options
-func (s *Spellchecker) WithOpts(opts ...OptionFunc) error {
-	s.mtx.Lock()
-	defer s.mtx.Unlock()
-
-	for _, o := range opts {
-		if err := o(s); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}

spellchecker_test.go

@@ -164,7 +164,6 @@ func benchmarkNorvig(b *testing.B, dataPath string) {
 	total := 0
 	ok := 0
 	for i := 0; i < b.N; i++ {
-
 		for _, item := range data {
 			for _, word := range item.words {
 				if word == "" {
@@ -236,6 +235,15 @@ func Test_Spellchecker_Fix(t *testing.T) {
 	require.Equal(t, "problem", result)
 }
 
+func Test_Spellchecker_Fix_ScoreFunc(t *testing.T) {
+	s := newSampleSpellchecker()
+	s.WithOpts(WithScoreFunc(defaultScoreFunc))
+
+	result, err := s.Fix("problam")
+	require.NoError(t, err)
+	require.Equal(t, "problem", result)
+}
+
 func Test_Spellchecker_Suggest(t *testing.T) {
 	s := newSampleSpellchecker()
 	result, err := s.Suggest("arang", 5)