FastSeekWpf/Core/SearchEngine.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;

namespace FastSeekWpf.Core;

public enum ResultKind
{
    App, Document, Image, Video, Audio, Archive, Folder, Other
}

public class SearchResult
{
    public string FullPath { get; set; } = string.Empty;
    public string Name { get; set; } = string.Empty;
    public byte Rank { get; set; }
    public bool IsDir { get; set; }
    public ResultKind Kind { get; set; }
}

public static class SearchEngine
{
    [ThreadStatic]
    private static StringBuilder? _pathBuilder;

    private static StringBuilder GetPathBuilder()
    {
        var sb = _pathBuilder;
        if (sb == null)
        {
            sb = new StringBuilder(512);
            _pathBuilder = sb;
        }
        sb.Clear();
        return sb;
    }

    private static readonly string[] AppExtensions = { "exe", "lnk", "msi", "appx", "msix" };
    private static readonly string[] AppPathMarkers = {
        "\\program files\\", "\\program files (x86)\\",
        "\\start menu\\", "\\desktop\\", "\\appdata\\"
    };

    private static readonly Dictionary<string, ResultKind> ExtensionMap = new(StringComparer.OrdinalIgnoreCase)
    {
        ["exe"] = ResultKind.App, ["lnk"] = ResultKind.App, ["msi"] = ResultKind.App,
        ["doc"] = ResultKind.Document, ["docx"] = ResultKind.Document, ["pdf"] = ResultKind.Document,
        ["txt"] = ResultKind.Document, ["xlsx"] = ResultKind.Document, ["xls"] = ResultKind.Document,
        ["pptx"] = ResultKind.Document, ["ppt"] = ResultKind.Document, ["odt"] = ResultKind.Document,
        ["ods"] = ResultKind.Document, ["odp"] = ResultKind.Document, ["rtf"] = ResultKind.Document,
        ["md"] = ResultKind.Document, ["csv"] = ResultKind.Document, ["json"] = ResultKind.Document,
        ["xml"] = ResultKind.Document, ["yaml"] = ResultKind.Document, ["toml"] = ResultKind.Document,
        ["ini"] = ResultKind.Document, ["log"] = ResultKind.Document,
        ["png"] = ResultKind.Image, ["jpg"] = ResultKind.Image, ["jpeg"] = ResultKind.Image,
        ["gif"] = ResultKind.Image, ["bmp"] = ResultKind.Image, ["webp"] = ResultKind.Image,
        ["svg"] = ResultKind.Image, ["ico"] = ResultKind.Image, ["tiff"] = ResultKind.Image,
        ["heic"] = ResultKind.Image, ["raw"] = ResultKind.Image, ["psd"] = ResultKind.Image,
        ["mp4"] = ResultKind.Video, ["mkv"] = ResultKind.Video, ["avi"] = ResultKind.Video,
        ["mov"] = ResultKind.Video, ["wmv"] = ResultKind.Video, ["flv"] = ResultKind.Video,
        ["webm"] = ResultKind.Video, ["m4v"] = ResultKind.Video,
        ["mp3"] = ResultKind.Audio, ["flac"] = ResultKind.Audio, ["wav"] = ResultKind.Audio,
        ["aac"] = ResultKind.Audio, ["ogg"] = ResultKind.Audio, ["m4a"] = ResultKind.Audio,
        ["wma"] = ResultKind.Audio, ["opus"] = ResultKind.Audio,
        ["zip"] = ResultKind.Archive, ["rar"] = ResultKind.Archive, ["7z"] = ResultKind.Archive,
        ["tar"] = ResultKind.Archive, ["gz"] = ResultKind.Archive, ["bz2"] = ResultKind.Archive,
        ["xz"] = ResultKind.Archive, ["zst"] = ResultKind.Archive,
    };

    public static ResultKind GetKind(string? ext, bool isDir, string fullPathLower)
    {
        if (isDir) return ResultKind.Folder;
        if (ext == null) return ResultKind.Other;
        if (ExtensionMap.TryGetValue(ext, out var kind))
        {
            if (kind == ResultKind.App && ContainsAny(fullPathLower, AppPathMarkers))
                return ResultKind.App;
            return kind;
        }
        return ResultKind.Other;
    }

    private static bool ContainsAny(string s, string[] markers)
    {
        foreach (var m in markers)
            if (s.Contains(m, StringComparison.Ordinal)) return true;
        return false;
    }

    private static bool SpanContains(ReadOnlySpan<char> haystack, ReadOnlySpan<char> needle)
    {
        if (needle.IsEmpty) return true;
        if (needle.Length > haystack.Length) return false;

        int limit = haystack.Length - needle.Length + 1;
        for (int i = 0; i < limit; i++)
        {
            if (haystack.Slice(i, needle.Length).SequenceEqual(needle))
                return true;
        }
        return false;
    }

    public static List<SearchResult> Search(
        IndexStore store,
        string query,
        int limit,
        bool caseSensitive,
        List<string> excludedDirs)
    {
        if (string.IsNullOrWhiteSpace(query))
            return new List<SearchResult>();

        var searchSw = System.Diagnostics.Stopwatch.StartNew();
        string q = caseSensitive ? query : query.ToLowerInvariant();
        ReadOnlySpan<char> qSpan = q.AsSpan();

        Span<(int idx, byte rank)> candidates = stackalloc (int, byte)[10000];
        int candidateCount = 0;

        lock (store)
        {
            int count = store.Entries.Count;
            for (int i = 0; i < count; i++)
            {
                string name = caseSensitive ? store.NameAt(i) : store.NameLowerAt(i);
                ReadOnlySpan<char> nameSpan = name.AsSpan();

                byte rank;
                if (nameSpan.SequenceEqual(qSpan)) rank = 1;
                else if (nameSpan.StartsWith(qSpan, StringComparison.Ordinal)) rank = 2;
                else if (SpanContains(nameSpan, qSpan)) rank = 3;
                else continue;

                if (candidateCount < candidates.Length)
                    candidates[candidateCount++] = (i, rank);
            }
        }

        candidates[..candidateCount].Sort((a, b) => a.rank.CompareTo(b.rank));

        int overshoot = Math.Max(limit * 5, 1000);
        if (candidateCount > overshoot)
            candidateCount = overshoot;

        var results = new List<SearchResult>(Math.Min(limit, candidateCount));
        var pathBuilder = GetPathBuilder();

        foreach (var (idx, baseRank) in candidates[..candidateCount])
        {
            var entry = store.Entries[idx];
            string fullPath = BuildPathFast(store, entry.FileRef, pathBuilder, entry.DriveIdx);
            string fullPathLower = fullPath.ToLowerInvariant();

            if (excludedDirs.Count > 0)
            {
                bool excluded = false;
                foreach (var ex in excludedDirs)
                {
                    if (fullPathLower.StartsWith(ex, StringComparison.Ordinal))
                    { excluded = true; break; }
                }
                if (excluded) continue;
            }

            string name = store.NameAt(idx);
            string nameLower = store.NameLowerAt(idx);

            byte rank = baseRank;
            if (baseRank <= 2)
            {
                int dotIdx = nameLower.LastIndexOf('.');
                if (dotIdx >= 0)
                {
                    string ext = nameLower[(dotIdx + 1)..];
                    if (Array.IndexOf(AppExtensions, ext) >= 0 && ContainsAny(fullPathLower, AppPathMarkers))
                        rank = 0;
                }
            }

            string? fileExt = null;
            if (!entry.IsDir)
            {
                int dot = name.LastIndexOf('.');
                if (dot >= 0) fileExt = name[(dot + 1)..].ToLowerInvariant();
            }

            results.Add(new SearchResult
            {
                FullPath = fullPath,
                Name = name,
                Rank = rank,
                IsDir = entry.IsDir,
                Kind = GetKind(fileExt, entry.IsDir, fullPathLower)
            });
        }

        results.Sort((a, b) => a.Rank.CompareTo(b.Rank));
        if (results.Count > limit)
            results.RemoveRange(limit, results.Count - limit);

        searchSw.Stop();
        Logger.Log($"Search '{query}': {candidateCount} candidates, {results.Count} results, {searchSw.ElapsedMilliseconds}ms (index={store.Count:N0})");

        return results;
    }

    public static List<SearchResult> SearchByExtension(
        IndexStore store,
        string ext,
        List<string> excludedDirs)
    {
        string dotExt = "." + ext.ToLowerInvariant();
        var results = new List<SearchResult>();
        var pathBuilder = GetPathBuilder();

        lock (store)
        {
            for (int i = 0; i < store.Entries.Count; i++)
            {
                string name = store.NameLowerAt(i);
                if (!name.EndsWith(dotExt)) continue;

                var entry = store.Entries[i];
                string fullPath = BuildPathFast(store, entry.FileRef, pathBuilder, entry.DriveIdx);
                string fullPathLower = fullPath.ToLowerInvariant();

                bool excluded = false;
                foreach (var ex in excludedDirs)
                {
                    if (fullPathLower.StartsWith(ex, StringComparison.Ordinal))
                    { excluded = true; break; }
                }
                if (excluded) continue;

                string? fileExt = null;
                int dot = name.LastIndexOf('.');
                if (dot >= 0) fileExt = name[(dot + 1)..];

                results.Add(new SearchResult
                {
                    FullPath = fullPath,
                    Name = store.NameAt(i),
                    Rank = 0,
                    IsDir = entry.IsDir,
                    Kind = GetKind(fileExt, entry.IsDir, fullPathLower)
                });

                if (results.Count >= 50) break;
            }
        }

        return results;
    }

    private static string BuildPathFast(IndexStore store, ulong fileRef, StringBuilder sb, byte startDriveIdx)
    {
        sb.Clear();
        string[] components = new string[64];
        int compCount = 0;
        ulong current = fileRef;
        byte driveIdx = startDriveIdx;

        for (int i = 0; i < 64; i++)
        {
            var idx = store.LookupIdx(current);
            if (idx == null) break;

            int eIdx = (int)idx;
            components[compCount++] = store.NameAt(eIdx);
            var entry = store.Entries[eIdx];
            driveIdx = entry.DriveIdx;
            if (entry.ParentRef == current) break;
            current = entry.ParentRef;
        }

        string driveRoot = driveIdx < store.DriveRoots.Count ? store.DriveRoots[driveIdx] : "C:\\";
        sb.Append(driveRoot);
        for (int i = compCount - 1; i >= 0; i--)
        {
            if (sb.Length > 0 && sb[sb.Length - 1] != '\\' && sb[sb.Length - 1] != '/')
                sb.Append('\\');
            sb.Append(components[i]);
        }
        return sb.ToString();
    }
}