Optimise FallbackIfEmpty (#464)

Co-authored-by: Stuart Turner <stuart@turner-isageek.com>
viceroypenguin · Jun 18, 2023 · 89ad40e · 89ad40e
1 parent 56f77c0
commit 89ad40e
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 8 deletions.
diff --git a/Source/SuperLinq/FallbackIfEmpty.cs b/Source/SuperLinq/FallbackIfEmpty.cs
@@ -40,16 +40,22 @@ public static IEnumerable<T> FallbackIfEmpty<T>(this IEnumerable<T> source, IEnu
 		Guard.IsNotNull(source);
 		Guard.IsNotNull(fallback);
 
-		return Core(source, fallback);
+		return source.TryGetCollectionCount() is not null
+			   && fallback.TryGetCollectionCount() is not null
+			? new FallbackIfEmptyCollectionIterator<T>(source, fallback)
+			: Core(source, fallback);
 
 		static IEnumerable<T> Core(IEnumerable<T> source, IEnumerable<T> fallback)
 		{
 			using (var e = source.GetEnumerator())
 			{
 				if (e.MoveNext())
 				{
-					do { yield return e.Current; }
-					while (e.MoveNext());
+					do
+					{
+						yield return e.Current;
+					} while (e.MoveNext());
+
 					yield break;
 				}
 			}
@@ -58,4 +64,28 @@ static IEnumerable<T> Core(IEnumerable<T> source, IEnumerable<T> fallback)
 				yield return item;
 		}
 	}
+
+	private sealed class FallbackIfEmptyCollectionIterator<T> : CollectionIterator<T>
+	{
+		private readonly IEnumerable<T> _source;
+		private readonly IEnumerable<T> _fallback;
+
+		public FallbackIfEmptyCollectionIterator(IEnumerable<T> source, IEnumerable<T> fallback)
+		{
+			_source = source;
+			_fallback = fallback;
+		}
+
+		public override int Count =>
+			_source.GetCollectionCount() == 0
+				? _fallback.Count()
+				: _source.GetCollectionCount();
+
+		protected override IEnumerable<T> GetEnumerable()
+		{
+			return _source.GetCollectionCount() == 0
+				? _fallback
+				: _source;
+		}
+	}
 }
diff --git a/Tests/SuperLinq.Test/FallbackIfEmptyTest.cs b/Tests/SuperLinq.Test/FallbackIfEmptyTest.cs
@@ -5,16 +5,60 @@ public class FallbackIfEmptyTest
 	[Fact]
 	public void FallbackIfEmptyWithEmptySequence()
 	{
-		using var source = Enumerable.Empty<int>().AsTestingSequence(maxEnumerations: 2);
+		using var source = Seq<int>().AsTestingSequence();
+		source.FallbackIfEmpty(12).AssertSequenceEqual(12);
+	}
+
+	[Fact]
+	public void FallbackIfEmptyWithCollectionSequence()
+	{
+		using var source = Seq<int>().AsTestingCollection();
 		source.FallbackIfEmpty(12).AssertSequenceEqual(12);
-		source.FallbackIfEmpty(12, 23).AssertSequenceEqual(12, 23);
 	}
 
 	[Fact]
 	public void FallbackIfEmptyWithNotEmptySequence()
 	{
-		using var source = Seq(1).AsTestingSequence(maxEnumerations: 2);
-		source.FallbackIfEmpty(12).AssertSequenceEqual(1);
-		source.FallbackIfEmpty(12, 23).AssertSequenceEqual(1);
+		using var source = Seq(1).AsTestingSequence();
+		source.FallbackIfEmpty(new BreakingSequence<int>()).AssertSequenceEqual(1);
+	}
+
+	[Fact]
+	public void FallbackIfEmptyWithNotEmptyCollectionSequence()
+	{
+		using var source = Seq(1).AsTestingCollection();
+		source.FallbackIfEmpty(new BreakingSequence<int>()).AssertSequenceEqual(1);
+	}
+
+	[Fact]
+	public void AssertFallbackIfEmptyCollectionBehaviorOnEmptyCollection()
+	{
+		using var source = Seq<int>().AsBreakingCollection();
+		using var fallback = Enumerable.Range(0, 10_000).AsBreakingCollection();
+
+		var result = source.FallbackIfEmpty(fallback);
+		result.AssertCollectionErrorChecking(10_000);
+	}
+
+	[Fact]
+	public void AssertFallbackIfEmptyCollectionBehaviorOnNonEmptyCollection()
+	{
+		using var source = Enumerable.Range(0, 10_000).AsBreakingCollection();
+		using var fallback = new BreakingCollection<int>();
+
+		var result = source.FallbackIfEmpty(fallback);
+		result.AssertCollectionErrorChecking(10_000);
+	}
+
+	[Fact]
+	public void FallbackIfEmptyUsesCollectionCountAtIterationTime()
+	{
+		var stack = new Stack<int>();
+
+		var result = stack.FallbackIfEmpty(4);
+		result.AssertSequenceEqual(4);
+
+		stack.Push(1);
+		result.AssertSequenceEqual(1);
 	}
 }