敏感词检测

2025-09-03 16:37:43 +08:00
parent 2e8310a3ff
commit f956b85ed3
6 changed files with 484 additions and 0 deletions
--- a/internal/lexicon/store.go
+++ b/internal/lexicon/store.go
@@ -0,0 +1,141 @@
+package lexicon
+
+import (
+	"bufio"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/ozeidan/fuzzy-patricia/patricia"
+)
+
+// Store holds the trie and statistics for the loaded lexicon.
+type Store struct {
+	mu   sync.RWMutex
+	trie *patricia.Trie
+	cnt  int
+}
+
+func NewStore() *Store {
+	return &Store{trie: patricia.NewTrie()}
+}
+
+// LoadFromDir loads all .txt files from dir into the trie.
+func (s *Store) LoadFromDir(dir string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	newTrie := patricia.NewTrie()
+	count := 0
+
+	walkErr := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if info.IsDir() {
+			return nil
+		}
+		if strings.ToLower(filepath.Ext(info.Name())) != ".txt" {
+			return nil
+		}
+		f, err := os.Open(path)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		scanner := bufio.NewScanner(f)
+		// Increase buffer for long lines
+		buf := make([]byte, 0, 1024*64)
+		scanner.Buffer(buf, 1024*1024)
+		for scanner.Scan() {
+			w := strings.TrimSpace(scanner.Text())
+			if w == "" || strings.HasPrefix(w, "#") {
+				continue
+			}
+			newTrie.Insert(patricia.Prefix(w), struct{}{})
+			count++
+		}
+		return scanner.Err()
+	})
+	if walkErr != nil {
+		return walkErr
+	}
+	if count == 0 {
+		return errors.New("no entries loaded")
+	}
+	// Swap in
+	s.trie = newTrie
+	s.cnt = count
+	return nil
+}
+
+func (s *Store) Stats() map[string]interface{} {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return map[string]interface{}{
+		"count": s.cnt,
+	}
+}
+
+// ForEachSubstringMatch visits any keys that contain the given substring.
+// It uses the library's substring search.
+func (s *Store) ForEachSubstringMatch(query string, visit func(word string) bool) {
+	s.mu.RLock()
+	tr := s.trie
+	s.mu.RUnlock()
+	if tr == nil || query == "" {
+		return
+	}
+	// second argument is caseSensitive; we use false by default
+	tr.VisitSubstring(patricia.Prefix(query), false, func(prefix patricia.Prefix, _ patricia.Item) error {
+		// The library does not expose a public stop error in all versions; ignore early stop
+		_ = visit(string(prefix))
+		return nil
+	})
+}
+
+// ForEachFuzzyMatch visits keys with fuzzy distance within maxDistance to query.
+func (s *Store) ForEachFuzzyMatch(query string, maxDistance int, visit func(word string, distance int) bool) {
+	s.mu.RLock()
+	tr := s.trie
+	s.mu.RUnlock()
+	if tr == nil || query == "" {
+		return
+	}
+	// signature in current lib: VisitFuzzy(prefix, caseSensitive bool, visitor)
+	tr.VisitFuzzy(patricia.Prefix(query), false, func(prefix patricia.Prefix, _ patricia.Item, dist int) error {
+		if dist <= maxDistance {
+			_ = visit(string(prefix), dist)
+		}
+		return nil
+	})
+}
+
+// HasAnyInText returns true if any lexicon word is a substring of the given text.
+// It scans each rune offset and visits prefixes against the trie.
+func (s *Store) HasAnyInText(text string) (bool, string) {
+	s.mu.RLock()
+	tr := s.trie
+	s.mu.RUnlock()
+	if tr == nil || text == "" {
+		return false, ""
+	}
+	runes := []rune(text)
+	n := len(runes)
+	for i := 0; i < n; i++ {
+		suffix := string(runes[i:])
+		foundWord := ""
+		tr.VisitPrefixes(patricia.Prefix(suffix), false, func(prefix patricia.Prefix, _ patricia.Item) error {
+			if foundWord == "" {
+				foundWord = string(prefix)
+			}
+			return nil
+		})
+		if foundWord != "" {
+			return true, foundWord
+		}
+	}
+	return false, ""
+}