' 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.Generic
Imports System.Threading
Imports Microsoft.Cci
Imports Microsoft.CodeAnalysis
Imports Microsoft.CodeAnalysis.Collections
Imports Microsoft.CodeAnalysis.Text
Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
Imports Microsoft.CodeAnalysis.VisualBasic.Syntax
 
Namespace Microsoft.CodeAnalysis.VisualBasic
 
    Partial Friend NotInheritable Class AsyncRewriter
        Inherits StateMachineRewriter(Of CapturedSymbolOrExpression)
 
        Partial Friend Class AsyncMethodToClassRewriter
            Inherits StateMachineMethodToClassRewriter
 
            Public Overrides Function VisitReturnStatement(node As BoundReturnStatement) As BoundNode
                Dim rewritten = DirectCast(MyBase.VisitReturnStatement(node), BoundReturnStatement)
                Dim expression As BoundExpression = rewritten.ExpressionOpt
 
                If expression IsNot Nothing Then
                    Debug.Assert(Me._asyncMethodKind = AsyncMethodKind.GenericTaskFunction)
 
                    If expression.Kind = BoundKind.SpillSequence Then
                        Dim spill = DirectCast(expression, BoundSpillSequence)
                        Debug.Assert(spill.ValueOpt IsNot Nothing)
                        Return Me.F.Block(
                                    RewriteSpillSequenceIntoBlock(
                                        spill,
                                        False,
                                        Me.F.Assignment(Me.F.Local(Me._exprRetValue, True), spill.ValueOpt)),
                                    Me.F.Goto(Me._exprReturnLabel))
                    Else
                        Return Me.F.Block(
                                    Me.F.Assignment(
                                        Me.F.Local(Me._exprRetValue, True), expression),
                                    Me.F.Goto(Me._exprReturnLabel))
                    End If
                End If
 
                Return F.Goto(Me._exprReturnLabel)
            End Function
 
            Public Overrides Function VisitExpressionStatement(node As BoundExpressionStatement) As BoundNode
                Dim rewritten = DirectCast(MyBase.VisitExpressionStatement(node), BoundExpressionStatement)
                Dim expression As BoundExpression = rewritten.Expression
 
                If expression.Kind <> BoundKind.SpillSequence Then
                    Return rewritten
                End If
 
                Return RewriteSpillSequenceIntoBlock(DirectCast(expression, BoundSpillSequence), True)
            End Function
 
            Public Overrides Function VisitThrowStatement(node As BoundThrowStatement) As BoundNode
                Dim rewritten = DirectCast(MyBase.VisitThrowStatement(node), BoundThrowStatement)
                Dim expression As BoundExpression = rewritten.ExpressionOpt
 
                If expression Is Nothing OrElse expression.Kind <> BoundKind.SpillSequence Then
                    Return rewritten
                End If
 
                Debug.Assert(expression.Kind = BoundKind.SpillSequence)
                Dim spill = DirectCast(expression, BoundSpillSequence)
                Debug.Assert(spill.ValueOpt IsNot Nothing)
 
                Return RewriteSpillSequenceIntoBlock(spill,
                                                     False,
                                                     rewritten.Update(spill.ValueOpt))
            End Function
 
            Public Overrides Function VisitConditionalGoto(node As BoundConditionalGoto) As BoundNode
                Dim rewritten = DirectCast(MyBase.VisitConditionalGoto(node), BoundConditionalGoto)
                Dim condition As BoundExpression = rewritten.Condition
 
                If condition.Kind <> BoundKind.SpillSequence Then
                    Return rewritten
                End If
 
                Dim spill = DirectCast(condition, BoundSpillSequence)
                Debug.Assert(spill.ValueOpt IsNot Nothing)
 
                Return RewriteSpillSequenceIntoBlock(spill,
                                                     False,
                                                     node.Update(spill.ValueOpt,
                                                                 node.JumpIfTrue,
                                                                 node.Label))
            End Function
 
        End Class
    End Class
 
End Namespace