RawMark/internal/search/search.go

// Package search 维护笔记的全文索引。索引存放在独立的数据库文件中，内容完全由仓库文件派生，
// 删除后会在下次启动时重建；结构变化时也直接重建而不做迁移。
package search

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"unicode"
	"unicode/utf8"

	_ "modernc.org/sqlite"

	"rawmark/internal/vault"
)

// 修改表结构时递增，旧索引会被丢弃重建
const schemaVersion = 1

const schema = `
CREATE TABLE docs (
  id   INTEGER PRIMARY KEY,
  path TEXT NOT NULL UNIQUE,
  rev  TEXT NOT NULL
);
-- trigram 分词不依赖空格断词，中文可按任意片段检索；不足 3 个字符的词改用 LIKE
CREATE VIRTUAL TABLE docs_fts USING fts5(path, body, tokenize = 'trigram');
`

const (
	maxTerms      = 8
	snippetRunes  = 160
	snippetBefore = 40
)

type Result struct {
	Path string `json:"path"`
	// 第一处匹配所在行（从 1 开始）与该行的摘录；只有文件名匹配时 Line 为 0
	Line    int    `json:"line"`
	Snippet string `json:"snippet"`
}

type doc struct {
	id  int64
	rev string
}

type Index struct {
	db *sql.DB
	v  *vault.Vault

	mu      sync.Mutex
	pending map[string]string // 待同步的最新快照（路径 -> rev）
	wake    chan struct{}

	known map[string]doc // 已索引的笔记，仅由同步协程访问
}

func Open(path string, v *vault.Vault) (*Index, error) {
	db, err := sql.Open("sqlite", "file:"+path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)")
	if err != nil {
		return nil, err
	}
	x := &Index{db: db, v: v, wake: make(chan struct{}, 1), known: map[string]doc{}}
	if err := x.init(); err != nil {
		db.Close()
		return nil, fmt.Errorf("初始化搜索索引失败: %w", err)
	}
	return x, nil
}

func (x *Index) init() error {
	var version int
	if err := x.db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil {
		return err
	}
	if version != schemaVersion {
		stmts := []string{`DROP TABLE IF EXISTS docs`, `DROP TABLE IF EXISTS docs_fts`, schema, fmt.Sprintf(`PRAGMA user_version = %d`, schemaVersion)}
		for _, s := range stmts {
			if _, err := x.db.Exec(s); err != nil {
				return err
			}
		}
	}
	rows, err := x.db.Query(`SELECT id, path, rev FROM docs`)
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var d doc
		var p string
		if err := rows.Scan(&d.id, &p, &d.rev); err != nil {
			return err
		}
		x.known[p] = d
	}
	return rows.Err()
}

func (x *Index) Close() error { return x.db.Close() }

// Update 提交仓库的最新快照（笔记路径 -> rev），立即返回；同步在后台进行，多次提交只处理最新一次
func (x *Index) Update(files map[string]string) {
	x.mu.Lock()
	x.pending = files
	x.mu.Unlock()
	select {
	case x.wake <- struct{}{}:
	default:
	}
}

// Run 执行后台同步，直到 ctx 结束
func (x *Index) Run(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-x.wake:
			x.mu.Lock()
			files := x.pending
			x.mu.Unlock()
			if err := x.sync(files); err != nil {
				slog.Error("更新搜索索引失败", "err", err)
			}
		}
	}
}

// sync 让索引与快照一致：只读取新增或 rev 变化的笔记
func (x *Index) sync(files map[string]string) error {
	tx, err := x.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	next := make(map[string]doc, len(x.known))
	for p, d := range x.known {
		next[p] = d
	}
	for p, d := range x.known {
		if _, ok := files[p]; !ok || vault.IsEncrypted(p) {
			if err := remove(tx, d.id); err != nil {
				return err
			}
			delete(next, p)
		}
	}
	for p, rev := range files {
		// 加密笔记的内容服务端无法读取，不入索引
		if vault.IsEncrypted(p) {
			continue
		}
		old, ok := x.known[p]
		if ok && old.rev == rev {
			continue
		}
		content, rev, err := x.v.Read(p)
		if errors.Is(err, vault.ErrNotFound) {
			continue // 快照之后被删除，下一次快照会处理
		}
		if err != nil {
			return err
		}
		if ok {
			if err := remove(tx, old.id); err != nil {
				return err
			}
		}
		res, err := tx.Exec(`INSERT INTO docs (path, rev) VALUES (?, ?)`, p, rev)
		if err != nil {
			return err
		}
		id, err := res.LastInsertId()
		if err != nil {
			return err
		}
		if _, err := tx.Exec(`INSERT INTO docs_fts (rowid, path, body) VALUES (?, ?, ?)`, id, p, string(content)); err != nil {
			return err
		}
		next[p] = doc{id: id, rev: rev}
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	x.known = next
	return nil
}

func remove(tx *sql.Tx, id int64) error {
	if _, err := tx.Exec(`DELETE FROM docs WHERE id = ?`, id); err != nil {
		return err
	}
	_, err := tx.Exec(`DELETE FROM docs_fts WHERE rowid = ?`, id)
	return err
}

// Search 查找同时包含所有词（空白分隔，不区分大小写）的笔记，词可出现在路径或正文中。
// 路径（文件名或目录名）命中第一个词的排在前面，其余按相关度排序
func (x *Index) Search(query string, limit int) ([]Result, error) {
	terms := strings.Fields(query)
	if len(terms) == 0 {
		return []Result{}, nil
	}
	if len(terms) > maxTerms {
		terms = terms[:maxTerms]
	}
	var where, phrases []string
	var args []any
	for _, t := range terms {
		if utf8.RuneCountInString(t) >= 3 {
			phrases = append(phrases, `"`+strings.ReplaceAll(t, `"`, `""`)+`"`)
		} else {
			like := "%" + escapeLike(t) + "%"
			where = append(where, `(path LIKE ? ESCAPE '\' OR body LIKE ? ESCAPE '\')`)
			args = append(args, like, like)
		}
	}
	order := "path"
	if len(phrases) > 0 {
		where = append(where, `docs_fts MATCH ?`)
		args = append(args, strings.Join(phrases, " AND "))
		order = "rank"
	}
	name := "%" + escapeLike(terms[0]) + "%"
	q := `SELECT path, body FROM docs_fts WHERE ` + strings.Join(where, " AND ") +
		` ORDER BY path LIKE ? ESCAPE '\' DESC, ` + order + ` LIMIT ?`
	args = append(args, name, limit)
	rows, err := x.db.Query(q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	results := []Result{}
	for rows.Next() {
		var r Result
		var body string
		if err := rows.Scan(&r.Path, &body); err != nil {
			return nil, err
		}
		r.Line, r.Snippet = snippet(body, terms)
		results = append(results, r)
	}
	return results, rows.Err()
}

func escapeLike(s string) string {
	return strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`).Replace(s)
}

// snippet 返回正文中第一处匹配所在的行号与摘录（过长的行截取匹配附近的一段）
func snippet(body string, terms []string) (int, string) {
	folded := make([][]rune, len(terms))
	for i, t := range terms {
		folded[i] = fold([]rune(t))
	}
	for n, line := range strings.Split(body, "\n") {
		runes := []rune(strings.TrimRight(line, "\r"))
		lower := fold(runes)
		at := -1
		for _, t := range folded {
			if i := indexRunes(lower, t); i >= 0 && (at < 0 || i < at) {
				at = i
			}
		}
		if at < 0 {
			continue
		}
		start := 0
		if len(runes) > snippetRunes {
			start = max(0, min(at-snippetBefore, len(runes)-snippetRunes))
		}
		end := min(len(runes), start+snippetRunes)
		text := strings.TrimSpace(string(runes[start:end]))
		if start > 0 {
			text = "…" + text
		}
		if end < len(runes) {
			text += "…"
		}
		return n + 1, text
	}
	return 0, ""
}

// fold 逐字符转小写，保持字符数不变，便于在原文中按位置截取
func fold(rs []rune) []rune {
	out := make([]rune, len(rs))
	for i, r := range rs {
		out[i] = unicode.ToLower(r)
	}
	return out
}

func indexRunes(s, sub []rune) int {
	for i := 0; i+len(sub) <= len(s); i++ {
		match := true
		for j := range sub {
			if s[i+j] != sub[j] {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}