// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

// =+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+=+
//
// ElementAtQueryOperator.cs
//
// =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-

using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Threading;

namespace System.Linq.Parallel
{
    /// <summary>
    /// ElementAt just retrieves an element at a specific index.  There is some cross-partition
    /// coordination to force partitions to stop looking once a partition has found the
    /// sought-after element.
    /// </summary>
    /// <typeparam name="TSource"></typeparam>
    internal sealed class ElementAtQueryOperator<TSource> : UnaryQueryOperator<TSource, TSource>
    {
        private readonly int _index; // The index that we're looking for.
        private readonly bool _prematureMerge; // Whether to prematurely merge the input of this operator.
        private readonly bool _limitsParallelism; // Whether this operator limits parallelism

        //---------------------------------------------------------------------------------------
        // Constructs a new instance of the contains search operator.
        //
        // Arguments:
        //     child       - the child tree to enumerate.
        //     index       - index we are searching for.
        //

        internal ElementAtQueryOperator(IEnumerable<TSource> child, int index)
            : base(child)
        {
            Debug.Assert(child != null, "child data source cannot be null");
            Debug.Assert(index >= 0, "index can't be less than 0");
            _index = index;

            OrdinalIndexState childIndexState = Child.OrdinalIndexState;
            if (ExchangeUtilities.IsWorseThan(childIndexState, OrdinalIndexState.Correct))
            {
                _prematureMerge = true;
                _limitsParallelism = childIndexState != OrdinalIndexState.Shuffled;
            }
        }

        //---------------------------------------------------------------------------------------
        // Just opens the current operator, including opening the child and wrapping it with
        // partitions as needed.
        //

        internal override QueryResults<TSource> Open(
            QuerySettings settings, bool preferStriping)
        {
            // We just open the child operator.
            QueryResults<TSource> childQueryResults = Child.Open(settings, false);
            return new UnaryQueryOperatorResults(childQueryResults, this, settings, preferStriping);
        }

        internal override void WrapPartitionedStream<TKey>(
            PartitionedStream<TSource, TKey> inputStream, IPartitionedStreamRecipient<TSource> recipient, bool preferStriping, QuerySettings settings)
        {
            // If the child OOP index is not correct, reindex.
            int partitionCount = inputStream.PartitionCount;

            PartitionedStream<TSource, int> intKeyStream;
            if (_prematureMerge)
            {
                intKeyStream = ExecuteAndCollectResults(inputStream, partitionCount, Child.OutputOrdered, preferStriping, settings).GetPartitionedStream();
                Debug.Assert(intKeyStream.OrdinalIndexState == OrdinalIndexState.Indexable);
            }
            else
            {
                intKeyStream = (PartitionedStream<TSource, int>)(object)inputStream;
            }

            // Create a shared cancellation variable and then return a possibly wrapped new enumerator.
            Shared<bool> resultFoundFlag = new Shared<bool>(false);

            PartitionedStream<TSource, int> outputStream = new PartitionedStream<TSource, int>(
                partitionCount, Util.GetDefaultComparer<int>(), OrdinalIndexState.Correct);

            for (int i = 0; i < partitionCount; i++)
            {
                outputStream[i] = new ElementAtQueryOperatorEnumerator(intKeyStream[i], _index, resultFoundFlag, settings.CancellationState.MergedCancellationToken);
            }

            recipient.Receive(outputStream);
        }

        //---------------------------------------------------------------------------------------
        // Returns an enumerable that represents the query executing sequentially.
        //

        [ExcludeFromCodeCoverage(Justification = "This method should never be called as fallback to sequential is handled in Aggregate()")]
        internal override IEnumerable<TSource> AsSequentialQuery(CancellationToken token)
        {
            Debug.Fail("This method should never be called as fallback to sequential is handled in Aggregate().");
            throw new NotSupportedException();
        }

        //---------------------------------------------------------------------------------------
        // Whether this operator performs a premature merge that would not be performed in
        // a similar sequential operation (i.e., in LINQ to Objects).
        //

        internal override bool LimitsParallelism
        {
            get { return _limitsParallelism; }
        }


        /// <summary>
        /// Executes the query, either sequentially or in parallel, depending on the query execution mode and
        /// whether a premature merge was inserted by this ElementAt operator.
        /// </summary>
        /// <param name="result">result</param>
        /// <param name="withDefaultValue">withDefaultValue</param>
        /// <returns>whether an element with this index exists</returns>
        internal bool Aggregate([MaybeNullWhen(false)] out TSource result, bool withDefaultValue)
        {
            // If we were to insert a premature merge before this ElementAt, and we are executing in conservative mode, run the whole query
            // sequentially.
            if (LimitsParallelism && SpecifiedQuerySettings.WithDefaults().ExecutionMode!.Value != ParallelExecutionMode.ForceParallelism)
            {
                CancellationState cancelState = SpecifiedQuerySettings.CancellationState;
                if (withDefaultValue)
                {
                    IEnumerable<TSource> childAsSequential = Child.AsSequentialQuery(cancelState.ExternalCancellationToken);
                    IEnumerable<TSource> childWithCancelChecks = CancellableEnumerable.Wrap(childAsSequential, cancelState.ExternalCancellationToken);
                    result = ExceptionAggregator.WrapEnumerable(childWithCancelChecks, cancelState).ElementAtOrDefault(_index)!;
                }
                else
                {
                    IEnumerable<TSource> childAsSequential = Child.AsSequentialQuery(cancelState.ExternalCancellationToken);
                    IEnumerable<TSource> childWithCancelChecks = CancellableEnumerable.Wrap(childAsSequential, cancelState.ExternalCancellationToken);
                    result = ExceptionAggregator.WrapEnumerable(childWithCancelChecks, cancelState).ElementAt(_index);
                }
                return true;
            }

            using (IEnumerator<TSource> e = GetEnumerator(ParallelMergeOptions.FullyBuffered))
            {
                if (e.MoveNext())
                {
                    TSource current = e.Current;
                    Debug.Assert(!e.MoveNext(), "expected enumerator to be empty");
                    result = current;
                    return true;
                }
            }

            result = default(TSource)!;
            return false;
        }


        //---------------------------------------------------------------------------------------
        // This enumerator performs the search for the element at the specified index.
        //

        private sealed class ElementAtQueryOperatorEnumerator : QueryOperatorEnumerator<TSource, int>
        {
            private readonly QueryOperatorEnumerator<TSource, int> _source; // The source data.
            private readonly int _index; // The index of the element to seek.
            private readonly Shared<bool> _resultFoundFlag; // Whether to cancel the operation.
            private readonly CancellationToken _cancellationToken;

            //---------------------------------------------------------------------------------------
            // Instantiates a new any/all search operator.
            //

            internal ElementAtQueryOperatorEnumerator(QueryOperatorEnumerator<TSource, int> source,
                                                      int index, Shared<bool> resultFoundFlag,
                CancellationToken cancellationToken)
            {
                Debug.Assert(source != null);
                Debug.Assert(index >= 0);
                Debug.Assert(resultFoundFlag != null);

                _source = source;
                _index = index;
                _resultFoundFlag = resultFoundFlag;
                _cancellationToken = cancellationToken;
            }

            //---------------------------------------------------------------------------------------
            // Enumerates the entire input until the element with the specified is found or another
            // partition has signaled that it found the element.
            //

            internal override bool MoveNext([MaybeNullWhen(false), AllowNull] ref TSource currentElement, ref int currentKey)
            {
                // Just walk the enumerator until we've found the element.
                int i = 0;
                while (_source.MoveNext(ref currentElement!, ref currentKey))
                {
                    if ((i++ & CancellationState.POLL_INTERVAL) == 0)
                        _cancellationToken.ThrowIfCancellationRequested();

                    if (_resultFoundFlag.Value)
                    {
                        // Another partition found the element.
                        break;
                    }

                    if (currentKey == _index)
                    {
                        // We have found the element. Cancel other searches and return true.
                        _resultFoundFlag.Value = true;
                        return true;
                    }
                }

                return false;
            }

            protected override void Dispose(bool disposing)
            {
                Debug.Assert(_source != null);
                _source.Dispose();
            }
        }
    }
}
