Reimplement SearchByPrefix to include distance in results

Author: andrewjsaid

src/Levenshtypo.Tests/LevenshtrieSearchByPrefixTests.cs

@@ -114,14 +114,51 @@ public void StackOverflow_Scenario2()
         Test(levenshtrie, "a", 0, entries);
     }
 
-    private static void Test(Levenshtrie<string> t, string query, int distance, IEnumerable<string> expectedResults)
+    private static void Test(Levenshtrie<string> t, string query, int distance, IEnumerable<string> expected)
     {
+        var list = expected.ToList();
+        if (list.Count > 10_000)
+        {
+            // Test will be too slow for Shouldly to compare with ignoreCase.
+            // N.B. to be clear: Levenshtypo can easily handle such numbers
+            return;
+        }
+
+        var expectedResults = list.Select(e => new LevenshtrieSearchResult<string>(CalculatePrefixDistance(query, e), e));
+
         t.SearchByPrefix(query, distance)
-            .OrderBy(x => x)
-            .ShouldBe(expectedResults.OrderBy(x => x));
+            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultEqualityComparer<string>());
 
         t.EnumerateSearchByPrefix(query, distance)
-            .OrderBy(x => x)
-            .ShouldBe(expectedResults.OrderBy(x => x));
+            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultEqualityComparer<string>());
+    }
+
+    private static int CalculatePrefixDistance(string query, string result)
+    {
+        int minDistance = int.MaxValue;
+
+        var automaton = LevenshtomatonFactory.Instance.Construct(query, maxEditDistance: 5);
+
+        var state = automaton.Start();
+
+        if (state.IsFinal && state.Distance < minDistance)
+        {
+            minDistance = state.Distance;
+        }
+
+        foreach (var r in result.EnumerateRunes())
+        {
+            if (!state.MoveNext(r, out state))
+            {
+                break;
+            }
+
+            if (state.IsFinal && state.Distance < minDistance)
+            {
+                minDistance = state.Distance;
+            }
+        }
+
+        return minDistance;
     }
 }

src/Levenshtypo.Tests/LevenshtrieSearchResultEqualityComparer.cs

@@ -0,0 +1,18 @@
+using System.Diagnostics.CodeAnalysis;
+
+namespace Levenshtypo.Tests;
+
+internal class LevenshtrieSearchResultEqualityComparer<T> : IEqualityComparer<LevenshtrieSearchResult<T>>
+{
+    public bool Equals(LevenshtrieSearchResult<T> x, LevenshtrieSearchResult<T> y)
+        => x.Distance == y.Distance
+        && (x.Result?.Equals(y.Result) ?? (y.Result is null));
+
+    public int GetHashCode([DisallowNull] LevenshtrieSearchResult<T> obj)
+    {
+        var hashCode = new HashCode();
+        hashCode.Add(obj.Distance);
+        hashCode.Add(obj.Result);
+        return hashCode.ToHashCode();
+    }
+}

src/Levenshtypo.Tests/LevenshtrieSearchTests.cs

@@ -1,5 +1,4 @@
-using System.Diagnostics.CodeAnalysis;
-using Shouldly;
+using Shouldly;
 
 namespace Levenshtypo.Tests;
 
@@ -112,6 +111,7 @@ public void MultiMap()
             .ShouldBe(["1", "2", "3", "5", "6", "7"], ignoreOrder: true, comparer: StringComparer.OrdinalIgnoreCase);
 
         t.SearchByPrefix("00", maxEditDistance: 0)
+            .Select(r => r.Result)
             .ToArray()
             .ShouldBe(["1", "2", "3", "4", "5", "6", "7"], ignoreOrder: true, comparer: StringComparer.OrdinalIgnoreCase);
 
@@ -127,24 +127,9 @@ private static void Test(Levenshtrie<string> t, string query, int distance, IEnu
         var expectedResults = expected.Select(e => new LevenshtrieSearchResult<string>(e.distance, e.word));
 
         t.Search(query, distance)
-            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultComparer<string>());
+            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultEqualityComparer<string>());
 
         t.EnumerateSearch(query, distance)
-            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultComparer<string>());
-    }
-
-    private class LevenshtrieSearchResultComparer<T> : IEqualityComparer<LevenshtrieSearchResult<T>>
-    {
-        public bool Equals(LevenshtrieSearchResult<T> x, LevenshtrieSearchResult<T> y)
-            => x.Distance == y.Distance
-            && (x.Result?.Equals(y.Result) ?? (y.Result is null));
-
-        public int GetHashCode([DisallowNull] LevenshtrieSearchResult<T> obj)
-        {
-            var hashCode = new HashCode();
-            hashCode.Add(obj.Distance);
-            hashCode.Add(obj.Result);
-            return hashCode.ToHashCode();
-        }
+            .ShouldBe(expectedResults, ignoreOrder: true, comparer: new LevenshtrieSearchResultEqualityComparer<string>());
     }
 }

src/Levenshtypo/ILevenshtrie.cs

@@ -25,16 +25,6 @@ public interface ILevenshtrie<T>
     LevenshtrieSearchResult<T>[] Search<TSearchState>(TSearchState searcher)
         where TSearchState : ILevenshtomatonExecutionState<TSearchState>;
 
-    /// <summary>
-    /// Searches for values whose keys begin with a prefix accepted by the search state,
-    /// but without considering further edits beyond the matched prefix.
-    /// </summary>
-    /// <typeparam name="TSearchState">The automaton state used to guide the prefix traversal.</typeparam>
-    /// <param name="searcher">The automaton execution state.</param>
-    /// <returns>An array of values with matching prefixes.</returns>
-    T[] SearchByPrefix<TSearchState>(TSearchState searcher)
-        where TSearchState : ILevenshtomatonExecutionState<TSearchState>;
-
     /// <summary>
     /// Lazily searches for values whose keys are accepted by the specified automaton execution state.
     /// Results are returned in arbitrary order and evaluated on-demand.
@@ -48,20 +38,6 @@ T[] SearchByPrefix<TSearchState>(TSearchState searcher)
     /// </remarks>
     IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearch<TSearchState>(TSearchState searcher)
         where TSearchState : ILevenshtomatonExecutionState<TSearchState>;
-    
-    /// <summary>
-    /// Lazily searches for values whose keys begin with a prefix accepted by the automaton execution state.
-    /// Results are returned in arbitrary order and evaluated on-demand.
-    /// </summary>
-    /// <typeparam name="TSearchState">The automaton state used to guide the traversal.</typeparam>
-    /// <param name="searcher">The automaton execution state.</param>
-    /// <returns>An enumerable of values with matching prefixes.</returns>
-    /// <remarks>
-    /// This method avoids allocating a full result array up front, and can be more efficient
-    /// when only a subset of results are consumed. However, it may be slower when all results are needed.
-    /// </remarks>
-    IEnumerable<T> EnumerateSearchByPrefix<TSearchState>(TSearchState searcher)
-        where TSearchState : ILevenshtomatonExecutionState<TSearchState>;
 }
 
 public static class LevenshtrieExtensions
@@ -113,6 +89,17 @@ public static LevenshtrieSearchResult<T>[] Search<T>(this ILevenshtrie<T> @this,
     public static LevenshtrieSearchResult<T>[] Search<T>(this ILevenshtrie<T> @this, LevenshtomatonExecutionState searcher)
         => @this.Search(searcher);
 
+    /// <summary>
+    /// Searches for values whose keys begin with a prefix accepted by the search state,
+    /// but without considering further edits beyond the matched prefix.
+    /// </summary>
+    /// <typeparam name="TSearchState">The automaton state used to guide the prefix traversal.</typeparam>
+    /// <param name="searcher">The automaton execution state.</param>
+    /// <returns>An array of matched values and their corresponding edit distances.</returns>
+    public static LevenshtrieSearchResult<T>[] SearchByPrefix<T, TSearchState>(this ILevenshtrie<T> @this, TSearchState searcher)
+        where TSearchState : ILevenshtomatonExecutionState<TSearchState>
+        => @this.Search(PrefixTrackingLevenshtomatonExecutionState<TSearchState>.Start(searcher));
+
     /// <summary>
     /// Searches for values whose keys begin with the specified prefix text,
     /// allowing up to the given edit distance. Results are returned in arbitrary order.
@@ -123,11 +110,12 @@ public static LevenshtrieSearchResult<T>[] Search<T>(this ILevenshtrie<T> @this,
     /// <param name="maxEditDistance">The maximum allowed edit distance.</param>
     /// <param name="metric">The edit distance metric to use.</param>
     /// <returns>An array of values whose keys begin with the given prefix.</returns>
-    public static T[] SearchByPrefix<T>(this ILevenshtrie<T> @this, string text, int maxEditDistance, LevenshtypoMetric metric = LevenshtypoMetric.Levenshtein)
+    public static LevenshtrieSearchResult<T>[] SearchByPrefix<T>(this ILevenshtrie<T> @this, string text, int maxEditDistance, LevenshtypoMetric metric = LevenshtypoMetric.Levenshtein)
     {
         var automaton = LevenshtomatonFactory.Instance.Construct(text, maxEditDistance, ignoreCase: @this.IgnoreCase, metric: metric);
         return @this.SearchByPrefix(automaton);
     }
+
     /// <summary>
     /// Searches for values whose keys are accepted by the given automaton, restricted to prefix matches.
     /// Results are returned in arbitrary order.
@@ -139,15 +127,15 @@ public static T[] SearchByPrefix<T>(this ILevenshtrie<T> @this, string text, int
     /// <exception cref="ArgumentException">
     /// Thrown if the automaton's case sensitivity does not match the trie's configuration.
     /// </exception>
-    public static T[] SearchByPrefix<T>(this ILevenshtrie<T> @this, Levenshtomaton automaton)
+    public static LevenshtrieSearchResult<T>[] SearchByPrefix<T>(this ILevenshtrie<T> @this, Levenshtomaton automaton)
     {
         if (automaton.IgnoreCase != @this.IgnoreCase)
         {
             throw new ArgumentException("Case sensitivity of automaton does not match.");
         }
 
-        var executor = @this as ILevenshtomatonExecutor<SearchByPrefixWrapper<T[]>> ?? new TrieExecutor<T>(@this);
-        return automaton.Execute(executor).Wrapped;
+        var executor = new TriePrefixExecutor<T>(@this);
+        return automaton.Execute<LevenshtrieSearchResult<T>[]>(executor);
     }
 
     /// <summary>
@@ -158,8 +146,8 @@ public static T[] SearchByPrefix<T>(this ILevenshtrie<T> @this, Levenshtomaton a
     /// <param name="this">The trie to search.</param>
     /// <param name="searcher">A boxed automaton execution state.</param>
     /// <returns>Values whose keys match the prefix condition.</returns>
-    public static T[] SearchByPrefix<T>(this ILevenshtrie<T> @this, LevenshtomatonExecutionState searcher)
-        => @this.SearchByPrefix<LevenshtomatonExecutionState>(searcher);
+    public static LevenshtrieSearchResult<T>[] SearchByPrefix<T>(this ILevenshtrie<T> @this, LevenshtomatonExecutionState searcher)
+        => @this.SearchByPrefix<T, LevenshtomatonExecutionState>(searcher);
 
     /// <summary>
     /// Lazily searches for approximate matches using a Levenshtein automaton, evaluated on-demand.
@@ -223,7 +211,7 @@ public static IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearch<T>(this IL
     /// Lazily evaluates prefix matches. Uses less memory than eager prefix search,
     /// and is efficient when only a subset of results is needed.
     /// </remarks>
-    public static IEnumerable<T> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, string text, int maxEditDistance, LevenshtypoMetric metric = LevenshtypoMetric.Levenshtein)
+    public static IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, string text, int maxEditDistance, LevenshtypoMetric metric = LevenshtypoMetric.Levenshtein)
     {
         var automaton = LevenshtomatonFactory.Instance.Construct(text, maxEditDistance, ignoreCase: @this.IgnoreCase, metric: metric);
         return @this.EnumerateSearchByPrefix(automaton);
@@ -239,15 +227,15 @@ public static IEnumerable<T> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @th
     /// <exception cref="ArgumentException">
     /// Thrown if the automaton's case sensitivity does not match the trie's configuration.
     /// </exception>
-    public static IEnumerable<T> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, Levenshtomaton automaton)
+    public static IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, Levenshtomaton automaton)
     {
         if (automaton.IgnoreCase != @this.IgnoreCase)
         {
             throw new ArgumentException("Case sensitivity of automaton does not match.");
         }
 
-        var executor = @this as ILevenshtomatonExecutor<SearchByPrefixWrapper<IEnumerable<T>>> ?? new TrieExecutor<T>(@this);
-        return automaton.Execute(executor).Wrapped;
+        var executor = new TriePrefixExecutor<T>(@this);
+        return automaton.Execute<IEnumerable<LevenshtrieSearchResult<T>>>(executor);
     }
 
     /// <summary>
@@ -257,25 +245,42 @@ public static IEnumerable<T> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @th
     /// <param name="this">The trie to search.</param>
     /// <param name="searcher">A boxed automaton execution state.</param>
     /// <returns>An enumerable of values with matching prefixes.</returns>
-    public static IEnumerable<T> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, LevenshtomatonExecutionState searcher)
-        => @this.EnumerateSearchByPrefix<LevenshtomatonExecutionState>(searcher);
+    public static IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearchByPrefix<T>(this ILevenshtrie<T> @this, LevenshtomatonExecutionState searcher)
+        => @this.EnumerateSearchByPrefix<T, LevenshtomatonExecutionState>(searcher);
+
+    /// <summary>
+    /// Lazily searches for values whose keys begin with a prefix accepted by the automaton execution state.
+    /// Results are returned in arbitrary order and evaluated on-demand.
+    /// </summary>
+    /// <typeparam name="TSearchState">The automaton state used to guide the traversal.</typeparam>
+    /// <param name="searcher">The automaton execution state.</param>
+    /// <returns>An array of matched values and their corresponding edit distances.</returns>
+    /// <remarks>
+    /// This method avoids allocating a full result array up front, and can be more efficient
+    /// when only a subset of results are consumed. However, it may be slower when all results are needed.
+    /// </remarks>
+    public static IEnumerable<LevenshtrieSearchResult<T>> EnumerateSearchByPrefix<T, TSearchState>(this ILevenshtrie<T> @this, TSearchState searcher)
+        where TSearchState : ILevenshtomatonExecutionState<TSearchState>
+        => @this.EnumerateSearch(PrefixTrackingLevenshtomatonExecutionState<TSearchState>.Start(searcher));
 
     private class TrieExecutor<T>(ILevenshtrie<T> trie) :
         ILevenshtomatonExecutor<LevenshtrieSearchResult<T>[]>,
-        ILevenshtomatonExecutor<IEnumerable<LevenshtrieSearchResult<T>>>,
-        ILevenshtomatonExecutor<SearchByPrefixWrapper<T[]>>,
-        ILevenshtomatonExecutor<SearchByPrefixWrapper<IEnumerable<T>>>
+        ILevenshtomatonExecutor<IEnumerable<LevenshtrieSearchResult<T>>>
     {
         public LevenshtrieSearchResult<T>[] ExecuteAutomaton<TState>(TState executionState) where TState : struct, ILevenshtomatonExecutionState<TState>
             => trie.Search(executionState);
 
         IEnumerable<LevenshtrieSearchResult<T>> ILevenshtomatonExecutor<IEnumerable<LevenshtrieSearchResult<T>>>.ExecuteAutomaton<TState>(TState executionState)
             => trie.EnumerateSearch(executionState);
+    }
+    private class TriePrefixExecutor<T>(ILevenshtrie<T> trie) :
+        ILevenshtomatonExecutor<LevenshtrieSearchResult<T>[]>,
+        ILevenshtomatonExecutor<IEnumerable<LevenshtrieSearchResult<T>>>
+    {
+        public LevenshtrieSearchResult<T>[] ExecuteAutomaton<TState>(TState executionState) where TState : struct, ILevenshtomatonExecutionState<TState>
+            => trie.SearchByPrefix(executionState);
 
-        SearchByPrefixWrapper<T[]> ILevenshtomatonExecutor<SearchByPrefixWrapper<T[]>>.ExecuteAutomaton<TState>(TState executionState)
-            => new(trie.SearchByPrefix(executionState));
-
-        SearchByPrefixWrapper<IEnumerable<T>> ILevenshtomatonExecutor<SearchByPrefixWrapper<IEnumerable<T>>>.ExecuteAutomaton<TState>(TState executionState)
-            => new(trie.EnumerateSearchByPrefix(executionState));
+        IEnumerable<LevenshtrieSearchResult<T>> ILevenshtomatonExecutor<IEnumerable<LevenshtrieSearchResult<T>>>.ExecuteAutomaton<TState>(TState executionState)
+            => trie.EnumerateSearchByPrefix(executionState);
     }
 }