Optimizing OrderBy for Skip/Take/etc.

Ix.NET/Source/System.Linq.Async/System/Linq/Operators/OrderedAsyncEnumerable.cs

@@ -10,15 +10,23 @@ using System.Threading.Tasks;
 
 namespace System.Linq
 {
-    // TODO: Add optimizations for First, Last, and ElementAt.
-    
-    internal abstract class OrderedAsyncEnumerable<TElement> : AsyncIterator<TElement>, IOrderedAsyncEnumerable<TElement>, IAsyncIListProvider<TElement>
+    // NB: Large portions of the implementation are based on https://github.com/dotnet/corefx/blob/master/src/System.Linq/src/System/Linq/OrderedEnumerable.cs.
+
+    internal abstract class OrderedAsyncEnumerable<TElement> : AsyncIterator<TElement>, IOrderedAsyncEnumerable<TElement>, IAsyncPartition<TElement>
     {
-        protected IAsyncEnumerable<TElement> _source;
+        protected readonly IAsyncEnumerable<TElement> _source;
+
         private TElement[] _buffer;
         private int[] _indexes;
         private int _index;
 
+        protected OrderedAsyncEnumerable(IAsyncEnumerable<TElement> source)
+        {
+            Debug.Assert(source != null);
+
+            _source = source;
+        }
+
         IOrderedAsyncEnumerable<TElement> IOrderedAsyncEnumerable<TElement>.CreateOrderedEnumerable<TKey>(Func<TElement, TKey> keySelector, IComparer<TKey> comparer, bool descending)
         {
             return new OrderedAsyncEnumerable<TElement, TKey>(_source, keySelector, comparer, descending, this);
@@ -39,9 +47,9 @@ namespace System.Linq
             switch (_state)
             {
                 case AsyncIteratorState.Allocated:
-                    _buffer = await _source.ToArrayAsync(_cancellationToken).ConfigureAwait(false);
+                    _buffer = await _source.ToArrayAsync(_cancellationToken).ConfigureAwait(false); // TODO: Use buffer.
 
-                    AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(next: null, _cancellationToken);
+                    AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(_cancellationToken);
                     _indexes = await sorter.Sort(_buffer, _buffer.Length).ConfigureAwait(false);
                     _index = 0;
 
@@ -72,6 +80,8 @@ namespace System.Linq
 
         internal abstract AsyncEnumerableSorter<TElement> GetAsyncEnumerableSorter(AsyncEnumerableSorter<TElement> next, CancellationToken cancellationToken);
 
+        internal AsyncEnumerableSorter<TElement> GetAsyncEnumerableSorter(CancellationToken cancellationToken) => GetAsyncEnumerableSorter(next: null, cancellationToken);
+
         public async ValueTask<TElement[]> ToArrayAsync(CancellationToken cancellationToken)
         {
             AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
@@ -101,6 +111,49 @@ namespace System.Linq
             return result;
         }
 
+        internal async ValueTask<TElement[]> ToArrayAsync(int minIndexInclusive, int maxIndexInclusive, CancellationToken cancellationToken)
+        {
+            AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
+
+            int count = elements.Length;
+
+            if (count <= minIndexInclusive)
+            {
+#if NO_ARRAY_EMPTY
+                return EmptyArray<TElement>.Value;
+#else
+                return Array.Empty<TElement>();
+#endif
+            }
+
+            if (count <= maxIndexInclusive)
+            {
+                maxIndexInclusive = count - 1;
+            }
+
+            TElement[] array = elements.Array;
+
+            if (minIndexInclusive == maxIndexInclusive)
+            {
+                AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
+
+                TElement element = await sorter.ElementAt(array, count, minIndexInclusive).ConfigureAwait(false);
+
+                return new TElement[] { element };
+            }
+
+            int[] map = await SortedMap(array, count, minIndexInclusive, maxIndexInclusive, cancellationToken).ConfigureAwait(false);
+
+            var result = new TElement[maxIndexInclusive - minIndexInclusive + 1];
+
+            for (int i = 0; minIndexInclusive <= maxIndexInclusive; i++)
+            {
+                result[i] = array[map[minIndexInclusive++]];
+            }
+
+            return result;
+        }
+
         public async ValueTask<List<TElement>> ToListAsync(CancellationToken cancellationToken)
         {
             AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
@@ -126,6 +179,45 @@ namespace System.Linq
             return result;
         }
 
+        internal async ValueTask<List<TElement>> ToListAsync(int minIndexInclusive, int maxIndexInclusive, CancellationToken cancellationToken)
+        {
+            AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
+
+            int count = elements.Length;
+
+            if (count <= minIndexInclusive)
+            {
+                return new List<TElement>(0);
+            }
+
+            if (count <= maxIndexInclusive)
+            {
+                maxIndexInclusive = count - 1;
+            }
+
+            TElement[] array = elements.Array;
+
+            if (minIndexInclusive == maxIndexInclusive)
+            {
+                AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
+
+                TElement element = await sorter.ElementAt(array, count, minIndexInclusive).ConfigureAwait(false);
+
+                return new List<TElement>(1) { element };
+            }
+
+            int[] map = await SortedMap(array, count, minIndexInclusive, maxIndexInclusive, cancellationToken).ConfigureAwait(false);
+
+            var list = new List<TElement>(maxIndexInclusive - minIndexInclusive + 1);
+
+            while (minIndexInclusive <= maxIndexInclusive)
+            {
+                list.Add(array[map[minIndexInclusive++]]);
+            }
+
+            return list;
+        }
+
         public async ValueTask<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken)
         {
             if (_source is IAsyncIListProvider<TElement> listProv)
@@ -136,12 +228,197 @@ namespace System.Linq
             return !onlyIfCheap || _source is ICollection<TElement> || _source is ICollection ? await _source.CountAsync(cancellationToken).ConfigureAwait(false) : -1;
         }
 
+        internal async ValueTask<int> GetCountAsync(int minIndexInclusive, int maxIndexInclusive, bool onlyIfCheap, CancellationToken cancellationToken)
+        {
+            int count = await GetCountAsync(onlyIfCheap, cancellationToken).ConfigureAwait(false);
+
+            if (count <= 0)
+            {
+                return count;
+            }
+
+            if (count <= minIndexInclusive)
+            {
+                return 0;
+            }
+
+            return (count <= maxIndexInclusive ? count : maxIndexInclusive + 1) - minIndexInclusive;
+        }
+
         private ValueTask<int[]> SortedMap(TElement[] elements, int count, CancellationToken cancellationToken)
         {
-            AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(next: null, cancellationToken);
+            AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
 
             return sorter.Sort(elements, count);
         }
+
+        private ValueTask<int[]> SortedMap(TElement[] elements, int count, int minIndexInclusive, int maxIndexInclusive, CancellationToken cancellationToken)
+        {
+            AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
+
+            return sorter.Sort(elements, count, minIndexInclusive, maxIndexInclusive);
+        }
+
+        private AsyncCachingComparer<TElement> GetComparer() => GetComparer(childComparer: null);
+
+        internal abstract AsyncCachingComparer<TElement> GetComparer(AsyncCachingComparer<TElement> childComparer);
+
+        public async ValueTask<Maybe<TElement>> TryGetFirstAsync(CancellationToken cancellationToken)
+        {
+            IAsyncEnumerator<TElement> e = _source.GetAsyncEnumerator(cancellationToken);
+
+            try
+            {
+                if (!await e.MoveNextAsync().ConfigureAwait(false))
+                {
+                    return new Maybe<TElement>();
+                }
+
+                TElement value = e.Current;
+
+                AsyncCachingComparer<TElement> comparer = GetComparer();
+
+                await comparer.SetElement(value, cancellationToken).ConfigureAwait(false);
+
+                while (await e.MoveNextAsync().ConfigureAwait(false))
+                {
+                    TElement x = e.Current;
+
+                    if (await comparer.Compare(x, cacheLower: true, cancellationToken).ConfigureAwait(false) < 0)
+                    {
+                        value = x;
+                    }
+                }
+
+                return new Maybe<TElement>(value);
+            }
+            finally
+            {
+                await e.DisposeAsync().ConfigureAwait(false);
+            }
+        }
+
+        public async ValueTask<Maybe<TElement>> TryGetLastAsync(CancellationToken cancellationToken)
+        {
+            IAsyncEnumerator<TElement> e = _source.GetAsyncEnumerator(cancellationToken);
+
+            try
+            {
+                if (!await e.MoveNextAsync().ConfigureAwait(false))
+                {
+                    return new Maybe<TElement>();
+                }
+
+                TElement value = e.Current;
+
+                AsyncCachingComparer<TElement> comparer = GetComparer();
+
+                await comparer.SetElement(value, cancellationToken).ConfigureAwait(false);
+
+                while (await e.MoveNextAsync().ConfigureAwait(false))
+                {
+                    TElement current = e.Current;
+
+                    if (await comparer.Compare(current, cacheLower: false, cancellationToken).ConfigureAwait(false) >= 0)
+                    {
+                        value = current;
+                    }
+                }
+
+                return new Maybe<TElement>(value);
+            }
+            finally
+            {
+                await e.DisposeAsync().ConfigureAwait(false);
+            }
+        }
+
+        internal async ValueTask<Maybe<TElement>> TryGetLastAsync(int minIndexInclusive, int maxIndexInclusive, CancellationToken cancellationToken)
+        {
+            AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
+
+            int count = elements.Length;
+
+            if (minIndexInclusive >= count)
+            {
+                return new Maybe<TElement>();
+            }
+
+            TElement[] array = elements.Array;
+
+            TElement last;
+
+            if (maxIndexInclusive < count - 1)
+            {
+                AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
+
+                last = await sorter.ElementAt(array, count, maxIndexInclusive).ConfigureAwait(false);
+            }
+            else
+            {
+                last = await Last(array, count, cancellationToken).ConfigureAwait(false);
+            }
+
+            return new Maybe<TElement>(last);
+        }
+
+        private async ValueTask<TElement> Last(TElement[] items, int count, CancellationToken cancellationToken)
+        {
+            TElement value = items[0];
+
+            AsyncCachingComparer<TElement> comparer = GetComparer();
+
+            await comparer.SetElement(value, cancellationToken).ConfigureAwait(false);
+
+            for (int i = 1; i != count; ++i)
+            {
+                TElement x = items[i];
+
+                if (await comparer.Compare(x, cacheLower: false, cancellationToken).ConfigureAwait(false) >= 0)
+                {
+                    value = x;
+                }
+            }
+
+            return value;
+        }
+
+        public IAsyncPartition<TElement> Skip(int count) => new OrderedAsyncPartition<TElement>(this, count, int.MaxValue);
+
+        public IAsyncPartition<TElement> Take(int count) => new OrderedAsyncPartition<TElement>(this, 0, count - 1);
+
+        public ValueTask<Maybe<TElement>> TryGetElementAtAsync(int index, CancellationToken cancellationToken)
+        {
+            if (index == 0)
+            {
+                return TryGetFirstAsync(cancellationToken);
+            }
+
+            if (index > 0)
+            {
+                return Core();
+
+                async ValueTask<Maybe<TElement>> Core()
+                {
+                    AsyncEnumerableHelpers.ArrayWithLength<TElement> elements = await AsyncEnumerableHelpers.ToArrayWithLength(_source, cancellationToken).ConfigureAwait(false);
+
+                    int count = elements.Length;
+
+                    if (index < count)
+                    {
+                        AsyncEnumerableSorter<TElement> sorter = GetAsyncEnumerableSorter(cancellationToken);
+
+                        TElement element = await sorter.ElementAt(elements.Array, count, index).ConfigureAwait(false);
+
+                        return new Maybe<TElement>(element);
+                    }
+
+                    return new Maybe<TElement>();
+                }
+            }
+
+            return new ValueTask<Maybe<TElement>>(new Maybe<TElement>());
+        }
     }
 
     internal sealed class OrderedAsyncEnumerable<TElement, TKey> : OrderedAsyncEnumerable<TElement>
@@ -152,11 +429,10 @@ namespace System.Linq
         private readonly OrderedAsyncEnumerable<TElement> _parent;
 
         public OrderedAsyncEnumerable(IAsyncEnumerable<TElement> source, Func<TElement, TKey> keySelector, IComparer<TKey> comparer, bool descending, OrderedAsyncEnumerable<TElement> parent)
+            : base(source)
         {
-            Debug.Assert(source != null);
             Debug.Assert(keySelector != null);
 
-            _source = source;
             _keySelector = keySelector;
             _comparer = comparer;
             _descending = descending;
@@ -179,6 +455,15 @@ namespace System.Linq
 
             return sorter;
         }
+
+        internal override AsyncCachingComparer<TElement> GetComparer(AsyncCachingComparer<TElement> childComparer)
+        {
+            AsyncCachingComparer<TElement> cmp = childComparer == null
+                ? new AsyncCachingComparer<TElement, TKey>(_keySelector, _comparer, _descending)
+                : new AsyncCachingComparerWithChild<TElement, TKey>(_keySelector, _comparer, _descending, childComparer);
+
+            return _parent != null ? _parent.GetComparer(cmp) : cmp;
+        }
     }
 
     internal sealed class OrderedAsyncEnumerableWithTask<TElement, TKey> : OrderedAsyncEnumerable<TElement>
@@ -189,11 +474,10 @@ namespace System.Linq
         private readonly OrderedAsyncEnumerable<TElement> _parent;
 
         public OrderedAsyncEnumerableWithTask(IAsyncEnumerable<TElement> source, Func<TElement, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending, OrderedAsyncEnumerable<TElement> parent)
+            : base(source)
         {
-            Debug.Assert(source != null);
             Debug.Assert(keySelector != null);
 
-            _source = source;
             _keySelector = keySelector;
             _comparer = comparer;
             _descending = descending;
@@ -216,6 +500,15 @@ namespace System.Linq
 
             return sorter;
         }
+
+        internal override AsyncCachingComparer<TElement> GetComparer(AsyncCachingComparer<TElement> childComparer)
+        {
+            AsyncCachingComparer<TElement> cmp = childComparer == null
+                ? new AsyncCachingComparerWithTask<TElement, TKey>(_keySelector, _comparer, _descending)
+                : new AsyncCachingComparerWithTaskAndChild<TElement, TKey>(_keySelector, _comparer, _descending, childComparer);
+
+            return _parent != null ? _parent.GetComparer(cmp) : cmp;
+        }
     }
 
 #if !NO_DEEP_CANCELLATION
@@ -227,11 +520,10 @@ namespace System.Linq
         private readonly OrderedAsyncEnumerable<TElement> _parent;
 
         public OrderedAsyncEnumerableWithTaskAndCancellation(IAsyncEnumerable<TElement> source, Func<TElement, CancellationToken, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending, OrderedAsyncEnumerable<TElement> parent)
+            : base(source)
         {
-            Debug.Assert(source != null);
             Debug.Assert(keySelector != null);
 
-            _source = source;
             _keySelector = keySelector;
             _comparer = comparer;
             _descending = descending;
@@ -254,6 +546,15 @@ namespace System.Linq
 
             return sorter;
         }
+
+        internal override AsyncCachingComparer<TElement> GetComparer(AsyncCachingComparer<TElement> childComparer)
+        {
+            AsyncCachingComparer<TElement> cmp = childComparer == null
+                ? new AsyncCachingComparerWithTaskAndCancellation<TElement, TKey>(_keySelector, _comparer, _descending)
+                : new AsyncCachingComparerWithTaskAndCancellationAndChild<TElement, TKey>(_keySelector, _comparer, _descending, childComparer);
+
+            return _parent != null ? _parent.GetComparer(cmp) : cmp;
+        }
     }
 #endif
 
@@ -261,9 +562,36 @@ namespace System.Linq
     {
         internal abstract ValueTask ComputeKeys(TElement[] elements, int count);
 
-        internal abstract int CompareKeys(int index1, int index2);
+        internal abstract int CompareAnyKeys(int index1, int index2);
 
         public async ValueTask<int[]> Sort(TElement[] elements, int count)
+        {
+            int[] map = await ComputeMap(elements, count).ConfigureAwait(false);
+
+            QuickSort(map, 0, count - 1);
+
+            return map;
+        }
+
+        public async ValueTask<int[]> Sort(TElement[] elements, int count, int minIndexInclusive, int maxIndexInclusive)
+        {
+            int[] map = await ComputeMap(elements, count).ConfigureAwait(false);
+
+            PartialQuickSort(map, 0, count - 1, minIndexInclusive, maxIndexInclusive);
+
+            return map;
+        }
+
+        public async ValueTask<TElement> ElementAt(TElement[] elements, int count, int index)
+        {
+            int[] map = await ComputeMap(elements, count).ConfigureAwait(false);
+
+            return index == 0 ?
+                elements[Min(map, count)] :
+                elements[QuickSelect(map, count - 1, index)];
+        }
+
+        private async ValueTask<int[]> ComputeMap(TElement[] elements, int count)
         {
             await ComputeKeys(elements, count).ConfigureAwait(false);
 
@@ -274,18 +602,57 @@ namespace System.Linq
                 map[i] = i;
             }
 
-            QuickSort(map, 0, count - 1);
-
             return map;
         }
 
-        private void QuickSort(int[] map, int left, int right)
+        protected abstract void QuickSort(int[] map, int left, int right);
+
+        protected abstract void PartialQuickSort(int[] map, int left, int right, int minIndexInclusive, int maxIndexInclusive);
+
+        protected abstract int QuickSelect(int[] map, int right, int idx);
+
+        protected abstract int Min(int[] map, int count);
+    }
+
+    internal abstract class AsyncEnumerableSorterBase<TElement, TKey> : AsyncEnumerableSorter<TElement>
+    {
+        private readonly IComparer<TKey> _comparer;
+        private readonly bool _descending;
+        protected readonly AsyncEnumerableSorter<TElement> _next;
+        protected TKey[] _keys;
+
+        public AsyncEnumerableSorterBase(IComparer<TKey> comparer, bool descending, AsyncEnumerableSorter<TElement> next)
+        {
+            _comparer = comparer ?? Comparer<TKey>.Default;
+            _descending = descending;
+            _next = next;
+        }
+
+        internal sealed override int CompareAnyKeys(int index1, int index2)
+        {
+            int c = _comparer.Compare(_keys[index1], _keys[index2]);
+
+            if (c == 0)
+            {
+                return _next == null ? index1 - index2 : _next.CompareAnyKeys(index1, index2);
+            }
+            else
+            {
+                return (_descending != (c > 0)) ? 1 : -1;
+            }
+        }
+
+        private int CompareKeys(int index1, int index2) => index1 == index2 ? 0 : CompareAnyKeys(index1, index2);
+
+        protected override void QuickSort(int[] map, int left, int right)
         {
+            // REVIEW: Consider using Array.Sort, see https://github.com/dotnet/corefx/commit/a6aff797a33e606a60ec0c9ca034a161c609620f#diff-d90239bd8284188e2bd210790483f5ca.
+
             do
             {
                 int i = left;
                 int j = right;
-                int x = map[i + (j - i >> 1)];
+                int x = map[i + ((j - i) >> 1)];
 
                 do
                 {
@@ -337,34 +704,156 @@ namespace System.Linq
             }
             while (left < right);
         }
-    }
-
-    internal abstract class AsyncEnumerableSorterBase<TElement, TKey> : AsyncEnumerableSorter<TElement>
-    {
-        private readonly IComparer<TKey> _comparer;
-        private readonly bool _descending;
-        protected readonly AsyncEnumerableSorter<TElement> _next;
-        protected TKey[] _keys;
 
-        public AsyncEnumerableSorterBase(IComparer<TKey> comparer, bool descending, AsyncEnumerableSorter<TElement> next)
+        protected override void PartialQuickSort(int[] map, int left, int right, int minIndexInclusive, int maxIndexInclusive)
         {
-            _comparer = comparer ?? Comparer<TKey>.Default;
-            _descending = descending;
-            _next = next;
+            do
+            {
+                int i = left;
+                int j = right;
+                int x = map[i + ((j - i) >> 1)];
+                do
+                {
+                    while (i < map.Length && CompareKeys(x, map[i]) > 0)
+                    {
+                        i++;
+                    }
+
+                    while (j >= 0 && CompareKeys(x, map[j]) < 0)
+                    {
+                        j--;
+                    }
+
+                    if (i > j)
+                    {
+                        break;
+                    }
+
+                    if (i < j)
+                    {
+                        int temp = map[i];
+                        map[i] = map[j];
+                        map[j] = temp;
+                    }
+
+                    i++;
+                    j--;
+                }
+                while (i <= j);
+
+                if (minIndexInclusive >= i)
+                {
+                    left = i + 1;
+                }
+                else if (maxIndexInclusive <= j)
+                {
+                    right = j - 1;
+                }
+
+                if (j - left <= right - i)
+                {
+                    if (left < j)
+                    {
+                        PartialQuickSort(map, left, j, minIndexInclusive, maxIndexInclusive);
+                    }
+
+                    left = i;
+                }
+                else
+                {
+                    if (i < right)
+                    {
+                        PartialQuickSort(map, i, right, minIndexInclusive, maxIndexInclusive);
+                    }
+
+                    right = j;
+                }
+            }
+            while (left < right);
         }
 
-        internal override int CompareKeys(int index1, int index2)
+        protected override int QuickSelect(int[] map, int right, int idx)
         {
-            int c = _comparer.Compare(_keys[index1], _keys[index2]);
-
-            if (c == 0)
+            int left = 0;
+            do
             {
-                return _next == null ? index1 - index2 : _next.CompareKeys(index1, index2);
+                int i = left;
+                int j = right;
+                int x = map[i + ((j - i) >> 1)];
+
+                do
+                {
+                    while (i < map.Length && CompareKeys(x, map[i]) > 0)
+                    {
+                        i++;
+                    }
+
+                    while (j >= 0 && CompareKeys(x, map[j]) < 0)
+                    {
+                        j--;
+                    }
+
+                    if (i > j)
+                    {
+                        break;
+                    }
+
+                    if (i < j)
+                    {
+                        int temp = map[i];
+                        map[i] = map[j];
+                        map[j] = temp;
+                    }
+
+                    i++;
+                    j--;
+                }
+                while (i <= j);
+
+                if (i <= idx)
+                {
+                    left = i + 1;
+                }
+                else
+                {
+                    right = j - 1;
+                }
+
+                if (j - left <= right - i)
+                {
+                    if (left < j)
+                    {
+                        right = j;
+                    }
+
+                    left = i;
+                }
+                else
+                {
+                    if (i < right)
+                    {
+                        left = i;
+                    }
+
+                    right = j;
+                }
             }
-            else
+            while (left < right);
+
+            return map[idx];
+        }
+
+        protected override int Min(int[] map, int count)
+        {
+            int index = 0;
+            for (int i = 1; i < count; i++)
             {
-                return (_descending != (c > 0)) ? 1 : -1;
+                if (CompareKeys(map[i], map[index]) < 0)
+                {
+                    index = i;
+                }
             }
+            return map[index];
         }
     }
 
@@ -449,4 +938,371 @@ namespace System.Linq
         }
     }
 #endif
+
+    internal sealed class OrderedAsyncPartition<TElement> : AsyncIterator<TElement>, IAsyncPartition<TElement>
+    {
+        private readonly OrderedAsyncEnumerable<TElement> _source;
+        private readonly int _minIndexInclusive;
+        private readonly int _maxIndexInclusive;
+
+        public OrderedAsyncPartition(OrderedAsyncEnumerable<TElement> source, int minIndexInclusive, int maxIndexInclusive)
+        {
+            _source = source;
+            _minIndexInclusive = minIndexInclusive;
+            _maxIndexInclusive = maxIndexInclusive;
+        }
+
+        public override AsyncIteratorBase<TElement> Clone() => new OrderedAsyncPartition<TElement>(_source, _minIndexInclusive, _maxIndexInclusive);
+
+        public ValueTask<int> GetCountAsync(bool onlyIfCheap, CancellationToken cancellationToken) =>
+            _source.GetCountAsync(_minIndexInclusive, _maxIndexInclusive, onlyIfCheap, cancellationToken);
+
+        public IAsyncPartition<TElement> Skip(int count)
+        {
+            int minIndex = unchecked(_minIndexInclusive + count);
+
+            if (unchecked((uint)minIndex > (uint)_maxIndexInclusive))
+            {
+                return AsyncEnumerable.EmptyAsyncIterator<TElement>.Instance;
+            }
+            
+            return new OrderedAsyncPartition<TElement>(_source, minIndex, _maxIndexInclusive);
+        }
+
+        public IAsyncPartition<TElement> Take(int count)
+        {
+            int maxIndex = unchecked(_minIndexInclusive + count - 1);
+
+            if (unchecked((uint)maxIndex >= (uint)_maxIndexInclusive))
+            {
+                return this;
+            }
+
+            return new OrderedAsyncPartition<TElement>(_source, _minIndexInclusive, maxIndex);
+        }
+
+        public ValueTask<TElement[]> ToArrayAsync(CancellationToken cancellationToken) =>
+            _source.ToArrayAsync(_minIndexInclusive, _maxIndexInclusive, cancellationToken);
+
+        public ValueTask<List<TElement>> ToListAsync(CancellationToken cancellationToken) =>
+            _source.ToListAsync(_minIndexInclusive, _maxIndexInclusive, cancellationToken);
+
+        public ValueTask<Maybe<TElement>> TryGetElementAtAsync(int index, CancellationToken cancellationToken)
+        {
+            if (unchecked((uint)index <= (uint)(_maxIndexInclusive - _minIndexInclusive)))
+            {
+                return _source.TryGetElementAtAsync(index + _minIndexInclusive, cancellationToken);
+            }
+
+            return new ValueTask<Maybe<TElement>>(new Maybe<TElement>());
+        }
+
+        public ValueTask<Maybe<TElement>> TryGetFirstAsync(CancellationToken cancellationToken) =>
+            _source.TryGetElementAtAsync(_minIndexInclusive, cancellationToken);
+
+        public ValueTask<Maybe<TElement>> TryGetLastAsync(CancellationToken cancellationToken) =>
+            _source.TryGetLastAsync(_minIndexInclusive, _maxIndexInclusive, cancellationToken);
+
+        // REVIEW: Consider to tear off an iterator object rather than storing this state here?
+
+        private TElement[] _buffer;
+        private int[] _indexes;
+        private int _minIndexIterator;
+        private int _maxIndexIterator;
+
+        protected override async ValueTask<bool> MoveNextCore()
+        {
+            switch (_state)
+            {
+                case AsyncIteratorState.Allocated:
+                    _buffer = await _source.ToArrayAsync(_cancellationToken).ConfigureAwait(false); // TODO: Use buffer.
+
+                    _minIndexIterator = _minIndexInclusive;
+                    _maxIndexIterator = _maxIndexInclusive;
+
+                    int count = _buffer.Length;
+
+                    if (count > _minIndexIterator)
+                    {
+                        if (count <= _maxIndexIterator)
+                        {
+                            _maxIndexIterator = count - 1;
+                        }
+
+                        AsyncEnumerableSorter<TElement> sorter = _source.GetAsyncEnumerableSorter(_cancellationToken);
+
+                        if (_minIndexIterator == _maxIndexIterator)
+                        {
+                            _current = await sorter.ElementAt(_buffer, _buffer.Length, _minIndexIterator).ConfigureAwait(false);
+
+                            _minIndexIterator = int.MaxValue;
+                            _maxIndexIterator = int.MinValue;
+
+                            _state = AsyncIteratorState.Iterating;
+                            return true;
+                        }
+                        else
+                        {
+                            _indexes = await sorter.Sort(_buffer, _buffer.Length, _minIndexIterator, _maxIndexIterator).ConfigureAwait(false);
+                        }
+
+                        _state = AsyncIteratorState.Iterating;
+                        goto case AsyncIteratorState.Iterating;
+                    }
+
+                    await DisposeAsync();
+                    break;
+
+                case AsyncIteratorState.Iterating:
+                    if (_minIndexIterator <= _maxIndexIterator)
+                    {
+                        _current = _buffer[_indexes[_minIndexIterator++]];
+                        return true;
+                    }
+
+                    await DisposeAsync().ConfigureAwait(false);
+                    break;
+            }
+
+            return false;
+        }
+
+        public override async ValueTask DisposeAsync()
+        {
+            _buffer = null;
+            _indexes = null;
+
+            await base.DisposeAsync().ConfigureAwait(false);
+        }
+    }
+
+    internal abstract class AsyncCachingComparer<TElement>
+    {
+        internal abstract ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken);
+
+        internal abstract ValueTask SetElement(TElement element, CancellationToken cancellationToken);
+    }
+
+    internal class AsyncCachingComparer<TElement, TKey> : AsyncCachingComparer<TElement>
+    {
+        protected readonly Func<TElement, TKey> _keySelector;
+        protected readonly IComparer<TKey> _comparer;
+        protected readonly bool _descending;
+        protected TKey _lastKey;
+
+        public AsyncCachingComparer(Func<TElement, TKey> keySelector, IComparer<TKey> comparer, bool descending)
+        {
+            _keySelector = keySelector;
+            _comparer = comparer;
+            _descending = descending;
+        }
+
+        internal override ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = _keySelector(element);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+            }
+
+            return new ValueTask<int>(cmp);
+        }
+
+        internal override ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            _lastKey = _keySelector(element);
+
+            return new ValueTask();
+        }
+    }
+
+    internal sealed class AsyncCachingComparerWithChild<TElement, TKey> : AsyncCachingComparer<TElement, TKey>
+    {
+        private readonly AsyncCachingComparer<TElement> _child;
+
+        public AsyncCachingComparerWithChild(Func<TElement, TKey> keySelector, IComparer<TKey> comparer, bool descending, AsyncCachingComparer<TElement> child)
+            : base(keySelector, comparer, descending)
+        {
+            _child = child;
+        }
+
+        internal override async ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = _keySelector(element);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cmp == 0)
+            {
+                return await _child.Compare(element, cacheLower, cancellationToken).ConfigureAwait(false);
+            }
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+
+                await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+            }
+
+            return cmp;
+        }
+
+        internal override async ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            await base.SetElement(element, cancellationToken).ConfigureAwait(false);
+
+            await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+    internal class AsyncCachingComparerWithTask<TElement, TKey> : AsyncCachingComparer<TElement>
+    {
+        protected readonly Func<TElement, ValueTask<TKey>> _keySelector;
+        protected readonly IComparer<TKey> _comparer;
+        protected readonly bool _descending;
+        protected TKey _lastKey;
+
+        public AsyncCachingComparerWithTask(Func<TElement, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending)
+        {
+            _keySelector = keySelector;
+            _comparer = comparer;
+            _descending = descending;
+        }
+
+        internal override async ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = await _keySelector(element).ConfigureAwait(false);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+            }
+
+            return cmp;
+        }
+
+        internal override async ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            _lastKey = await _keySelector(element).ConfigureAwait(false);
+        }
+    }
+
+    internal sealed class AsyncCachingComparerWithTaskAndChild<TElement, TKey> : AsyncCachingComparerWithTask<TElement, TKey>
+    {
+        private readonly AsyncCachingComparer<TElement> _child;
+
+        public AsyncCachingComparerWithTaskAndChild(Func<TElement, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending, AsyncCachingComparer<TElement> child)
+            : base(keySelector, comparer, descending)
+        {
+            _child = child;
+        }
+
+        internal override async ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = await _keySelector(element).ConfigureAwait(false);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cmp == 0)
+            {
+                return await _child.Compare(element, cacheLower, cancellationToken).ConfigureAwait(false);
+            }
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+
+                await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+            }
+
+            return cmp;
+        }
+
+        internal override async ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            await base.SetElement(element, cancellationToken).ConfigureAwait(false);
+
+            await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+#if !NO_DEEP_CANCELLATION
+    internal class AsyncCachingComparerWithTaskAndCancellation<TElement, TKey> : AsyncCachingComparer<TElement>
+    {
+        protected readonly Func<TElement, CancellationToken, ValueTask<TKey>> _keySelector;
+        protected readonly IComparer<TKey> _comparer;
+        protected readonly bool _descending;
+        protected TKey _lastKey;
+
+        public AsyncCachingComparerWithTaskAndCancellation(Func<TElement, CancellationToken, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending)
+        {
+            _keySelector = keySelector;
+            _comparer = comparer;
+            _descending = descending;
+        }
+
+        internal override async ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = await _keySelector(element, cancellationToken).ConfigureAwait(false);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+            }
+
+            return cmp;
+        }
+
+        internal override async ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            _lastKey = await _keySelector(element, cancellationToken).ConfigureAwait(false);
+        }
+    }
+
+    internal sealed class AsyncCachingComparerWithTaskAndCancellationAndChild<TElement, TKey> : AsyncCachingComparerWithTaskAndCancellation<TElement, TKey>
+    {
+        private readonly AsyncCachingComparer<TElement> _child;
+
+        public AsyncCachingComparerWithTaskAndCancellationAndChild(Func<TElement, CancellationToken, ValueTask<TKey>> keySelector, IComparer<TKey> comparer, bool descending, AsyncCachingComparer<TElement> child)
+            : base(keySelector, comparer, descending)
+        {
+            _child = child;
+        }
+
+        internal override async ValueTask<int> Compare(TElement element, bool cacheLower, CancellationToken cancellationToken)
+        {
+            TKey newKey = await _keySelector(element, cancellationToken).ConfigureAwait(false);
+
+            int cmp = _descending ? _comparer.Compare(_lastKey, newKey) : _comparer.Compare(newKey, _lastKey);
+
+            if (cmp == 0)
+            {
+                return await _child.Compare(element, cacheLower, cancellationToken).ConfigureAwait(false);
+            }
+
+            if (cacheLower == cmp < 0)
+            {
+                _lastKey = newKey;
+
+                await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+            }
+
+            return cmp;
+        }
+
+        internal override async ValueTask SetElement(TElement element, CancellationToken cancellationToken)
+        {
+            await base.SetElement(element, cancellationToken).ConfigureAwait(false);
+
+            await _child.SetElement(element, cancellationToken).ConfigureAwait(false);
+        }
+    }
+#endif
 }