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

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;
        }
    }
} 

Answer 2

Тot a direct answer to the question. I want to show how you can use SkipWhile to simulate inclusive TakeWhile logic.

IEnumerable<string> list = new List<string> { "1", "2", "3", "4", "5" };
var result = list
    .Reverse()
    .SkipWhile(item => item != "3")
    .Reverse();
// result will be {"1", "2", "3"}

However note that doing it this way will be less efficient, because essentially, list in this case has to to be reversed twice. So make sure performance will not be an issue.