' 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.Runtime.CompilerServices
Imports System.Runtime.InteropServices
Imports Microsoft.CodeAnalysis.PooledObjects
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
Imports Scanner = Microsoft.CodeAnalysis.VisualBasic.Syntax.InternalSyntax.Scanner

Namespace Microsoft.CodeAnalysis.VisualBasic

    Friend Module SyntaxNodeExtensions

        <Extension()>
        Public Function WithAnnotations(Of TNode As VisualBasicSyntaxNode)(node As TNode, ParamArray annotations As SyntaxAnnotation()) As TNode
            Return DirectCast(node.Green.SetAnnotations(annotations).CreateRed(), TNode)
        End Function

        ''' <summary>
        ''' Find enclosing WithStatement if it exists.
        ''' </summary>
        ''' <param name="node"></param>
        ''' <returns></returns>
        ''' <remarks></remarks>
        <Extension()> _
        Public Function ContainingWithStatement(node As VisualBasicSyntaxNode) As WithStatementSyntax
            Debug.Assert(node IsNot Nothing)

            If node Is Nothing Then
                Return Nothing
            End If

            node = node.Parent
            While node IsNot Nothing
                Select Case node.Kind
                    Case SyntaxKind.WithBlock
                        Return DirectCast(node, WithBlockSyntax).WithStatement

                    Case SyntaxKind.SubBlock,
                        SyntaxKind.FunctionBlock,
                        SyntaxKind.PropertyBlock,
                        SyntaxKind.ConstructorBlock,
                        SyntaxKind.OperatorBlock,
                        SyntaxKind.EventBlock
                        ' Don't look outside the current method/property/operator/event
                        Exit While

                End Select

                node = node.Parent
            End While

            Return Nothing
        End Function

        <Extension()> _
        Public Sub GetAncestors(Of T As VisualBasicSyntaxNode, C As VisualBasicSyntaxNode)(node As VisualBasicSyntaxNode, result As ArrayBuilder(Of T))

            Dim current = node.Parent
            Do While current IsNot Nothing AndAlso Not (TypeOf current Is C)
                If TypeOf current Is T Then
                    result.Add(DirectCast(current, T))
                End If
                current = current.Parent
            Loop

            result.ReverseContents()
        End Sub

        <Extension()> _
        Public Function GetAncestorOrSelf(Of T As VisualBasicSyntaxNode)(node As VisualBasicSyntaxNode) As T

            Do While node IsNot Nothing
                Dim result = TryCast(node, T)
                If result IsNot Nothing Then
                    Return result
                End If
                node = node.Parent
            Loop

            Return Nothing
        End Function

        <Extension()>
        Public Function IsLambdaExpressionSyntax(this As SyntaxNode) As Boolean
            Select Case this.Kind
                Case SyntaxKind.SingleLineFunctionLambdaExpression,
                     SyntaxKind.SingleLineSubLambdaExpression,
                     SyntaxKind.MultiLineFunctionLambdaExpression,
                     SyntaxKind.MultiLineSubLambdaExpression
                    Return True
            End Select

            Return False
        End Function

        ''' <summary>
        ''' Simplified version of ExtractAnonymousTypeMemberName implemented on inner tokens.
        ''' </summary>
        <Extension()>
        Friend Function ExtractAnonymousTypeMemberName(input As ExpressionSyntax, <Out()> ByRef failedToInferFromXmlName As XmlNameSyntax) As SyntaxToken
            ' TODO: revise and remove code duplication
            failedToInferFromXmlName = Nothing

TryAgain:
            Select Case input.Kind
                Case SyntaxKind.IdentifierName
                    Return DirectCast(input, IdentifierNameSyntax).Identifier

                Case SyntaxKind.XmlName
                    Dim xmlNameInferredFrom = DirectCast(input, XmlNameSyntax)
                    If Not Scanner.IsIdentifier(xmlNameInferredFrom.LocalName.ToString) Then
                        failedToInferFromXmlName = xmlNameInferredFrom
                        Return Nothing
                    End If

                    Return xmlNameInferredFrom.LocalName

                Case SyntaxKind.XmlBracketedName
                    ' handles something like <a-a>
                    Dim xmlNameInferredFrom = DirectCast(input, XmlBracketedNameSyntax)
                    input = xmlNameInferredFrom.Name
                    GoTo TryAgain

                Case SyntaxKind.SimpleMemberAccessExpression,
                     SyntaxKind.DictionaryAccessExpression

                    Dim memberAccess = DirectCast(input, MemberAccessExpressionSyntax)

                    If input.Kind = SyntaxKind.SimpleMemberAccessExpression Then
                        ' See if this is an identifier qualified with XmlElementAccessExpression or XmlDescendantAccessExpression
                        Dim receiver As ExpressionSyntax = If(memberAccess.Expression, GetCorrespondingConditionalAccessReceiver(memberAccess))

                        If receiver IsNot Nothing Then
                            Select Case receiver.Kind
                                Case SyntaxKind.XmlElementAccessExpression,
                                    SyntaxKind.XmlDescendantAccessExpression

                                    input = receiver
                                    GoTo TryAgain
                            End Select
                        End If
                    End If

                    input = memberAccess.Name
                    GoTo TryAgain

                Case SyntaxKind.XmlElementAccessExpression,
                     SyntaxKind.XmlAttributeAccessExpression,
                     SyntaxKind.XmlDescendantAccessExpression

                    Dim xmlAccess = DirectCast(input, XmlMemberAccessExpressionSyntax)

                    input = xmlAccess.Name
                    GoTo TryAgain

                Case SyntaxKind.InvocationExpression
                    Dim invocation = DirectCast(input, InvocationExpressionSyntax)

                    Dim target As ExpressionSyntax = If(invocation.Expression, GetCorrespondingConditionalAccessReceiver(invocation))

                    If target Is Nothing Then
                        Exit Select
                    End If

                    If invocation.ArgumentList Is Nothing OrElse invocation.ArgumentList.Arguments.Count = 0 Then
                        input = target
                        GoTo TryAgain
                    End If

                    Debug.Assert(invocation.ArgumentList IsNot Nothing)

                    If invocation.ArgumentList.Arguments.Count = 1 Then
                        ' See if this is an indexed XmlElementAccessExpression or XmlDescendantAccessExpression
                        Select Case target.Kind
                            Case SyntaxKind.XmlElementAccessExpression,
                                SyntaxKind.XmlDescendantAccessExpression
                                input = target
                                GoTo TryAgain
                        End Select
                    End If

                Case SyntaxKind.ConditionalAccessExpression
                    input = DirectCast(input, ConditionalAccessExpressionSyntax).WhenNotNull
                    GoTo TryAgain
            End Select

            Return Nothing
        End Function

        Private Function GetCorrespondingConditionalAccessReceiver(node As ExpressionSyntax) As ExpressionSyntax
            Dim access As ConditionalAccessExpressionSyntax = GetCorrespondingConditionalAccessExpression(node)

            If access IsNot Nothing Then
                Return access.Expression
            End If

            Return Nothing
        End Function

        <Extension>
        Friend Function GetCorrespondingConditionalAccessExpression(node As ExpressionSyntax) As ConditionalAccessExpressionSyntax
            Dim access As VisualBasicSyntaxNode = node
            Dim parent As VisualBasicSyntaxNode = access.Parent

            While parent IsNot Nothing
                Select Case parent.Kind
                    Case SyntaxKind.DictionaryAccessExpression,
                         SyntaxKind.SimpleMemberAccessExpression

                        If DirectCast(parent, MemberAccessExpressionSyntax).Expression IsNot access Then
                            Return Nothing
                        End If

                    Case SyntaxKind.XmlElementAccessExpression,
                         SyntaxKind.XmlDescendantAccessExpression,
                         SyntaxKind.XmlAttributeAccessExpression

                        If DirectCast(parent, XmlMemberAccessExpressionSyntax).Base IsNot access Then
                            Return Nothing
                        End If

                    Case SyntaxKind.InvocationExpression

                        If DirectCast(parent, InvocationExpressionSyntax).Expression IsNot access Then
                            Return Nothing
                        End If

                    Case SyntaxKind.ConditionalAccessExpression

                        Dim conditional = DirectCast(parent, ConditionalAccessExpressionSyntax)
                        If conditional.WhenNotNull Is access Then
                            Return conditional
                        ElseIf conditional.Expression IsNot access Then
                            Return Nothing
                        End If

                    Case Else
                        Return Nothing
                End Select

                access = parent
                parent = access.Parent
            End While

            Return Nothing
        End Function

        <Extension>
        Friend Function GetLeafAccess(conditionalAccess As ConditionalAccessExpressionSyntax) As ExpressionSyntax
            Dim access As ExpressionSyntax = conditionalAccess.WhenNotNull

            Do
                Select Case access.Kind
                    Case SyntaxKind.DictionaryAccessExpression,
                         SyntaxKind.SimpleMemberAccessExpression

                        Dim memberAccess = DirectCast(access, MemberAccessExpressionSyntax)
                        If memberAccess.Expression Is Nothing Then
                            Return memberAccess
                        Else
                            access = memberAccess.Expression
                        End If

                    Case SyntaxKind.XmlElementAccessExpression,
                         SyntaxKind.XmlDescendantAccessExpression,
                         SyntaxKind.XmlAttributeAccessExpression

                        Dim memberAccess = DirectCast(access, XmlMemberAccessExpressionSyntax)
                        If memberAccess.Base Is Nothing Then
                            Return memberAccess
                        Else
                            access = memberAccess.Base
                        End If

                    Case SyntaxKind.InvocationExpression

                        Dim invocation = DirectCast(access, InvocationExpressionSyntax)
                        If invocation.Expression Is Nothing Then
                            Return invocation
                        Else
                            access = invocation.Expression
                        End If

                    Case SyntaxKind.ConditionalAccessExpression

                        access = DirectCast(access, ConditionalAccessExpressionSyntax).Expression

                        If access Is Nothing Then
                            ' Must be a syntax error
                            Return Nothing
                        End If

                    Case Else
                        Return Nothing
                End Select
            Loop
        End Function

        ''' <summary>
        ''' Returns true if all arguments are of the specified kind and they are also missing.
        ''' </summary>
        <Extension()>
        Public Function AllAreMissing(arguments As IEnumerable(Of VisualBasicSyntaxNode), kind As SyntaxKind) As Boolean
            Return Not arguments.Any(Function(arg) Not (arg.Kind = kind AndAlso DirectCast(arg, IdentifierNameSyntax).IsMissing))
        End Function

        ''' <summary>
        ''' Returns true if all arguments are missing.
        ''' </summary>
        ''' <param name="arguments"></param>
        <Extension()>
        Public Function AllAreMissingIdentifierName(arguments As IEnumerable(Of VisualBasicSyntaxNode)) As Boolean
            Return arguments.AllAreMissing(SyntaxKind.IdentifierName)
        End Function

        ''' <summary>
        ''' Given a syntax node of query clause returns its leading keyword
        ''' </summary>
        <Extension()>
        Public Function QueryClauseKeywordOrRangeVariableIdentifier(syntax As SyntaxNode) As SyntaxToken
            Select Case syntax.Kind

                Case SyntaxKind.CollectionRangeVariable
                    Return DirectCast(syntax, CollectionRangeVariableSyntax).Identifier.Identifier

                Case SyntaxKind.ExpressionRangeVariable
                    Return DirectCast(syntax, ExpressionRangeVariableSyntax).NameEquals.Identifier.Identifier

                Case SyntaxKind.FromClause
                    Return DirectCast(syntax, FromClauseSyntax).FromKeyword

                Case SyntaxKind.FromClause
                    Return DirectCast(syntax, FromClauseSyntax).FromKeyword

                Case SyntaxKind.LetClause
                    Return DirectCast(syntax, LetClauseSyntax).LetKeyword

                Case SyntaxKind.AggregateClause
                    Return DirectCast(syntax, AggregateClauseSyntax).AggregateKeyword

                Case SyntaxKind.DistinctClause
                    Return DirectCast(syntax, DistinctClauseSyntax).DistinctKeyword

                Case SyntaxKind.WhereClause
                    Return DirectCast(syntax, WhereClauseSyntax).WhereKeyword

                Case SyntaxKind.SkipWhileClause, SyntaxKind.TakeWhileClause
                    Return DirectCast(syntax, PartitionWhileClauseSyntax).SkipOrTakeKeyword

                Case SyntaxKind.SkipClause, SyntaxKind.TakeClause
                    Return DirectCast(syntax, PartitionClauseSyntax).SkipOrTakeKeyword

                Case SyntaxKind.GroupByClause
                    Return DirectCast(syntax, GroupByClauseSyntax).GroupKeyword

                Case SyntaxKind.GroupJoinClause
                    Return DirectCast(syntax, GroupJoinClauseSyntax).GroupKeyword

                Case SyntaxKind.SimpleJoinClause
                    Return DirectCast(syntax, SimpleJoinClauseSyntax).JoinKeyword

                Case SyntaxKind.OrderByClause
                    Return DirectCast(syntax, OrderByClauseSyntax).OrderKeyword

                Case SyntaxKind.SelectClause
                    Return DirectCast(syntax, SelectClauseSyntax).SelectKeyword

                Case Else
                    Throw ExceptionUtilities.UnexpectedValue(syntax.Kind)
            End Select
        End Function

        <Extension>
        Friend Function EnclosingStructuredTrivia(node As VisualBasicSyntaxNode) As StructuredTriviaSyntax
            While node IsNot Nothing
                If node.IsStructuredTrivia Then
                    Return DirectCast(node, StructuredTriviaSyntax)
                Else
                    node = node.Parent
                End If
            End While
            Return Nothing
        End Function

    End Module
End Namespace