// 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 System;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using Microsoft.CodeAnalysis.Operations;
using Microsoft.CodeAnalysis.PooledObjects;
using Microsoft.CodeAnalysis.Shared.Extensions;

namespace Microsoft.CodeAnalysis.FlowAnalysis.SymbolUsageAnalysis
{
    internal static partial class SymbolUsageAnalysis
    {
        /// <summary>
        /// Core analysis data to drive the operation <see cref="Walker"/>
        /// for operation tree based analysis OR control flow graph based analysis.
        /// </summary>
        private abstract class AnalysisData : IDisposable
        {
            /// <summary>
            /// Pooled <see cref="BasicBlockAnalysisData"/> allocated during analysis with the
            /// current <see cref="AnalysisData"/> instance, which will be freed during <see cref="Dispose"/>.
            /// </summary>
            private readonly ArrayBuilder<BasicBlockAnalysisData> _allocatedBasicBlockAnalysisDatas;

            /// <summary>
            /// Set of locals/parameters which are passed by reference to other method calls.
            /// </summary>
            private readonly PooledHashSet<ISymbol> _referenceTakenSymbolsBuilder;

            protected AnalysisData()
            {
                _allocatedBasicBlockAnalysisDatas = ArrayBuilder<BasicBlockAnalysisData>.GetInstance();
                _referenceTakenSymbolsBuilder = PooledHashSet<ISymbol>.GetInstance();
                CurrentBlockAnalysisData = CreateBlockAnalysisData();
                AdditionalConditionalBranchAnalysisData = CreateBlockAnalysisData();
            }

            /// <summary>
            /// Map from each (symbol, write) to a boolean indicating if the value assigned
            /// at the write is read on some control flow path.
            /// For example, consider the following code:
            /// <code>
            ///     int x = 0;
            ///     x = 1;
            ///     Console.WriteLine(x);
            /// </code>
            /// This map will have two entries for 'x':
            ///     1. Key = (symbol: x, write: 'int x = 0')
            ///        Value = 'false', because value assigned to 'x' here **is never** read. 
            ///     2. Key = (symbol: x, write: 'x = 1')
            ///        Value = 'true', because value assigned to 'x' here **may be** read on
            ///        some control flow path.
            /// </summary>
            protected abstract PooledDictionary<(ISymbol symbol, IOperation operation), bool> SymbolsWriteBuilder { get; }

            /// <summary>
            /// Set of locals/parameters that are read at least once.
            /// </summary>
            protected abstract PooledHashSet<ISymbol> SymbolsReadBuilder { get; }

            /// <summary>
            /// Set of lambda/local functions whose invocations are currently being analyzed to prevent
            /// infinite recursion for analyzing code with recursive lambda/local function calls.
            /// </summary>
            protected abstract PooledHashSet<IMethodSymbol> LambdaOrLocalFunctionsBeingAnalyzed { get; }

            /// <summary>
            /// Current block analysis data used for analysis.
            /// </summary>
            public BasicBlockAnalysisData CurrentBlockAnalysisData { get; }

            /// <summary>
            /// Block analysis data used for an additional conditional branch.
            /// </summary>
            public BasicBlockAnalysisData AdditionalConditionalBranchAnalysisData { get; }

            /// <summary>
            /// Creates an immutable <see cref="SymbolUsageResult"/> for the current analysis data.
            /// </summary>
            public SymbolUsageResult ToResult()
                => new(SymbolsWriteBuilder.ToImmutableDictionary(),
                                         SymbolsReadBuilder.ToImmutableHashSet());

            public BasicBlockAnalysisData AnalyzeLocalFunctionInvocation(IMethodSymbol localFunction, CancellationToken cancellationToken)
            {
                Debug.Assert(localFunction.IsLocalFunction());

                // Use the original definition of the local function for flow analysis.
                localFunction = localFunction.OriginalDefinition;

                if (!LambdaOrLocalFunctionsBeingAnalyzed.Add(localFunction))
                {
                    ResetState();
                    return CurrentBlockAnalysisData;
                }
                else
                {
                    var result = AnalyzeLocalFunctionInvocationCore(localFunction, cancellationToken);
                    LambdaOrLocalFunctionsBeingAnalyzed.Remove(localFunction);
                    return result;
                }
            }

            public BasicBlockAnalysisData AnalyzeLambdaInvocation(IFlowAnonymousFunctionOperation lambda, CancellationToken cancellationToken)
            {
                if (!LambdaOrLocalFunctionsBeingAnalyzed.Add(lambda.Symbol))
                {
                    ResetState();
                    return CurrentBlockAnalysisData;
                }
                else
                {
                    var result = AnalyzeLambdaInvocationCore(lambda, cancellationToken);
                    LambdaOrLocalFunctionsBeingAnalyzed.Remove(lambda.Symbol);
                    return result;
                }
            }

            protected abstract BasicBlockAnalysisData AnalyzeLocalFunctionInvocationCore(IMethodSymbol localFunction, CancellationToken cancellationToken);
            protected abstract BasicBlockAnalysisData AnalyzeLambdaInvocationCore(IFlowAnonymousFunctionOperation lambda, CancellationToken cancellationToken);

            // Methods specific to flow capture analysis for CFG based dataflow analysis.
            public abstract bool IsLValueFlowCapture(CaptureId captureId);
            public abstract bool IsRValueFlowCapture(CaptureId captureId);
            public abstract void OnLValueCaptureFound(ISymbol symbol, IOperation operation, CaptureId captureId);
            public abstract void OnLValueDereferenceFound(CaptureId captureId);

            // Methods specific to delegate analysis to track potential delegate invocation targets for CFG based dataflow analysis.
            public abstract bool IsTrackingDelegateCreationTargets { get; }
            public abstract void SetTargetsFromSymbolForDelegate(IOperation write, ISymbol symbol);
            public abstract void SetLambdaTargetForDelegate(IOperation write, IFlowAnonymousFunctionOperation lambdaTarget);
            public abstract void SetLocalFunctionTargetForDelegate(IOperation write, IMethodReferenceOperation localFunctionTarget);
            public abstract void SetEmptyInvocationTargetsForDelegate(IOperation write);
            public abstract bool TryGetDelegateInvocationTargets(IOperation write, out ImmutableHashSet<IOperation> targets);

            protected static PooledDictionary<(ISymbol Symbol, IOperation Write), bool> CreateSymbolsWriteMap(
                ImmutableArray<IParameterSymbol> parameters)
            {
                var symbolsWriteMap = PooledDictionary<(ISymbol Symbol, IOperation Write), bool>.GetInstance();
                return UpdateSymbolsWriteMap(symbolsWriteMap, parameters);
            }

            protected static PooledDictionary<(ISymbol Symbol, IOperation Write), bool> UpdateSymbolsWriteMap(
                PooledDictionary<(ISymbol Symbol, IOperation Write), bool> symbolsWriteMap,
                ImmutableArray<IParameterSymbol> parameters)
            {
                // Mark parameters as being written from the value provided at the call site.
                // Note that the write operation is "null" as there is no corresponding IOperation for parameter definition.
                foreach (var parameter in parameters)
                {
                    (ISymbol, IOperation) key = (parameter, null);
                    if (!symbolsWriteMap.ContainsKey(key))
                    {
                        symbolsWriteMap.Add(key, false);
                    }
                }

                return symbolsWriteMap;
            }

            public BasicBlockAnalysisData CreateBlockAnalysisData()
            {
                var instance = BasicBlockAnalysisData.GetInstance();
                TrackAllocatedBlockAnalysisData(instance);
                return instance;
            }

            public void TrackAllocatedBlockAnalysisData(BasicBlockAnalysisData allocatedData)
                => _allocatedBasicBlockAnalysisDatas.Add(allocatedData);

            public void OnReadReferenceFound(ISymbol symbol)
            {
                if (symbol.Kind == SymbolKind.Discard)
                {
                    return;
                }

                // Mark all the current reaching writes of symbol as read.
                if (SymbolsWriteBuilder.Count != 0)
                {
                    CurrentBlockAnalysisData.ForEachCurrentWrite(
                        symbol,
                        static (write, arg) =>
                        {
                            arg.self.SymbolsWriteBuilder[(arg.symbol, write)] = true;
                        },
                        (symbol, self: this));
                }

                // Mark the current symbol as read.
                SymbolsReadBuilder.Add(symbol);
            }

            public void OnWriteReferenceFound(ISymbol symbol, IOperation operation, bool maybeWritten, bool isRef)
            {
                var symbolAndWrite = (symbol, operation);
                if (symbol.Kind == SymbolKind.Discard)
                {
                    // Skip discard symbols and also for already processed writes (back edge from loops).
                    return;
                }

                if (_referenceTakenSymbolsBuilder.Contains(symbol))
                {
                    // Skip tracking writes for reference taken symbols as the written value may be read from a different variable.
                    return;
                }

                // Add a new write for the given symbol at the given operation.
                CurrentBlockAnalysisData.OnWriteReferenceFound(symbol, operation, maybeWritten);

                if (isRef)
                {
                    _referenceTakenSymbolsBuilder.Add(symbol);
                }

                // Only mark as unused write if we are processing it for the first time (not from back edge for loops)
                if (!SymbolsWriteBuilder.ContainsKey(symbolAndWrite) &&
                    !maybeWritten)
                {
                    SymbolsWriteBuilder.Add((symbol, operation), false);
                }
            }

            /// <summary>
            /// Resets all the currently tracked symbol writes to be conservatively marked as read.
            /// </summary>
            public void ResetState()
            {
                foreach (var symbol in SymbolsWriteBuilder.Keys.Select(d => d.symbol).ToArray())
                {
                    OnReadReferenceFound(symbol);
                }
            }

            public void SetCurrentBlockAnalysisDataFrom(BasicBlockAnalysisData newBlockAnalysisData)
            {
                Debug.Assert(newBlockAnalysisData != null);
                CurrentBlockAnalysisData.SetAnalysisDataFrom(newBlockAnalysisData);
            }

            public virtual void Dispose()
            {
                foreach (var instance in _allocatedBasicBlockAnalysisDatas)
                {
                    instance.Dispose();
                }

                _allocatedBasicBlockAnalysisDatas.Free();
                _referenceTakenSymbolsBuilder.Free();
            }
        }
    }
}
