// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

#nullable disable

using Microsoft.CodeAnalysis.CSharp.Symbols;
using Roslyn.Utilities;
using System.Collections.Generic;

namespace Microsoft.CodeAnalysis.CSharp
{
    internal partial class IteratorMethodToStateMachineRewriter
    {
        /// <summary>
        /// Analyzes method body for yields in try blocks and labels that they contain.
        /// </summary>
        private sealed class YieldsInTryAnalysis : LabelCollector
        {
            // all try blocks with yields in them and complete set of labels inside those try blocks
            // NOTE: non-yielding try blocks are transparently ignored - i.e. their labels are included
            //       in the label set of the nearest yielding-try parent  
            private Dictionary<BoundTryStatement, HashSet<LabelSymbol>> _labelsInYieldingTrys;

            // transient accumulators.
            private bool _seenYield;

            public YieldsInTryAnalysis(BoundStatement body)
            {
                _seenYield = false;
                this.Visit(body);
            }

            /// <summary>
            /// Returns true if given try or any of its nested try blocks contain yields
            /// </summary>
            public bool ContainsYields(BoundTryStatement statement)
            {
                return _labelsInYieldingTrys != null && _labelsInYieldingTrys.ContainsKey(statement);
            }

            /// <summary>
            /// Returns true if body contains yield returns within try blocks.
            /// </summary>
            public bool ContainsYieldsInTrys()
            {
                return _labelsInYieldingTrys != null;
            }

            /// <summary>
            /// Labels reachable from within this frame without invoking its finally. 
            /// null if there are none such labels.
            /// </summary>
            internal HashSet<LabelSymbol> Labels(BoundTryStatement statement)
            {
                return _labelsInYieldingTrys[statement];
            }

            public override BoundNode VisitTryStatement(BoundTryStatement node)
            {
                var origSeenYield = _seenYield;
                var origLabels = this.currentLabels;

                // sibling try blocks do not see each other's yields
                _seenYield = false;
                this.currentLabels = null;

                base.VisitTryStatement(node);

                if (_seenYield)
                {
                    // this try yields !

                    var yieldingTryLabels = _labelsInYieldingTrys;
                    if (yieldingTryLabels == null)
                    {
                        _labelsInYieldingTrys = yieldingTryLabels = new Dictionary<BoundTryStatement, HashSet<LabelSymbol>>();
                    }

                    yieldingTryLabels.Add(node, currentLabels);
                    currentLabels = origLabels;
                }
                else
                {
                    // this is a boring non-yielding try

                    // currentLabels = currentLabels U origLabels ;
                    if (currentLabels == null)
                    {
                        currentLabels = origLabels;
                    }
                    else if (origLabels != null)
                    {
                        currentLabels.UnionWith(origLabels);
                    }
                }

                _seenYield = _seenYield | origSeenYield;
                return null;
            }

            public override BoundNode VisitYieldReturnStatement(BoundYieldReturnStatement node)
            {
                _seenYield = true;
                return base.VisitYieldReturnStatement(node);
            }

            public override BoundNode VisitExpressionStatement(BoundExpressionStatement node)
            {
                // expressions cannot contain labels, branches or yields.
                return null;
            }
        }
    }

    /// <summary>
    /// Analyzes method body for labels.
    /// </summary>
    internal abstract class LabelCollector : BoundTreeWalkerWithStackGuardWithoutRecursionOnTheLeftOfBinaryOperator
    {
        // transient accumulator.
        protected HashSet<LabelSymbol> currentLabels;

        public override BoundNode VisitLabelStatement(BoundLabelStatement node)
        {
            CollectLabel(node.Label);
            return base.VisitLabelStatement(node);
        }

        private void CollectLabel(LabelSymbol label)
        {
            if ((object)label != null)
            {
                var currentLabels = this.currentLabels;
                if (currentLabels == null)
                {
                    this.currentLabels = currentLabels = new HashSet<LabelSymbol>();
                }
                currentLabels.Add(label);
            }
        }
    }
}