kitty/kittens/choose_files/scan.go

package choose_files

import (
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"sync"
	"time"
	"unicode"
	"unsafe"

	"github.com/kovidgoyal/kitty/tools/fzf"
	"github.com/kovidgoyal/kitty/tools/utils"
)

var _ = fmt.Print

type ResultItem struct {
	text, ltext, abspath string
	ftype                fs.FileMode
	positions            []int // may be nil
	score                float64
	positions_sorted     bool
}

func (r ResultItem) String() string {
	return fmt.Sprintf("{text: %#v, abspath: %#v, is_dir: %v, positions: %#v}", r.text, r.abspath, r.ftype.IsDir(), r.positions)
}

func (r *ResultItem) sorted_positions() []int {
	if !r.positions_sorted {
		r.positions_sorted = true
		if len(r.positions) > 1 {
			sort.Ints(r.positions)
		}
	}
	return r.positions
}

func get_modified_score(abspath string, score float64, score_patterns []ScorePattern) float64 {
	for _, sp := range score_patterns {
		if sp.pat.MatchString(abspath) {
			score = sp.op(score, sp.val)
		}
	}
	return score
}

func count_uppercase(s string) int {
	count := 0
	for _, r := range s {
		if unicode.IsUpper(r) {
			count++
		}
	}
	return count
}

type ScanRequest struct {
	root_dir string
}

type ScanResult struct {
	root_dir    string
	items       []ResultItem
	err         error
	is_finished bool
}

type ScoreRequest struct {
	root_dir, query              string
	is_last_for_current_root_dir bool
	items                        []ResultItem
}

type ScoreResult struct {
	query, root_dir              string
	is_last_for_current_root_dir bool
	items                        []ResultItem
}

type Settings interface {
	OnlyDirs() bool
	ScorePatterns() []ScorePattern
	CurrentDir() string
	SearchText() string
}

type ResultManager struct {
	current_root_dir               string
	current_root_dir_scan_complete bool
	results_for_current_root_dir   []ResultItem
	scan_requests                  chan ScanRequest
	scan_results                   chan ScanResult

	current_query                  string
	current_query_scoring_complete bool
	matches_for_current_query      []ResultItem
	score_queries                  chan ScoreRequest
	score_results                  chan ScoreResult
	report_errors                  chan error

	renderable_results []ResultItem

	mutex    sync.Mutex
	scorer   *fzf.FuzzyMatcher
	settings Settings

	WakeupMainThread func() bool
}

func NewResultManager(err_chan chan error, settings Settings, WakeupMainThread func() bool) *ResultManager {
	ans := &ResultManager{
		scan_requests:    make(chan ScanRequest, 256),
		scan_results:     make(chan ScanResult, 256),
		score_queries:    make(chan ScoreRequest, 256),
		score_results:    make(chan ScoreResult, 256),
		report_errors:    err_chan,
		scorer:           fzf.NewFuzzyMatcher(fzf.PATH_SCHEME),
		settings:         settings,
		WakeupMainThread: WakeupMainThread,
	}
	go ans.scan_worker()
	go ans.scan_result_handler()
	go ans.score_worker()
	go ans.sort_worker()
	return ans
}

func (m *ResultManager) lock() {
	m.mutex.Lock()
}

func (m *ResultManager) unlock() {
	m.mutex.Unlock()
}

func (m *ResultManager) scan(dir, root_dir string, level int) (err error) {
	defer func() {
		if r := recover(); r != nil {
			st, qerr := utils.Format_stacktrace_on_panic(r)
			err = fmt.Errorf("%w\n%s", qerr, st)
		}
	}()
	items, err := os.ReadDir(dir)
	if err != nil {
		if level == 0 {
			return fmt.Errorf("failed to read directory: %s with error: %w", dir, err)
		}
		return nil
	}
	m.scan_results <- ScanResult{root_dir: root_dir, items: utils.Map(func(x os.DirEntry) ResultItem {
		return ResultItem{abspath: filepath.Join(dir, x.Name()), ftype: x.Type()}
	}, items)}
	for _, x := range items {
		if x.IsDir() {
			if !m.is_root_dir_current(root_dir) {
				return
			}
			if err = m.scan(filepath.Join(dir, x.Name()), root_dir, level+1); err != nil {
				return
			}
		}
	}
	return
}

func (m *ResultManager) scan_worker() {
	for r := range m.scan_requests {
		st := time.Now()
		if err := m.scan(r.root_dir, r.root_dir, 0); err == nil {
			m.scan_results <- ScanResult{root_dir: r.root_dir, is_finished: true}
		}
		debugprintln(111111111, time.Now().Sub(st), len(m.results_for_current_root_dir))
	}
}

func (m *ResultManager) create_score_query(items []ResultItem, is_finished bool) ScoreRequest {
	return ScoreRequest{root_dir: m.current_root_dir, query: m.current_query, items: utils.Map(
		func(r ResultItem) ResultItem {
			text, err := filepath.Rel(m.current_root_dir, r.abspath)
			if err != nil {
				text = r.abspath
			}
			return ResultItem{abspath: r.abspath, text: text, ltext: strings.ToLower(text), ftype: r.ftype}
		}, items), is_last_for_current_root_dir: is_finished}
}

func (m *ResultManager) scan_result_handler() {
	one := func(r ScanResult) {
		if !m.is_root_dir_current(r.root_dir) {
			return
		}
		var sqr ScoreRequest
		has_items := len(r.items) > 0
		m.lock()
		if has_items {
			m.results_for_current_root_dir = append(m.results_for_current_root_dir, r.items...)
		}
		sqr = m.create_score_query(r.items, r.is_finished)
		if r.is_finished {
			m.current_root_dir_scan_complete = true
		}
		m.unlock()
		m.score_queries <- sqr
	}
	for r := range m.scan_results {
		if r.err != nil {
			m.report_errors <- r.err
			continue
		}
		one(r)
	}
}

func (m *ResultManager) score(r ScoreRequest) (err error) {
	items := r.items
	m.lock()
	only_dirs := m.settings.OnlyDirs()
	sp := m.settings.ScorePatterns()
	m.unlock()
	if only_dirs {
		items = utils.Filter(items, func(r ResultItem) bool { return r.ftype.IsDir() })
	}
	res := ScoreResult{query: r.query, items: items, root_dir: r.root_dir, is_last_for_current_root_dir: r.is_last_for_current_root_dir}
	if r.query != "" {
		var r []fzf.Result
		if r, err = m.scorer.ScoreWithCache(utils.Map(func(r ResultItem) string { return r.text }, items), res.query); err != nil {
			return
		}
		for i, x := range r {
			items[i].positions = x.Positions
			items[i].score = get_modified_score(items[i].abspath, float64(x.Score), sp)
		}
		items = utils.Filter(items, func(r ResultItem) bool { return r.score > 0 })
	}
	m.score_results <- res
	return
}

func int_cmp(a, b int) int {
	if a < b {
		return -1
	}
	if a > b {
		return 1
	}
	return 0
}

func str_cmp(a, b string) int {
	if a < b {
		return -1
	}
	if a > b {
		return 1
	}
	return 0
}

func float_cmp(a, b float64) int { // deliberately doesnt handle NaN
	if a < b {
		return -1
	}
	if a > b {
		return 1
	}
	return 0
}

func bool_as_int(b bool) int {
	return *(*int)(unsafe.Pointer(&b))
}

func bool_cmp(a, b bool) int {
	return bool_as_int(a) - bool_as_int(b)
}

func (m *ResultManager) score_worker() {
	for r := range m.score_queries {
		if m.is_query_current(r.query, r.root_dir) {
			if err := m.score(r); err != nil {
				m.report_errors <- err
			}
		}
	}
}

func cmp_with_score(a, b ResultItem) (ans int) {
	ans = float_cmp(b.score, a.score)
	if ans == 0 {
		ans = int_cmp(len(a.text), len(b.text))
		if ans == 0 {
			ans = int_cmp(count_uppercase(a.text), count_uppercase(b.text))
		}
	}
	return
}

func cmp_without_score(a, b ResultItem) (ans int) {
	ans = bool_cmp(a.ftype.IsDir(), b.ftype.IsDir())
	if ans == 0 {
		ans = str_cmp(a.ltext, b.ltext)
		if ans == 0 {
			ans = int_cmp(count_uppercase(a.text), count_uppercase(b.text))
		}
	}
	return
}

func merge_sorted_slices(a, b []ResultItem, cmp func(a, b ResultItem) int) []ResultItem {
	result := make([]ResultItem, 0, len(a)+len(b))
	i, j := 0, 0
	for i < len(a) && j < len(b) {
		if cmp(a[i], b[j]) <= 0 {
			result = append(result, a[i])
			i++
		} else {
			result = append(result, b[j])
			j++
		}
	}
	result = append(result, a[i:]...)
	return append(result, b[j:]...)
}

func (m *ResultManager) add_score_results(r ScoreResult) (err error) {
	cmp := utils.IfElse(r.query == "", cmp_without_score, cmp_with_score)
	slices.SortStableFunc(r.items, cmp)
	m.lock()
	defer func() {
		m.unlock()
		if r := recover(); r != nil {
			st, qerr := utils.Format_stacktrace_on_panic(r)
			err = fmt.Errorf("%w\n%s", qerr, st)
		}
	}()
	m.renderable_results = merge_sorted_slices(m.renderable_results, r.items, cmp)
	if r.is_last_for_current_root_dir {
		m.current_query_scoring_complete = true
	}
	return
}

func (m *ResultManager) sort_worker() {
	last_wakeup_at := time.Now()
	for r := range m.score_results {
		if m.is_query_current(r.query, r.root_dir) {
			if err := m.add_score_results(r); err != nil {
				m.report_errors <- err
			} else {
				m.lock()
				is_complete := m.current_root_dir_scan_complete && m.current_query_scoring_complete
				m.unlock()
				if is_complete || time.Now().Sub(last_wakeup_at) > time.Millisecond*50 {
					m.WakeupMainThread()
					last_wakeup_at = time.Now()
				}
			}
		}
	}
}

func (m *ResultManager) is_root_dir_current(root_dir string) bool {
	m.lock()
	defer m.unlock()
	return root_dir == m.current_root_dir
}

func (m *ResultManager) set_root_dir(root_dir string) {
	var err error
	if root_dir == "" || root_dir == "." {
		if root_dir, err = os.Getwd(); err != nil {
			return
		}
	}
	root_dir = utils.Expanduser(root_dir)
	if root_dir, err = filepath.Abs(root_dir); err != nil {
		return
	}
	m.lock()
	defer m.unlock()
	if m.current_root_dir == root_dir {
		return
	}
	m.current_root_dir = root_dir
	m.results_for_current_root_dir = nil
	m.matches_for_current_query = nil
	m.renderable_results = nil
	m.current_query_scoring_complete = false
	m.current_root_dir_scan_complete = false
	m.scan_requests <- ScanRequest{root_dir: m.current_root_dir}
}

func (m *ResultManager) is_query_current(query, root_dir string) bool {
	m.lock()
	defer m.unlock()
	return root_dir == m.current_root_dir && query == m.current_query
}

func (m *ResultManager) set_query(query string) {
	var sqr *ScoreRequest
	m.lock()
	defer func() {
		m.unlock()
		if sqr != nil {
			m.score_queries <- *sqr
		}
	}()
	if query == m.current_query {
		return
	}
	m.current_query = query
	m.matches_for_current_query = nil
	m.renderable_results = nil
	m.current_query_scoring_complete = false
	if m.results_for_current_root_dir != nil {
		s := m.create_score_query(m.results_for_current_root_dir, m.current_root_dir_scan_complete)
		sqr = &s
	} else if m.current_root_dir_scan_complete {
		m.current_query_scoring_complete = true
	}
}

func (h *Handler) get_results() (ans []ResultItem, in_progress bool) {
	h.result_manager.lock()
	defer h.result_manager.unlock()
	ans = h.result_manager.renderable_results
	in_progress = !h.result_manager.current_query_scoring_complete || !h.result_manager.current_root_dir_scan_complete
	return
}