' 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 Microsoft.CodeAnalysis.PooledObjects
Imports Microsoft.CodeAnalysis.Syntax.InternalSyntax
Imports Microsoft.CodeAnalysis.Text
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
Imports ReferenceEqualityComparer = Roslyn.Utilities.ReferenceEqualityComparer
'-----------------------------------------------------------------------------
' Contains the definition of the DeclarationContext
'-----------------------------------------------------------------------------
 
Namespace Microsoft.CodeAnalysis.VisualBasic.Syntax.InternalSyntax
 
    Friend NotInheritable Class CompilationUnitContext
        Inherits NamespaceBlockContext
 
        Private _optionStmts As CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of OptionStatementSyntax)
        Private _importsStmts As CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of ImportsStatementSyntax)
        Private _attributeStmts As CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of AttributesStatementSyntax)
        Private _state As SyntaxKind
 
        Friend Sub New(parser As Parser)
            MyBase.New(SyntaxKind.CompilationUnit, Nothing, Nothing)
 
            Me.Parser = parser
            _statements = _parser._pool.Allocate(Of StatementSyntax)()
            _state = SyntaxKind.OptionStatement
        End Sub
 
        Friend Overrides ReadOnly Property IsWithinAsyncMethodOrLambda As Boolean
            Get
                Return Parser.IsScript
            End Get
        End Property
 
        Friend Overrides Function ProcessSyntax(node As VisualBasicSyntaxNode) As BlockContext
            Do
                Select Case _state
                    Case SyntaxKind.OptionStatement
                        If node.Kind = SyntaxKind.OptionStatement Then
                            Add(node)
                            Return Me
                        End If
                        _optionStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of OptionStatementSyntax)(Body.Node)
                        _state = SyntaxKind.ImportsStatement
 
                    Case SyntaxKind.ImportsStatement
                        If node.Kind = SyntaxKind.ImportsStatement Then
                            Add(node)
                            Return Me
                        End If
                        _importsStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of ImportsStatementSyntax)(Body.Node)
                        _state = SyntaxKind.AttributesStatement
 
                    Case SyntaxKind.AttributesStatement
                        If node.Kind = SyntaxKind.AttributesStatement Then
                            Add(node)
                            Return Me
                        End If
                        _attributeStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of AttributesStatementSyntax)(Body.Node)
                        _state = SyntaxKind.None
 
                    Case Else
                        ' only allow executable statements in top-level script code
                        If _parser.IsScript Then
                            Dim newContext = TryProcessExecutableStatement(node)
                            If newContext IsNot Nothing Then
                                Return newContext
                            End If
                        End If
 
                        Return MyBase.ProcessSyntax(node)
                End Select
            Loop
        End Function
 
        Friend Overrides Function TryLinkSyntax(node As VisualBasicSyntaxNode, ByRef newContext As BlockContext) As LinkResult
            If _parser.IsScript Then
                Return TryLinkStatement(node, newContext)
            Else
                Return MyBase.TryLinkSyntax(node, newContext)
            End If
        End Function
 
        Friend Overrides Function CreateBlockSyntax(endStmt As StatementSyntax) As VisualBasicSyntaxNode
            Throw ExceptionUtilities.Unreachable
        End Function
 
        Friend Function CreateCompilationUnit(optionalTerminator As PunctuationSyntax,
                                              notClosedIfDirectives As ArrayBuilder(Of IfDirectiveTriviaSyntax),
                                              notClosedRegionDirectives As ArrayBuilder(Of RegionDirectiveTriviaSyntax),
                                              haveRegionDirectives As Boolean,
                                              notClosedExternalSourceDirective As ExternalSourceDirectiveTriviaSyntax) As CompilationUnitSyntax
 
            Debug.Assert(optionalTerminator Is Nothing OrElse optionalTerminator.Kind = SyntaxKind.EndOfFileToken)
 
            If _state <> SyntaxKind.None Then
                Select Case _state
                    Case SyntaxKind.OptionStatement
                        _optionStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of OptionStatementSyntax)(Body.Node)
 
                    Case SyntaxKind.ImportsStatement
                        _importsStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of ImportsStatementSyntax)(Body.Node)
 
                    Case SyntaxKind.AttributesStatement
                        _attributeStmts = New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of AttributesStatementSyntax)(Body.Node)
                End Select
                _state = SyntaxKind.None
            End If
 
            Dim declarations = Body()
 
            Dim result = SyntaxFactory.CompilationUnit(_optionStmts,
                                                       _importsStmts,
                                                       _attributeStmts,
                                                       declarations,
                                                       optionalTerminator)
 
            Dim regionsAreAllowedEverywhere = Not haveRegionDirectives OrElse Parser.CheckFeatureAvailability(Feature.RegionsEverywhere)
 
            If notClosedIfDirectives IsNot Nothing OrElse notClosedRegionDirectives IsNot Nothing OrElse notClosedExternalSourceDirective IsNot Nothing OrElse
               Not regionsAreAllowedEverywhere Then
                result = DiagnosticRewriter.Rewrite(result, notClosedIfDirectives, notClosedRegionDirectives, regionsAreAllowedEverywhere, notClosedExternalSourceDirective, Parser)
 
                If notClosedIfDirectives IsNot Nothing Then
                    notClosedIfDirectives.Free()
                End If
                If notClosedRegionDirectives IsNot Nothing Then
                    notClosedRegionDirectives.Free()
                End If
            End If
 
            FreeStatements()
            Return result
        End Function
 
        Private Class DiagnosticRewriter
            Inherits VisualBasicSyntaxRewriter
 
            Private _notClosedIfDirectives As HashSet(Of IfDirectiveTriviaSyntax) = Nothing
            Private _notClosedRegionDirectives As HashSet(Of RegionDirectiveTriviaSyntax) = Nothing
            Private _notClosedExternalSourceDirective As ExternalSourceDirectiveTriviaSyntax = Nothing
            Private _regionsAreAllowedEverywhere As Boolean
 
            Private _parser As Parser
            Private _declarationBlocksBeingVisited As ArrayBuilder(Of VisualBasicSyntaxNode) ' CompilationUnitSyntax is treated as a declaration block for our purposes
            Private _parentsOfRegionDirectivesAwaitingClosure As ArrayBuilder(Of VisualBasicSyntaxNode) ' Nodes are coming from _declarationBlocksBeingVisited
            Private _tokenWithDirectivesBeingVisited As SyntaxToken
 
            Private Sub New()
                MyBase.New()
            End Sub
 
            Public Shared Function Rewrite(compilationUnit As CompilationUnitSyntax,
                                           notClosedIfDirectives As ArrayBuilder(Of IfDirectiveTriviaSyntax),
                                           notClosedRegionDirectives As ArrayBuilder(Of RegionDirectiveTriviaSyntax),
                                           regionsAreAllowedEverywhere As Boolean,
                                           notClosedExternalSourceDirective As ExternalSourceDirectiveTriviaSyntax,
                                           parser As Parser) As CompilationUnitSyntax
 
                Dim rewriter As New DiagnosticRewriter()
 
                If notClosedIfDirectives IsNot Nothing Then
                    rewriter._notClosedIfDirectives =
                        New HashSet(Of IfDirectiveTriviaSyntax)(ReferenceEqualityComparer.Instance)
 
                    For Each node In notClosedIfDirectives
                        rewriter._notClosedIfDirectives.Add(node)
                    Next
                End If
 
                If notClosedRegionDirectives IsNot Nothing Then
                    rewriter._notClosedRegionDirectives =
                        New HashSet(Of RegionDirectiveTriviaSyntax)(ReferenceEqualityComparer.Instance)
 
                    For Each node In notClosedRegionDirectives
                        rewriter._notClosedRegionDirectives.Add(node)
                    Next
                End If
 
                rewriter._parser = parser
                rewriter._regionsAreAllowedEverywhere = regionsAreAllowedEverywhere
 
                If Not regionsAreAllowedEverywhere Then
                    rewriter._declarationBlocksBeingVisited = ArrayBuilder(Of VisualBasicSyntaxNode).GetInstance()
                    rewriter._parentsOfRegionDirectivesAwaitingClosure = ArrayBuilder(Of VisualBasicSyntaxNode).GetInstance()
                End If
 
                rewriter._notClosedExternalSourceDirective = notClosedExternalSourceDirective
 
                Dim result = DirectCast(rewriter.Visit(compilationUnit), CompilationUnitSyntax)
 
                If Not regionsAreAllowedEverywhere Then
                    Debug.Assert(rewriter._declarationBlocksBeingVisited.Count = 0)
                    Debug.Assert(rewriter._parentsOfRegionDirectivesAwaitingClosure.Count = 0) ' We never add parents of not closed #Region directives into this stack.
                    rewriter._declarationBlocksBeingVisited.Free()
                    rewriter._parentsOfRegionDirectivesAwaitingClosure.Free()
                End If
 
                Return result
            End Function
 
#If DEBUG Then
            ' NOTE: the logic is heavily relying on the fact that green nodes in
            ' NOTE: one single tree are not reused, the following code assert this
            Private ReadOnly _processedNodesWithoutDuplication As HashSet(Of VisualBasicSyntaxNode) = New HashSet(Of VisualBasicSyntaxNode)(ReferenceEqualityComparer.Instance)
#End If
 
            Public Overrides Function VisitCompilationUnit(node As CompilationUnitSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitCompilationUnit(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitMethodBlock(node As MethodBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitMethodBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitConstructorBlock(node As ConstructorBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitConstructorBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitOperatorBlock(node As OperatorBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitOperatorBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitAccessorBlock(node As AccessorBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitAccessorBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitNamespaceBlock(node As NamespaceBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitNamespaceBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitClassBlock(node As ClassBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitClassBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitStructureBlock(node As StructureBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitStructureBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitModuleBlock(node As ModuleBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitModuleBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitInterfaceBlock(node As InterfaceBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitInterfaceBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitEnumBlock(node As EnumBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitEnumBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitPropertyBlock(node As PropertyBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitPropertyBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitEventBlock(node As EventBlockSyntax) As VisualBasicSyntaxNode
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    _declarationBlocksBeingVisited.Push(node)
                End If
 
                Dim result = MyBase.VisitEventBlock(node)
 
                If _declarationBlocksBeingVisited IsNot Nothing Then
                    Dim n = _declarationBlocksBeingVisited.Pop()
                    Debug.Assert(n Is node)
                End If
 
                Return result
            End Function
 
            Public Overrides Function VisitIfDirectiveTrivia(node As IfDirectiveTriviaSyntax) As VisualBasicSyntaxNode
#If DEBUG Then
                Debug.Assert(_processedNodesWithoutDuplication.Add(node))
#End If
                Dim rewritten = MyBase.VisitIfDirectiveTrivia(node)
                If Me._notClosedIfDirectives IsNot Nothing AndAlso Me._notClosedIfDirectives.Contains(node) Then
                    rewritten = Parser.ReportSyntaxError(rewritten, ERRID.ERR_LbExpectedEndIf)
                End If
 
                Return rewritten
            End Function
 
            Public Overrides Function VisitRegionDirectiveTrivia(node As RegionDirectiveTriviaSyntax) As VisualBasicSyntaxNode
#If DEBUG Then
                Debug.Assert(_processedNodesWithoutDuplication.Add(node))
#End If
                Dim rewritten = MyBase.VisitRegionDirectiveTrivia(node)
                If Me._notClosedRegionDirectives IsNot Nothing AndAlso Me._notClosedRegionDirectives.Contains(node) Then
                    rewritten = Parser.ReportSyntaxError(rewritten, ERRID.ERR_ExpectedEndRegion)
 
                ElseIf Not _regionsAreAllowedEverywhere
                    rewritten = VerifyRegionPlacement(node, rewritten)
                End If
 
                Return rewritten
            End Function
 
            Private Function VerifyRegionPlacement(original As VisualBasicSyntaxNode, rewritten As VisualBasicSyntaxNode) As VisualBasicSyntaxNode
                Dim containingBlock = _declarationBlocksBeingVisited.Peek()
 
                ' Ensure that the directive is inside the block, rather than is attached to it as a leading/trailing trivia
                Debug.Assert(_declarationBlocksBeingVisited.Count > 1 OrElse containingBlock.Kind = SyntaxKind.CompilationUnit)
 
                If _declarationBlocksBeingVisited.Count > 1 Then
 
                    If _tokenWithDirectivesBeingVisited Is containingBlock.GetFirstToken() Then
                        Dim leadingTrivia = _tokenWithDirectivesBeingVisited.GetLeadingTrivia()
 
                        If leadingTrivia IsNot Nothing AndAlso New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of VisualBasicSyntaxNode)(leadingTrivia).Nodes.Contains(original) Then
                            containingBlock = _declarationBlocksBeingVisited(_declarationBlocksBeingVisited.Count - 2)
                        End If
 
                    ElseIf _tokenWithDirectivesBeingVisited Is containingBlock.GetLastToken() Then
                        Dim trailingTrivia = _tokenWithDirectivesBeingVisited.GetTrailingTrivia()
 
                        If trailingTrivia IsNot Nothing AndAlso New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of VisualBasicSyntaxNode)(trailingTrivia).Nodes.Contains(original) Then
                            containingBlock = _declarationBlocksBeingVisited(_declarationBlocksBeingVisited.Count - 2)
                        End If
                    End If
                End If
 
                Dim reportAnError = Not IsValidContainingBlockForRegionInVB12(containingBlock)
 
                If original.Kind = SyntaxKind.RegionDirectiveTrivia Then
                    _parentsOfRegionDirectivesAwaitingClosure.Push(containingBlock)
                Else
                    Debug.Assert(original.Kind = SyntaxKind.EndRegionDirectiveTrivia)
 
                    If _parentsOfRegionDirectivesAwaitingClosure.Count > 0 Then
                        Dim regionBeginContainingBlock = _parentsOfRegionDirectivesAwaitingClosure.Pop()
 
                        If regionBeginContainingBlock IsNot containingBlock AndAlso IsValidContainingBlockForRegionInVB12(regionBeginContainingBlock) Then
                            reportAnError = True
                        End If
                    End If
                End If
 
                If reportAnError Then
                    rewritten = _parser.ReportFeatureUnavailable(Feature.RegionsEverywhere, rewritten)
                End If
 
                Return rewritten
            End Function
 
            Private Shared Function IsValidContainingBlockForRegionInVB12(containingBlock As VisualBasicSyntaxNode) As Boolean
                Select Case containingBlock.Kind
                    Case SyntaxKind.FunctionBlock,
                         SyntaxKind.SubBlock,
                         SyntaxKind.ConstructorBlock,
                         SyntaxKind.OperatorBlock,
                         SyntaxKind.SetAccessorBlock,
                         SyntaxKind.GetAccessorBlock,
                         SyntaxKind.AddHandlerAccessorBlock,
                         SyntaxKind.RemoveHandlerAccessorBlock,
                         SyntaxKind.RaiseEventAccessorBlock
 
                        Return False
                End Select
 
                Return True
            End Function
 
            Public Overrides Function VisitEndRegionDirectiveTrivia(node As EndRegionDirectiveTriviaSyntax) As VisualBasicSyntaxNode
#If DEBUG Then
                Debug.Assert(_processedNodesWithoutDuplication.Add(node))
#End If
                Dim rewritten = MyBase.VisitEndRegionDirectiveTrivia(node)
 
                If Not _regionsAreAllowedEverywhere Then
                    rewritten = VerifyRegionPlacement(node, rewritten)
                End If
 
                Return rewritten
            End Function
 
            Public Overrides Function VisitExternalSourceDirectiveTrivia(node As ExternalSourceDirectiveTriviaSyntax) As VisualBasicSyntaxNode
#If DEBUG Then
                Debug.Assert(_processedNodesWithoutDuplication.Add(node))
#End If
                Dim rewritten = MyBase.VisitExternalSourceDirectiveTrivia(node)
                If Me._notClosedExternalSourceDirective Is node Then
                    rewritten = Parser.ReportSyntaxError(rewritten, ERRID.ERR_ExpectedEndExternalSource)
                End If
                Return rewritten
            End Function
 
            Public Overrides Function Visit(node As VisualBasicSyntaxNode) As VisualBasicSyntaxNode
                If node Is Nothing OrElse Not node.ContainsDirectives Then
                    Return node
                End If
 
                Return node.Accept(Me)
            End Function
 
            Public Overrides Function VisitSyntaxToken(token As SyntaxToken) As SyntaxToken
                If token Is Nothing OrElse Not token.ContainsDirectives Then
                    Return token
                End If
 
                Debug.Assert(_tokenWithDirectivesBeingVisited Is Nothing)
                _tokenWithDirectivesBeingVisited = token
 
                Dim leadingTrivia = token.GetLeadingTrivia()
                Dim trailingTrivia = token.GetTrailingTrivia()
 
                If leadingTrivia IsNot Nothing Then
                    Dim rewritten = VisitList(New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of VisualBasicSyntaxNode)(leadingTrivia)).Node
                    If leadingTrivia IsNot rewritten Then
                        token = DirectCast(token.WithLeadingTrivia(rewritten), SyntaxToken)
                    End If
                End If
 
                If trailingTrivia IsNot Nothing Then
                    Dim rewritten = VisitList(New CodeAnalysis.Syntax.InternalSyntax.SyntaxList(Of VisualBasicSyntaxNode)(trailingTrivia)).Node
                    If trailingTrivia IsNot rewritten Then
                        token = DirectCast(token.WithTrailingTrivia(rewritten), SyntaxToken)
                    End If
                End If
 
                _tokenWithDirectivesBeingVisited = Nothing
                Return token
            End Function
 
        End Class
    End Class
 
End Namespace