whitemagic/WhiteMagic/Discovery/PatternScannerCache.cs

using System.Collections.Concurrent;
using System.Diagnostics;

namespace WhiteMagic.Discovery;

/// <summary>
/// Caches pattern scan results to avoid repeated scans of the same memory range.
/// </summary>
public sealed class PatternScannerCache
{
    private readonly ConcurrentDictionary<CacheKey, IntPtr> _cache = new();
    private readonly MemoryBase _memory;

    /// <summary>
    /// Creates a new cache for the given memory accessor.
    /// </summary>
    /// <param name="memory">The memory accessor to scan.</param>
    public PatternScannerCache(MemoryBase memory)
    {
        ArgumentNullException.ThrowIfNull(memory);
        _memory = memory;
    }

    /// <summary>
    /// Finds a pattern, returning a cached result if available.
    /// </summary>
    /// <param name="pattern">The byte pattern to search for.</param>
    /// <param name="mask">
    /// A mask string where 'x' means "match this byte exactly" and '?' means "wildcard".
    /// If <see langword="null"/>, all bytes are treated as 'x' (exact match).
    /// </param>
    /// <param name="start">The starting address of the scan range.</param>
    /// <param name="end">The ending address (exclusive) of the scan range.</param>
    /// <returns>
    /// The address of the first match from cache or memory, or <see cref="IntPtr.Zero"/> if not found.
    /// </returns>
    public IntPtr FindCached(
        byte[] pattern,
        string? mask,
        IntPtr start,
        IntPtr end)
    {
        var key = new CacheKey(pattern, mask, start, end);

        // Try to get from cache first
        if (_cache.TryGetValue(key, out IntPtr cached))
            return cached;

        // Not in cache, perform the scan
        IntPtr found = PatternScanner.Find(_memory, pattern, mask, start, end);

        // Cache the result (even if Zero)
        _cache[key] = found;

        return found;
    }

    /// <summary>
    /// Finds a pattern within a module, returning a cached result if available.
    /// </summary>
    /// <param name="pattern">The byte pattern to search for.</param>
    /// <param name="mask">
    /// A mask string where 'x' means "match this byte exactly" and '?' means "wildcard".
    /// If <see langword="null"/>, all bytes are treated as 'x' (exact match).
    /// </param>
    /// <param name="module">The module to scan.</param>
    /// <returns>
    /// The address of the first match from cache or memory, or <see cref="IntPtr.Zero"/> if not found.
    /// </returns>
    public IntPtr FindInModuleCached(
        byte[] pattern,
        string? mask,
        ProcessModule module)
    {
        ArgumentNullException.ThrowIfNull(module);

        IntPtr start = module.BaseAddress;
        IntPtr end = start + module.ModuleMemorySize;

        return FindCached(pattern, mask, start, end);
    }

    /// <summary>
    /// Finds a pattern across multiple modules, returning a cached result if available.
    /// </summary>
    /// <param name="pattern">The byte pattern to search for.</param>
    /// <param name="mask">
    /// A mask string where 'x' means "match this byte exactly" and '?' means "wildcard".
    /// If <see langword="null"/>, all bytes are treated as 'x' (exact match).
    /// </param>
    /// <param name="modules">The modules to scan, in order.</param>
    /// <returns>
    /// The address of the first match from cache or memory, or <see cref="IntPtr.Zero"/> if not found.
    /// </returns>
    public IntPtr FindInModulesCached(
        byte[] pattern,
        string? mask,
        IEnumerable<ProcessModule> modules)
    {
        // For multiple modules, we use a combined key (all modules hashed together)
        // This is less granular but still useful for repeated queries
        var moduleList = modules.ToList();
        var key = new CacheKey(pattern, mask, IntPtr.Zero, IntPtr.Zero, Modules: moduleList);
        if (_cache.TryGetValue(key, out IntPtr cached))
            return cached;

        IntPtr found = PatternScanner.FindInModules(_memory, pattern, mask, moduleList);

        _cache[key] = found;
        return found;
    }

    /// <summary>
    /// Clears all cached scan results.
    /// </summary>
    public void Clear()
    {
        _cache.Clear();
    }

    /// <summary>
    /// Cache key combining pattern, mask, and address range.
    /// </summary>
    private sealed record CacheKey(
        byte[] Pattern,
        string? Mask,
        IntPtr Start,
        IntPtr End,
        IReadOnlyList<ProcessModule>? Modules = null) : IEquatable<CacheKey>
    {
        public bool Equals(CacheKey? other)
        {
            if (other is null)
                return false;
            if (Start != other.Start || End != other.End || Mask != other.Mask)
                return false;
            if (!Pattern.AsSpan().SequenceEqual(other.Pattern))
                return false;

            if (Modules is null)
                return other.Modules is null;
            if (other.Modules is null || Modules.Count != other.Modules.Count)
                return false;

            for (int i = 0; i < Modules.Count; i++)
            {
                if (Modules[i].BaseAddress != other.Modules[i].BaseAddress)
                    return false;
            }

            return true;
        }

        // Override GetHashCode to hash the contents, not references
        public override int GetHashCode()
        {
            var hash = new HashCode();

            // Hash pattern bytes
            foreach (byte b in Pattern)
                hash.Add(b);

            // Hash mask
            hash.Add(Mask?.GetHashCode() ?? 0);

            // Hash address range
            hash.Add(Start.GetHashCode());
            hash.Add(End.GetHashCode());

            // Hash modules if present (by base address)
            if (Modules is not null)
            {
                foreach (var m in Modules)
                    hash.Add(m.BaseAddress.GetHashCode());
            }

            return hash.ToHashCode();
        }
    }
}