Custom Inclusive TakeWhile(), is there a better way?

Question

I've written a custom LINQ extension method that extends the TakeWhile() method to be inclusive, rather than exclusive when the predicate is false.

        public static IEnumerable<T> TakeWhile<T>(this IEnumerable<T> source, Func<T, bool> predicate, bool inclusive)
        {
            source.ThrowIfNull("source");
            predicate.ThrowIfNull("predicate");

            if (!inclusive)
                return source.TakeWhile(predicate);

            var totalCount = source.Count();
            var count = source.TakeWhile(predicate).Count();

            if (count == totalCount)
                return source;
            else
                return source.Take(count + 1);
        }

While this works, I'm sure there is a better way of doing it. I'm fairly sure that this doesn't work in terms of deferred execution/loading.

ThrowIfNull() is a extension method for ArgumentNullException checking

Can the community provide some some hints or re-writes? :)

Answer 1

I'm not sure if there's an easy way to do this using existing LINQ primitives; what you're actually wanting to do is store state as part of a LINQ expression, and that's always difficult since LINQ is a functional system.

You could rewrite it so that it doesn't traverse the source sequence multiple times:

 private static IEnumerable<T> TakeWhileInclusive<T>(this IEnumerable<T> source, Func<T, bool> predicate)
 {
     using (var enumerator = source.GetEnumerator())
     {
         bool last = false;
         while (enumerator.MoveNext())
         {
             if (last)
             {
                 yield return enumerator.Current;
                 yield break;
             }
             else if (predicate(enumerator.Current))
             {
                 yield return enumerator.Current;
             }
             else
             {
                 last = true;
             }
         }
     }
 }

Answer 2

You are correct; this is not friendly to deferred execution (calling Count requires a full enumeration of the source).

You could, however, do this:

public static IEnumerable<T> TakeWhile<T>(this IEnumerable<T> source, Func<T, bool> predicate, bool inclusive)
{
    foreach(T item in source)
    {
        if(predicate(item)) 
        {
            yield return item;
        }
        else
        {
            if(inclusive) yield return item;

            yield break;
        }
    }
}