' 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.

Imports System.Collections.Immutable
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols

Namespace Microsoft.CodeAnalysis.VisualBasic.ExpressionEvaluator
    Friend NotInheritable Class CapturedVariableRewriter
        Inherits BoundTreeRewriterWithStackGuardWithoutRecursionOnTheLeftOfBinaryOperator

        Friend Shared Function Rewrite(
            targetMethodMeParameter As ParameterSymbol,
            displayClassVariables As ImmutableDictionary(Of String, DisplayClassVariable),
            node As BoundNode,
            diagnostics As DiagnosticBag) As BoundNode

            Dim rewriter = New CapturedVariableRewriter(targetMethodMeParameter, displayClassVariables, diagnostics)
            Return rewriter.Visit(node)
        End Function

        Private ReadOnly _targetMethodMeParameter As ParameterSymbol
        Private ReadOnly _displayClassVariables As ImmutableDictionary(Of String, DisplayClassVariable)
        Private ReadOnly _diagnostics As DiagnosticBag

        Private Sub New(
            targetMethodMeParameter As ParameterSymbol,
            displayClassVariables As ImmutableDictionary(Of String, DisplayClassVariable),
            diagnostics As DiagnosticBag)

            _targetMethodMeParameter = targetMethodMeParameter
            _displayClassVariables = displayClassVariables
            _diagnostics = diagnostics
        End Sub

        Public Overrides Function Visit(node As BoundNode) As BoundNode
            ' Ignore nodes that will be rewritten to literals in the LocalRewriter.
            If TryCast(node, BoundExpression)?.ConstantValueOpt IsNot Nothing Then
                Return node
            End If

            Return MyBase.Visit(node)
        End Function

        Public Overrides Function VisitBlock(node As BoundBlock) As BoundNode
            Dim rewrittenLocals = node.Locals.WhereAsArray(AddressOf IncludeLocal)
            Dim rewrittenStatements = VisitList(node.Statements)
            Return node.Update(node.StatementListSyntax, rewrittenLocals, rewrittenStatements)
        End Function

        Private Function IncludeLocal(local As LocalSymbol) As Boolean
            Return Not local.IsStatic AndAlso
                (local.IsCompilerGenerated OrElse local.Name Is Nothing OrElse GetVariable(local.Name) Is Nothing)
        End Function

        Public Overrides Function VisitLocal(node As BoundLocal) As BoundNode
            Dim local = node.LocalSymbol
            If Not local.IsCompilerGenerated Then
                Dim syntax = node.Syntax
                Dim staticLocal = TryCast(local, EEStaticLocalSymbol)
                If staticLocal IsNot Nothing Then
                    Dim receiver = If(_targetMethodMeParameter Is Nothing,
                        Nothing,
                        GetRewrittenMeParameter(syntax, New BoundMeReference(syntax, _targetMethodMeParameter.Type)))
                    Dim result = staticLocal.ToBoundExpression(receiver, syntax, node.IsLValue)
                    Debug.Assert(TypeSymbol.Equals(node.Type, result.Type, TypeCompareKind.ConsiderEverything))
                    Return result
                End If
                Dim variable = GetVariable(local.Name)
                If variable IsNot Nothing Then
                    Dim result = variable.ToBoundExpression(syntax, node.IsLValue, node.SuppressVirtualCalls)
                    Debug.Assert(TypeSymbol.Equals(node.Type, result.Type, TypeCompareKind.ConsiderEverything))
                    Return result
                End If
            End If
            Return node
        End Function

        Public Overrides Function VisitParameter(node As BoundParameter) As BoundNode
            Return RewriteParameter(node.Syntax, node.ParameterSymbol, node)
        End Function

        Public Overrides Function VisitMeReference(node As BoundMeReference) As BoundNode
            Return Me.GetRewrittenMeParameter(node.Syntax, node)
        End Function

        Public Overrides Function VisitMyBaseReference(node As BoundMyBaseReference) As BoundNode
            ' Rewrite as a "Me" reference with a conversion
            ' to the base type and with no virtual calls.
            Debug.Assert(node.SuppressVirtualCalls)

            Dim syntax = node.Syntax

            Dim meParameter = Me.GetRewrittenMeParameter(syntax, node)
            Debug.Assert(meParameter.Type.TypeKind = TypeKind.Class) ' Illegal in structures and modules.

            Dim baseType = node.Type
            Debug.Assert(baseType.TypeKind = TypeKind.Class) ' Illegal in structures and modules.

            Dim result = New BoundDirectCast(
                syntax:=syntax,
                operand:=meParameter,
                conversionKind:=ConversionKind.WideningReference, ' From a class to its base class.
                suppressVirtualCalls:=node.SuppressVirtualCalls,
                constantValueOpt:=Nothing,
                relaxationLambdaOpt:=Nothing,
                type:=baseType)
            Debug.Assert(TypeSymbol.Equals(result.Type, node.Type, TypeCompareKind.ConsiderEverything))
            Return result
        End Function

        Public Overrides Function VisitMyClassReference(node As BoundMyClassReference) As BoundNode
            ' MyClass is just Me with virtual calls suppressed.
            Debug.Assert(node.SuppressVirtualCalls)
            Return Me.GetRewrittenMeParameter(node.Syntax, node)
        End Function

        Public Overrides Function VisitValueTypeMeReference(node As BoundValueTypeMeReference) As BoundNode
            ' ValueTypeMe is just Me with IsLValue true.
            Debug.Assert(node.IsLValue)
            Return Me.GetRewrittenMeParameter(node.Syntax, node)
        End Function

        Private Function GetRewrittenMeParameter(syntax As SyntaxNode, node As BoundExpression) As BoundExpression
            If _targetMethodMeParameter Is Nothing Then
                ReportMissingMe(node.Syntax)
                Return node
            End If

            Dim result = RewriteParameter(syntax, _targetMethodMeParameter, node)
            Debug.Assert(result IsNot Nothing)
            Return result
        End Function

        Private Function RewriteParameter(syntax As SyntaxNode, symbol As ParameterSymbol, node As BoundExpression) As BoundExpression
            Dim name As String = symbol.Name
            Dim variable = Me.GetVariable(name)
            If variable Is Nothing Then
                ' The state machine case is for async lambdas.  The state machine
                ' will have a hoisted "me" field if it needs access to the containing
                ' display class, but the display class may not have a "me" field.
                If symbol.Type.IsClosureOrStateMachineType() AndAlso
                   GeneratedNameParser.GetKind(name) <> GeneratedNameKind.TransparentIdentifier Then

                    ReportMissingMe(syntax)
                End If

                Return If(TryCast(node, BoundParameter), New BoundParameter(syntax, symbol, node.IsLValue, node.SuppressVirtualCalls, symbol.Type))
            End If

            Dim result = variable.ToBoundExpression(syntax, node.IsLValue, node.SuppressVirtualCalls)
            Debug.Assert(TypeSymbol.Equals(result.Type, node.Type, TypeCompareKind.ConsiderEverything) OrElse
                         TypeSymbol.Equals(result.Type.BaseTypeNoUseSiteDiagnostics, node.Type, TypeCompareKind.ConsiderEverything))
            Return result
        End Function

        Private Sub ReportMissingMe(syntax As SyntaxNode)
            _diagnostics.Add(New VBDiagnostic(ErrorFactory.ErrorInfo(ERRID.ERR_UseOfKeywordNotInInstanceMethod1, syntax.ToString()), syntax.GetLocation()))
        End Sub

        Private Function GetVariable(name As String) As DisplayClassVariable
            Dim variable As DisplayClassVariable = Nothing
            _displayClassVariables.TryGetValue(name, variable)
            Return variable
        End Function

    End Class
End Namespace
