Refine hoisted local scopes in async methods

Old: Scopes followed the bound tree, so the outermost scope was within the
try-catch block synthesized in async MovedNext methods.

New: If the entire try block would be a hoisted local scope, we remove the
marker node (BoundStateMachineScope) and insert a new one around the
entire method body.

Upshot: hoisted locals are still in scope when stopped on method closing
braces.

src/Compilers/CSharp/Portable/Lowering/AsyncRewriter/AsyncMethodToStateMachineRewriter.cs

@@ -118,6 +118,9 @@ internal void GenerateMoveNext(BoundStatement body, MethodSymbol moveNextMethod)
 
             BoundStatement rewrittenBody = (BoundStatement)Visit(body);
 
+            ImmutableArray<StateMachineFieldSymbol> rootScopeHoistedLocals;
+            TryUnwrapBoundStateMachineScope(ref rewrittenBody, out rootScopeHoistedLocals);
+
             var bodyBuilder = ArrayBuilder<BoundStatement>.GetInstance();
 
             bodyBuilder.Add(F.HiddenSequencePoint());
@@ -187,18 +190,25 @@ internal void GenerateMoveNext(BoundStatement body, MethodSymbol moveNextMethod)
             bodyBuilder.Add(F.Label(_exitLabel));
             bodyBuilder.Add(F.Return());
 
-            var newBody = bodyBuilder.ToImmutableAndFree();
+            var newStatements = bodyBuilder.ToImmutableAndFree();
 
             var locals = ArrayBuilder<LocalSymbol>.GetInstance();
             locals.Add(cachedState);
             if ((object)_exprRetValue != null) locals.Add(_exprRetValue);
 
-            F.CloseMethod(
+            var newBody =
                 F.SequencePoint(
-                    body.Syntax,
+                    body.Syntax, 
                     F.Block(
-                        locals.ToImmutableAndFree(),
-                        newBody)));
+                        locals.ToImmutableAndFree(), 
+                        newStatements));
+
+            if (rootScopeHoistedLocals.Length > 0)
+            {
+                newBody = MakeStateMachineScope(rootScopeHoistedLocals, newBody);
+            }
+            
+            F.CloseMethod(newBody);
         }
 
         protected override BoundStatement GenerateReturn(bool finished)

src/Compilers/CSharp/Portable/Lowering/StateMachineRewriter/MethodToStateMachineRewriter.cs

@@ -322,7 +322,7 @@ private BoundStatement PossibleIteratorScope(ImmutableArray<LocalSymbol> locals,
             // wrap the node in an iterator scope for debugging
             if (hoistedLocalsWithDebugScopes.Count != 0)
             {
-                translatedStatement = F.Block(new BoundStateMachineScope(F.Syntax, hoistedLocalsWithDebugScopes.ToImmutable(), translatedStatement));
+                translatedStatement = MakeStateMachineScope(hoistedLocalsWithDebugScopes.ToImmutable(), translatedStatement);
             }
 
             hoistedLocalsWithDebugScopes.Free();
@@ -330,6 +330,36 @@ private BoundStatement PossibleIteratorScope(ImmutableArray<LocalSymbol> locals,
             return translatedStatement;
         }
 
+        /// <remarks>
+        /// Must remain in sync with <see cref="TryUnwrapBoundStateMachineScope"/>.
+        /// </remarks>
+        internal BoundBlock MakeStateMachineScope(ImmutableArray<StateMachineFieldSymbol> hoistedLocals, BoundStatement statement)
+        {
+            return F.Block(new BoundStateMachineScope(F.Syntax, hoistedLocals, statement));
+        }
+
+        /// <remarks>
+        /// Must remain in sync with <see cref="MakeStateMachineScope"/>.
+        /// </remarks>
+        internal bool TryUnwrapBoundStateMachineScope(ref BoundStatement statement, out ImmutableArray<StateMachineFieldSymbol> hoistedLocals)
+        {
+            if (statement.Kind == BoundKind.Block)
+            {
+                var rewrittenBlock = (BoundBlock)statement;
+                var rewrittenStatements = rewrittenBlock.Statements;
+                if (rewrittenStatements.Length == 1 && rewrittenStatements[0].Kind == BoundKind.StateMachineScope)
+                {
+                    var stateMachineScope = (BoundStateMachineScope)rewrittenStatements[0];
+                    statement = stateMachineScope.Statement;
+                    hoistedLocals = stateMachineScope.Fields;
+                    return true;
+                }
+            }
+
+            hoistedLocals = ImmutableArray<StateMachineFieldSymbol>.Empty;
+            return false;
+        }
+
         private void AddVariableCleanup(ArrayBuilder<BoundAssignmentOperator> cleanup, FieldSymbol field)
         {
             if (MightContainReferences(field.Type))

src/Compilers/CSharp/Test/Emit/PDB/PDBAsyncTests.cs

@@ -462,7 +462,7 @@ private Task<int> GetNextInt(Random random)
       <customDebugInfo>
         <forward declaringType=""ConsoleApplication1.Program"" methodName=""Main"" parameterNames=""args"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0x11"" endOffset=""0x11e"" />
+          <slot startOffset=""0x0"" endOffset=""0x14e"" />
           <slot startOffset=""0x0"" endOffset=""0x0"" />
           <slot startOffset=""0x0"" endOffset=""0x0"" />
           <slot startOffset=""0x41"" endOffset=""0xed"" />
@@ -848,7 +848,7 @@ static async Task M(int b)
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0x11"" endOffset=""0xe0"" />
+          <slot startOffset=""0x0"" endOffset=""0x10d"" />
         </hoistedLocalScopes>
         <encLocalSlotMap>
           <slot kind=""27"" offset=""0"" />
@@ -945,7 +945,7 @@ static async Task M(int b)
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0xa"" endOffset=""0xc0"" />
+          <slot startOffset=""0x0"" endOffset=""0xeb"" />
         </hoistedLocalScopes>
       </customDebugInfo>
       <sequencePoints>
@@ -1034,7 +1034,7 @@ static async Task M(int b)
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0x11"" endOffset=""0xcf"" />
+          <slot startOffset=""0x0"" endOffset=""0xfc"" />
         </hoistedLocalScopes>
         <encLocalSlotMap>
           <slot kind=""27"" offset=""0"" />
@@ -1127,7 +1127,7 @@ static async Task M()
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0xe"" endOffset=""0xdc"" />
+          <slot startOffset=""0x0"" endOffset=""0x109"" />
         </hoistedLocalScopes>
         <encLocalSlotMap>
           <slot kind=""27"" offset=""0"" />
@@ -1290,7 +1290,7 @@ static async Task M()
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0x11"" endOffset=""0xdc"" />
+          <slot startOffset=""0x0"" endOffset=""0x109"" />
         </hoistedLocalScopes>
         <encLocalSlotMap>
           <slot kind=""27"" offset=""0"" />
@@ -1587,7 +1587,7 @@ .maxstack  3
       <customDebugInfo>
         <forward declaringType=""C"" methodName=""F"" />
         <hoistedLocalScopes>
-          <slot startOffset=""0xe"" endOffset=""0xed"" />
+          <slot startOffset=""0x0"" endOffset=""0x11d"" />
           <slot startOffset=""0x29"" endOffset=""0x32"" />
         </hoistedLocalScopes>
         <encLocalSlotMap>

src/Compilers/VisualBasic/Portable/Lowering/AsyncRewriter/AsyncRewriter.AsyncMethodToClassRewriter.vb

@@ -117,6 +117,9 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
 
                 Dim rewrittenBody As BoundStatement = DirectCast(Visit(body), BoundStatement)
 
+                Dim rootScopeHoistedLocals As ImmutableArray(Of FieldSymbol) = Nothing
+                TryUnwrapBoundStateMachineScope(rewrittenBody, rootScopeHoistedLocals)
+
                 Dim bodyBuilder = ArrayBuilder(Of BoundStatement).GetInstance()
 
                 ' NOTE: We don't need to create/use any label after Dispatch inside Try block below 
@@ -201,14 +204,18 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
                 bodyBuilder.Add(Me.F.Label(Me._exitLabel))
                 bodyBuilder.Add(Me.F.Return())
 
-                Dim newBody As ImmutableArray(Of BoundStatement) = bodyBuilder.ToImmutableAndFree()
-
-                Me._owner.CloseMethod(
-                    Me.F.Block(
+                Dim newStatements As ImmutableArray(Of BoundStatement) = bodyBuilder.ToImmutableAndFree()
+                Dim newBody = Me.F.Block(
                         If(Me._exprRetValue IsNot Nothing,
                            ImmutableArray.Create(Me._exprRetValue, Me.CachedState),
                            ImmutableArray.Create(Me.CachedState)),
-                       newBody))
+                       newStatements)
+
+                If rootScopeHoistedLocals.Length > 0 Then
+                    newBody = MakeStateMachineScope(rootScopeHoistedLocals, newBody)
+                End If
+
+                Me._owner.CloseMethod(newBody)
             End Sub
 
             Protected Overrides ReadOnly Property ResumeLabelName As String

src/Compilers/VisualBasic/Portable/Lowering/StateMachineRewriter/StateMachineRewriter.StateMachineMethodToClassRewriter.vb

 ' Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 Imports System.Collections.Immutable
+Imports System.Runtime.InteropServices
 Imports Microsoft.CodeAnalysis.CodeGen
 Imports Microsoft.CodeAnalysis.VisualBasic.Symbols
 
@@ -220,9 +221,7 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
                 ' Wrap the node in a StateMachineScope for debugging
                 Dim translatedStatement As BoundNode = wrapped
                 If hoistedLocalsWithDebugScopes.Count > 0 Then
-                    translatedStatement = Me.F.Block(New BoundStateMachineScope(Me.F.Syntax,
-                                                             hoistedLocalsWithDebugScopes.ToImmutable(),
-                                                             DirectCast(translatedStatement, BoundStatement)).MakeCompilerGenerated)
+                    translatedStatement = MakeStateMachineScope(hoistedLocalsWithDebugScopes.ToImmutable(), DirectCast(translatedStatement, BoundStatement))
                 End If
 
                 hoistedLocalsWithDebugScopes.Free()
@@ -230,6 +229,32 @@ Namespace Microsoft.CodeAnalysis.VisualBasic
                 Return translatedStatement
             End Function
 
+            ''' <remarks>
+            ''' Must remain in sync with <see cref="TryUnwrapBoundStateMachineScope"/>.
+            ''' </remarks>
+            Friend Function MakeStateMachineScope(hoistedLocals As ImmutableArray(Of FieldSymbol), statement As BoundStatement) As BoundBlock
+                Return Me.F.Block(New BoundStateMachineScope(Me.F.Syntax, hoistedLocals, statement).MakeCompilerGenerated)
+            End Function
+
+            ''' <remarks>
+            ''' Must remain in sync with <see cref="MakeStateMachineScope"/>.
+            ''' </remarks>
+            Friend Function TryUnwrapBoundStateMachineScope(ByRef statement As BoundStatement, <Out> ByRef hoistedLocals As ImmutableArray(Of FieldSymbol)) As Boolean
+                If statement.Kind = BoundKind.Block Then
+                    Dim rewrittenBlock = DirectCast(statement, BoundBlock)
+                    Dim rewrittenStatements = rewrittenBlock.Statements
+                    If rewrittenStatements.Length = 1 AndAlso rewrittenStatements(0).Kind = BoundKind.StateMachineScope Then
+                        Dim stateMachineScope = DirectCast(rewrittenStatements(0), BoundStateMachineScope)
+                        statement = stateMachineScope.Statement
+                        hoistedLocals = stateMachineScope.Fields
+                        Return True
+                    End If
+                End If
+
+                hoistedLocals = ImmutableArray(Of FieldSymbol).Empty
+                Return False
+            End Function
+
             Private Function NeedsProxy(localOrParameter As Symbol) As Boolean
                 Debug.Assert(localOrParameter.Kind = SymbolKind.Local OrElse localOrParameter.Kind = SymbolKind.Parameter)
                 Return _hoistedVariables.Contains(localOrParameter)

src/Compilers/VisualBasic/Test/Emit/PDB/PDBAsyncTests.vb

@@ -441,10 +441,8 @@ End Class
             </sequencePoints>
             <scope startOffset="0x0" endOffset="0x13a">
                 <importsforward declaringType="C+_Closure$__1-0" methodName="_Lambda$__0"/>
-                <scope startOffset="0x11" endOffset="0xfd">
-                    <local name="$VB$ResumableLocal_$VB$Closure_$0" il_index="0" il_start="0x11" il_end="0xfd" attributes="0"/>
-                    <local name="$VB$ResumableLocal_a$1" il_index="1" il_start="0x11" il_end="0xfd" attributes="0"/>
-                </scope>
+                <local name="$VB$ResumableLocal_$VB$Closure_$0" il_index="0" il_start="0x0" il_end="0x13a" attributes="0"/>
+                <local name="$VB$ResumableLocal_a$1" il_index="1" il_start="0x0" il_end="0x13a" attributes="0"/>
             </scope>
             <asyncInfo>
                 <kickoffMethod declaringType="C" methodName="Async_Lambda"/>
@@ -506,9 +504,7 @@ End Class
             </sequencePoints>
             <scope startOffset="0x0" endOffset="0x10f">
                 <importsforward declaringType="C+_Closure$__1-0" methodName="_Lambda$__0"/>
-                <scope startOffset="0xa" endOffset="0xd6">
-                    <local name="$VB$ResumableLocal_$VB$Closure_$0" il_index="0" il_start="0xa" il_end="0xd6" attributes="0"/>
-                </scope>
+                <local name="$VB$ResumableLocal_$VB$Closure_$0" il_index="0" il_start="0x0" il_end="0x10f" attributes="0"/>
             </scope>
             <asyncInfo>
                 <kickoffMethod declaringType="C" methodName="Async_Lambda"/>
@@ -579,10 +575,8 @@ End Class
                 <namespace name="System" importlevel="file"/>
                 <namespace name="System.Threading.Tasks" importlevel="file"/>
                 <currentnamespace name=""/>
-                <scope startOffset="0xe" endOffset="0xba">
-                    <local name="$VB$ResumableLocal_x$0" il_index="0" il_start="0xe" il_end="0xba" attributes="0"/>
-                    <local name="$VB$ResumableLocal_y$1" il_index="1" il_start="0xe" il_end="0xba" attributes="0"/>
-                </scope>
+                <local name="$VB$ResumableLocal_x$0" il_index="0" il_start="0x0" il_end="0xf7" attributes="0"/>
+                <local name="$VB$ResumableLocal_y$1" il_index="1" il_start="0x0" il_end="0xf7" attributes="0"/>
             </scope>
             <asyncInfo>
                 <kickoffMethod declaringType="C" methodName="Async_NoLambda"/>
@@ -643,10 +637,8 @@ End Class
                 <namespace name="System" importlevel="file"/>
                 <namespace name="System.Threading.Tasks" importlevel="file"/>
                 <currentnamespace name=""/>
-                <scope startOffset="0xa" endOffset="0xaa">
-                    <local name="$VB$ResumableLocal_x$0" il_index="0" il_start="0xa" il_end="0xaa" attributes="0"/>
-                    <local name="$VB$ResumableLocal_y$1" il_index="1" il_start="0xa" il_end="0xaa" attributes="0"/>
-                </scope>
+                <local name="$VB$ResumableLocal_x$0" il_index="0" il_start="0x0" il_end="0xe3" attributes="0"/>
+                <local name="$VB$ResumableLocal_y$1" il_index="1" il_start="0x0" il_end="0xe3" attributes="0"/>
             </scope>
             <asyncInfo>
                 <kickoffMethod declaringType="C" methodName="Async_NoLambda"/>

src/ExpressionEvaluator/CSharp/Test/ExpressionCompiler/LocalsTests.cs

@@ -2749,6 +2749,34 @@ .maxstack  1
             locals.Free();
         }
 
+        [WorkItem(2336, "https://github.com/dotnet/roslyn/issues/2336")]
+        [Fact]
+        public void LocalsOnAsyncMethodClosingBrace()
+        {
+            var source =
+@"using System;
+using System.Threading.Tasks;
+class C
+{
+    async void M()
+    {
+        string s = null;
+#line 999
+    }
+}";
+            var compilation = CreateCompilationWithMscorlib45(source, options: TestOptions.DebugDll);
+            var runtime = CreateRuntimeInstance(compilation);
+            var context = CreateMethodContext(runtime, "C.<M>d__0.MoveNext()", atLineNumber: 999);
+            var locals = ArrayBuilder<LocalAndMethod>.GetInstance();
+            string typeName;
+            var testData = new CompilationTestData();
+            context.CompileGetLocals(locals, argumentsOnly: false, typeName: out typeName, testData: testData);
+            Assert.Equal(locals.Count, 2);
+            VerifyLocal(testData, "<>x", locals[0], "<>m0", "this");
+            VerifyLocal(testData, "<>x", locals[1], "<>m1", "s");
+            locals.Free();
+        }
+
         [WorkItem(1139013, "DevDiv")]
         [Fact]
         public void TransparentIdentifiers_FromParameter()

src/ExpressionEvaluator/VisualBasic/Test/ExpressionCompiler/LocalsTests.vb

@@ -2729,6 +2729,37 @@ End Class"
             locals.Free()
         End Sub
 
+        <WorkItem(2336, "https://github.com/dotnet/roslyn/issues/2336")>
+        <Fact>
+        Public Sub LocalsOnAsyncEndSub()
+            Const source = "
+Imports System
+Imports System.Threading.Tasks
+
+Class C
+    Async Sub M()
+        Dim s As String = Nothing
+#ExternalSource(""test"", 999)
+    End Sub
+#End ExternalSource
+End Class
+"
+            Dim comp = CreateCompilationWithReferences(
+                MakeSources(source),
+                {MscorlibRef_v4_0_30316_17626, MsvbRef_v4_0_30319_17929, SystemCoreRef_v4_0_30319_17929},
+                TestOptions.DebugDll)
+            Dim runtime = CreateRuntimeInstance(comp)
+            Dim context = CreateMethodContext(runtime, methodName:="C.VB$StateMachine_1_M.MoveNext", atLineNumber:=999)
+            Dim testData As New CompilationTestData()
+            Dim locals = ArrayBuilder(Of LocalAndMethod).GetInstance()
+            Dim typeName As String = Nothing
+            context.CompileGetLocals(locals, argumentsOnly:=False, typeName:=typeName, testData:=testData)
+            Assert.Equal(2, locals.Count)
+            VerifyLocal(testData, typeName, locals(0), "<>m0", "Me")
+            VerifyLocal(testData, typeName, locals(1), "<>m1", "s")
+            locals.Free()
+        End Sub
+
         <WorkItem(1139013, "DevDiv")>
         <Fact>
         Public Sub TransparentIdentifiers_FromParameter()