kitty/kittens/choose_files/scan.go

package choose_files

import (
	"bytes"
	"cmp"
	"fmt"
	"io/fs"
	"math"
	"os"
	"path/filepath"
	"slices"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unicode"
	"unicode/utf8"

	"github.com/kovidgoyal/kitty/tools/fzf"
	"github.com/kovidgoyal/kitty/tools/utils"
	"golang.org/x/sys/unix"
)

var _ = fmt.Print

func (c CombinedScore) String() string {
	return fmt.Sprintf("{score: %d length: %d index: %d}", c.Score(), c.Length(), c.Index())
}

type ResultItem struct {
	text      string
	ftype     fs.FileMode
	positions []int // may be nil
	score     CombinedScore
}
type ResultsType []*ResultItem

func (r *ResultItem) SetScoreResult(x fzf.Result) {
	r.positions = x.Positions
	r.score.Set_score(uint16(math.MaxUint16 - uint16(x.Score)))
}

func (r ResultItem) IsMatching() bool {
	return r.score.Score() < uint16(math.MaxUint16)
}

func (r ResultItem) String() string {
	return fmt.Sprintf("{text: %#v, %s, positions: %#v}", r.text, r.score, r.positions)
}

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

type FileSystemScanner struct {
	listeners               []chan bool
	in_progress, keep_going atomic.Bool
	root_dir                string
	mutex                   sync.Mutex
	results                 []ResultItem
	dir_reader              func(path string) ([]fs.DirEntry, error)
	err                     error
}

func NewFileSystemScanner(root_dir string, notify chan bool) (fss *FileSystemScanner) {
	ans := &FileSystemScanner{root_dir: root_dir, listeners: []chan bool{notify}, results: make([]ResultItem, 0, 1024)}
	ans.in_progress.Store(true)
	ans.keep_going.Store(true)
	ans.dir_reader = os.ReadDir
	return ans
}

type Scanner interface {
	Start()
	Cancel()
	AddListener(chan bool)
	Len() int
	Batch(offset int) []ResultItem
	Finished() bool
	Error() error
}

func (fss *FileSystemScanner) lock()   { fss.mutex.Lock() }
func (fss *FileSystemScanner) unlock() { fss.mutex.Unlock() }

func (fss *FileSystemScanner) Error() error {
	fss.lock()
	defer fss.unlock()
	return fss.err
}

func (fss *FileSystemScanner) Start() {
	go fss.worker()
}

func (fss *FileSystemScanner) Cancel() {
	fss.keep_going.Store(false)
}

func (fss *FileSystemScanner) AddListener(x chan bool) {
	fss.lock()
	defer fss.unlock()
	if !fss.in_progress.Load() {
		close(x)
	} else {
		fss.listeners = append(fss.listeners, x)
	}
}

func (fss *FileSystemScanner) Len() int {
	fss.lock()
	defer fss.unlock()
	return len(fss.results)
}

func (fss *FileSystemScanner) Batch(offset int) []ResultItem {
	fss.lock()
	defer fss.unlock()
	if offset >= len(fss.results) {
		return nil
	}
	limit := min(len(fss.results), offset+4096)
	return fss.results[offset:limit]
}

func (fss *FileSystemScanner) Finished() bool {
	return !fss.in_progress.Load()
}

type sortable_dir_entry struct {
	name     string
	ftype    fs.FileMode
	sort_key []byte
	buf      [unix.NAME_MAX + 1]byte
}

const SymlinkToDir = 1

// lowercase a string into a pre-existing byte buffer with speedups for ASCII
func as_lower(s string, output []byte) int {
	limit := min(len(s), len(output))
	found_non_ascii := false
	pos := 0
	for i := range limit {
		c := s[i]
		if 'A' <= c && c <= 'Z' {
			c += 'a' - 'A'
			if pos < i {
				copy(output[pos:i], s[pos:i])
			}
			output[i] = c
			pos = i + 1
		} else if c >= utf8.RuneSelf {
			if pos < i {
				copy(output[pos:i], s[pos:i])
			}
			found_non_ascii = true
			pos = i
			break
		}
	}
	if !found_non_ascii {
		if pos < limit {
			copy(output[pos:limit], s[pos:limit])
		}
		return limit
	}
	buf := [4]byte{}
	var n int
	for _, r := range s[pos:] {
		o := output[pos:]
		r = unicode.ToLower(r)
		if len(o) > 3 {
			n = utf8.EncodeRune(o, r)
		} else {
			n = utf8.EncodeRune(buf[:], r)
			n = copy(o, buf[:n])
		}
		pos += n
		if pos >= len(output) {
			break
		}
	}
	return pos
}

func (fss *FileSystemScanner) worker() {
	defer func() {
		fss.lock()
		defer fss.unlock()
		fss.in_progress.Store(false)
		if r := recover(); r != nil {
			st, qerr := utils.Format_stacktrace_on_panic(r)
			fss.err = fmt.Errorf("%w\n%s", qerr, st)
		}
		for _, l := range fss.listeners {
			close(l)
		}
	}()
	root_dir, _ := filepath.Abs(fss.root_dir)
	dir := root_dir
	if !strings.HasSuffix(dir, string(os.PathSeparator)) {
		dir += string(os.PathSeparator)
	}
	base := ""
	pos := 0
	var arena []sortable_dir_entry
	var sortable []*sortable_dir_entry
	var idx uint32
	// do a breadth first traversal of the filesystem
	is_root := true
	for dir != "" {
		if !fss.keep_going.Load() {
			break
		}
		entries, err := fss.dir_reader(dir)
		if err != nil {
			if is_root {
				fss.keep_going.Store(false)
				fss.lock()
				fss.err = err
				fss.unlock()
			}
			entries = nil
		}
		if cap(arena) < len(entries) {
			arena = make([]sortable_dir_entry, 0, max(1024, len(entries), 2*cap(arena)))
			sortable = make([]*sortable_dir_entry, 0, cap(arena))
		}
		arena = arena[:len(entries)]
		sortable = sortable[:len(entries)]
		for i, e := range entries {
			arena[i].name = e.Name()
			ftype := e.Type()
			if ftype&fs.ModeSymlink != 0 {
				if st, serr := os.Stat(dir + arena[i].name); serr == nil && st.IsDir() {
					ftype |= SymlinkToDir
				}
			}
			arena[i].ftype = ftype
			if ftype&fs.ModeDir != 0 {
				arena[i].buf[0] = '0'
			} else {
				arena[i].buf[0] = '1'
			}
			n := as_lower(arena[i].name, arena[i].buf[1:])
			arena[i].sort_key = arena[i].buf[:1+n]
			sortable[i] = &arena[i]
		}
		slices.SortFunc(sortable, func(a, b *sortable_dir_entry) int { return bytes.Compare(a.sort_key, b.sort_key) })
		ns := fss.results
		new_sz := len(ns) + len(entries)
		if cap(ns) < new_sz {
			ns = make([]ResultItem, len(ns), max(1024, new_sz, cap(ns)*2))
			copy(ns, fss.results)
		}
		new_items := ns[len(ns):new_sz]
		for i, e := range sortable {
			new_items[i].ftype = e.ftype
			new_items[i].text = base + e.name
			new_items[i].score.Set_index(idx)
			idx++
		}
		ns = ns[0:new_sz]
		fss.lock()
		fss.results = ns
		listeners := fss.listeners
		fss.unlock()
		for _, l := range listeners {
			select {
			case l <- true:
			default:
			}
		}
		dir = ""
		for pos < len(fss.results) && dir == "" {
			if fss.results[pos].ftype&fs.ModeDir != 0 {
				base = fss.results[pos].text + string(os.PathSeparator)
				dir = root_dir + string(os.PathSeparator) + base
			}
			pos++
		}
		is_root = false
	}
}

type FileSystemScorer struct {
	scanner                 Scanner
	keep_going, is_complete atomic.Bool
	root_dir, query         string
	only_dirs               bool
	mutex                   sync.Mutex
	renderable_results      []*ResultItem
	on_results              func(error, bool)
	current_worker_wait     *sync.WaitGroup
	scorer                  *fzf.FuzzyMatcher
}

func NewFileSystemScorer(root_dir, query string, only_dirs bool, on_results func(error, bool)) (ans *FileSystemScorer) {
	return &FileSystemScorer{
		query: query, root_dir: root_dir, only_dirs: only_dirs, on_results: on_results,
		scorer: fzf.NewFuzzyMatcher(fzf.PATH_SCHEME)}
}

func (fss *FileSystemScorer) lock()   { fss.mutex.Lock() }
func (fss *FileSystemScorer) unlock() { fss.mutex.Unlock() }

func (fss *FileSystemScorer) Start() {
	on_results := make(chan bool)
	fss.is_complete.Store(false)
	fss.keep_going.Store(true)
	if fss.scanner == nil {
		fss.scanner = NewFileSystemScanner(fss.root_dir, on_results)
		fss.scanner.Start()
	} else {
		fss.scanner.AddListener(on_results)
	}
	fss.current_worker_wait = &sync.WaitGroup{}
	fss.current_worker_wait.Add(1)
	go fss.worker(on_results, fss.current_worker_wait)
}

func (fss *FileSystemScorer) Change_query(query string) {
	if fss.query == query {
		return
	}
	fss.keep_going.Store(false)
	if fss.current_worker_wait != nil {
		fss.current_worker_wait.Wait()
	}
	fss.lock()
	fss.query = query
	fss.renderable_results = nil
	fss.unlock()
	fss.Start()
}

func (fss *FileSystemScorer) worker(on_results chan bool, worker_wait *sync.WaitGroup) {
	defer func() {
		fss.is_complete.Store(true)
		defer worker_wait.Done()
		if r := recover(); r != nil {
			if fss.keep_going.Load() {
				st, qerr := utils.Format_stacktrace_on_panic(r)
				fss.on_results(fmt.Errorf("%w\n%s", qerr, st), true)
			}
		} else {
			if fss.keep_going.Load() {
				fss.on_results(fss.scanner.Error(), true)
			}
		}
	}()
	global_min_score, global_max_score := CombinedScore(math.MaxUint64), CombinedScore(0)
	handle_batch := func(results []ResultItem) (err error) {
		if err = fss.scanner.Error(); err != nil {
			return
		}
		var rp []*ResultItem
		if fss.only_dirs {
			rp = make([]*ResultItem, 0, len(results))
			for i, r := range results {
				if r.ftype.IsDir() {
					rp = append(rp, &results[i])
				}
			}
		} else {
			rp = make([]*ResultItem, len(results))
			for i := range len(rp) {
				rp[i] = &results[i]
			}
		}
		if len(rp) > 0 {
			if fss.query != "" {
				scores, err := fss.scorer.Score(utils.Map(func(r *ResultItem) string { return r.text }, rp), fss.query)
				if err != nil {
					return err
				}
				for i, r := range rp {
					r.SetScoreResult(scores[i])
					r.score.Set_length(uint16(len(r.text)))
				}
				rp = utils.Filter(rp, func(r *ResultItem) bool { return r.IsMatching() })
			} else {
				for _, r := range rp {
					r.score &= 0b11111111111111111111111111111111 // only preserve index
					r.positions = nil
				}
			}
		}
		min_score, max_score := CombinedScore(math.MaxUint64), CombinedScore(0)
		if len(rp) > 0 {
			slices.SortFunc(rp, func(a, b *ResultItem) int { return cmp.Compare(a.score, b.score) })
			min_score, max_score = rp[0].score, rp[len(rp)-1].score
		}
		var rr []*ResultItem
		fss.lock()
		existing := fss.renderable_results
		fss.unlock()
		switch {
		case min_score >= global_max_score:
			rr = append(existing, rp...)
		case max_score < global_min_score:
			rr = make([]*ResultItem, len(existing)+len(rp), max(16*1024, len(existing)+len(rp), 2*cap(existing)))
			copy(rr, rp)
			copy(rr[len(rp):], existing)
		default:
			rr = merge_sorted_slices(existing, rp)
		}
		global_min_score = min(global_min_score, min_score)
		global_max_score = max(global_max_score, max_score)
		fss.lock()
		fss.renderable_results = rr
		fss.unlock()
		return
	}

	offset := 0
	for range on_results {
		if !fss.keep_going.Load() {
			break
		}
		results := fss.scanner.Batch(offset)
		if len(results) > 0 || fss.scanner.Error() != nil {
			offset += len(results)
			fss.on_results(handle_batch(results), false)
		}
	}
	for fss.keep_going.Load() {
		b := fss.scanner.Batch(offset)
		if len(b) == 0 {
			break
		}
		offset += len(b)
		fss.on_results(handle_batch(b), false)
	}
}

func (fss *FileSystemScorer) Results() (ans ResultsType, is_finished bool) {
	fss.lock()
	defer fss.unlock()
	return fss.renderable_results, fss.is_complete.Load()
}

func (fss *FileSystemScorer) Cancel() {
	fss.keep_going.Store(false)
	fss.scanner.Cancel()
}

type Settings interface {
	OnlyDirs() bool
	CurrentDir() string
	SearchText() string
}

type ResultManager struct {
	report_errors    chan error
	WakeupMainThread func() bool
	settings         Settings

	scorer         *FileSystemScorer
	mutex          sync.Mutex
	last_wakeup_at time.Time
}

func NewResultManager(err_chan chan error, settings Settings, WakeupMainThread func() bool) *ResultManager {
	ans := &ResultManager{
		report_errors:    err_chan,
		settings:         settings,
		WakeupMainThread: WakeupMainThread,
	}
	return ans
}

func (m *ResultManager) on_results(err error, is_finished bool) {
	if err != nil {
		m.report_errors <- err
		m.WakeupMainThread()
		return
	}
	m.mutex.Lock()
	defer m.mutex.Unlock()
	if is_finished || time.Since(m.last_wakeup_at) > time.Millisecond*50 {
		m.WakeupMainThread()
		m.last_wakeup_at = time.Now()
	}
}

func merge_sorted_slices(a, b []*ResultItem) []*ResultItem {
	result := make([]*ResultItem, 0, 2*(len(a)+len(b)))
	i, j := 0, 0
	for i < len(a) && j < len(b) {
		if a[i].score <= b[j].score {
			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) 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
	}
	if m.scorer != nil {
		m.scorer.Cancel()
	}
	m.scorer = NewFileSystemScorer(root_dir, "", m.settings.OnlyDirs(), m.on_results)
	m.scorer.Start()
}

func (m *ResultManager) set_query(query string) {
	if m.scorer == nil {
		m.scorer = NewFileSystemScorer(".", "", m.settings.OnlyDirs(), m.on_results)
		m.scorer.Start()
	} else {
		m.scorer.Change_query(query)
	}
}

func (h *Handler) get_results() (ans ResultsType, is_complete bool) {
	if h.result_manager.scorer == nil {
		return
	}
	return h.result_manager.scorer.Results()
}