Implement module cancellation and stack probing instrumentations

src/Compilers/CSharp/Portable/BoundTree/BoundExpression.cs

@@ -507,14 +507,15 @@ public override Symbol ExpressionSymbol
         /// <summary>
         /// Build an object creation expression without performing any rewriting
         /// </summary>
-        internal BoundObjectCreationExpression UpdateArgumentsAndInitializer(
+        internal BoundObjectCreationExpression Update(
+            MethodSymbol constructor,
             ImmutableArray<BoundExpression> newArguments,
             ImmutableArray<RefKind> newRefKinds,
             BoundObjectInitializerExpressionBase? newInitializerExpression,
             TypeSymbol? changeTypeOpt = null)
         {
             return Update(
-                constructor: Constructor,
+                constructor: constructor,
                 arguments: newArguments,
                 argumentNamesOpt: default(ImmutableArray<string?>),
                 argumentRefKindsOpt: newRefKinds,

src/Compilers/CSharp/Portable/BoundTree/BoundNode.cs

@@ -525,24 +525,31 @@ private void CheckDeclared(LocalSymbol local)
 
             public override BoundNode? VisitBlock(BoundBlock node)
             {
-                if (node.Instrumentation != null)
+                var instrumentation = node.Instrumentation;
+                if (instrumentation != null)
                 {
-                    var added = DeclaredLocals.Add(node.Instrumentation.Local);
-                    Debug.Assert(added);
+                    foreach (var local in instrumentation.Locals)
+                    {
+                        var added = DeclaredLocals.Add(local);
+                        Debug.Assert(added);
+                    }
 
-                    _ = Visit(node.Instrumentation.Prologue);
+                    _ = Visit(instrumentation.Prologue);
                 }
 
                 AddAll(node.Locals);
                 base.VisitBlock(node);
                 RemoveAll(node.Locals);
 
-                if (node.Instrumentation != null)
+                if (instrumentation != null)
                 {
-                    _ = Visit(node.Instrumentation.Epilogue);
+                    _ = Visit(instrumentation.Epilogue);
 
-                    var removed = DeclaredLocals.Remove(node.Instrumentation.Local);
-                    Debug.Assert(removed);
+                    foreach (var local in instrumentation.Locals)
+                    {
+                        var removed = DeclaredLocals.Remove(local);
+                        Debug.Assert(removed);
+                    }
                 }
 
                 return null;

src/Compilers/CSharp/Portable/BoundTree/BoundNodes.xml

@@ -706,9 +706,15 @@
 
   <!-- Instrumentation info attached to BoundBlock -->
   <Node Name="BoundBlockInstrumentation" Base="BoundNode">
-    <Field Name="Local" Type="LocalSymbol" Null="disallow"/>
-    <Field Name="Prologue" Type="BoundStatement" Null="disallow"/>
-    <Field Name="Epilogue" Type="BoundStatement" Null="disallow"/>
+    <!-- Local variables with scope that cover the entire block including the instrumentation prologue and epilogue. -->
+    <Field Name="Locals" Type="OneOrMany&lt;LocalSymbol&gt;" Null="disallow"/>
+    <!-- Optional prologue emitted in front of any instructions of the block, so that it always executes. -->
+    <Field Name="Prologue" Type="BoundStatement?" Null="allow"/>
+    <!--
+      If specified, a try-finally block is emitted around the instrumented block. The epilogue is emitted to the finally clause,
+      so that it always executes when the execution leaves the instrumented block.
+     -->
+    <Field Name="Epilogue" Type="BoundStatement?" Null="allow"/>
   </Node>
 
   <!-- Represents the raw metadata RowId value for a method definition.
@@ -759,6 +765,17 @@
     <Field Name="Type" Type="TypeSymbol" Override="true" Null="disallow"/>
   </Node>
 
+  <!-- Call to PrivateImplementationDetails.ModuleCancellationToken.ThrowIfCancellationRequested().
+       PrivateImplementationDetails.ModuleCancellationToken has no language-level symbol. -->
+  <Node Name="BoundThrowIfModuleCancellationRequested" Base="BoundExpression">
+    <Field Name="Type" Type="TypeSymbol" Override="true" Null="disallow"/>
+  </Node>
+
+  <!-- Returns the value of PrivateImplementationDetails.ModuleCancellationToken. It has no language-level symbol. -->
+  <Node Name="ModuleCancellationTokenExpression" Base="BoundExpression">
+    <Field Name="Type" Type="TypeSymbol" Override="true" Null="disallow"/>
+  </Node>
+
   <!-- Represents the GUID that is the current module's MVID.
        Implemented as a reference to PrivateImplementationDetails.MVID, which has no language-level symbol. -->
   <Node Name="BoundModuleVersionId" Base="BoundExpression">

src/Compilers/CSharp/Portable/CodeGen/EmitExpression.cs

@@ -238,6 +238,16 @@ private void EmitExpressionCore(BoundExpression expression, bool used)
                     EmitModuleVersionIdStringLoad();
                     break;
 
+                case BoundKind.ThrowIfModuleCancellationRequested:
+                    Debug.Assert(!used);
+                    EmitThrowIfModuleCancellationRequested(expression.Syntax);
+                    break;
+
+                case BoundKind.ModuleCancellationTokenExpression:
+                    Debug.Assert(used);
+                    EmitModuleCancellationTokenLoad(expression.Syntax);
+                    break;
+
                 case BoundKind.InstrumentationPayloadRoot:
                     Debug.Assert(used);
                     EmitInstrumentationPayloadRootLoad((BoundInstrumentationPayloadRoot)expression);
@@ -3571,7 +3581,43 @@ private void EmitModuleVersionIdStore(BoundModuleVersionId node)
 
         private void EmitModuleVersionIdToken(BoundModuleVersionId node)
         {
-            _builder.EmitToken(_module.GetModuleVersionId(_module.Translate(node.Type, node.Syntax, _diagnostics.DiagnosticBag), node.Syntax, _diagnostics.DiagnosticBag), node.Syntax, _diagnostics.DiagnosticBag);
+            _builder.EmitToken(
+                _module.GetModuleVersionId(_module.Translate(node.Type, node.Syntax, _diagnostics.DiagnosticBag), node.Syntax, _diagnostics.DiagnosticBag),
+                node.Syntax,
+                _diagnostics.DiagnosticBag);
+        }
+
+        private void EmitThrowIfModuleCancellationRequested(SyntaxNode syntax)
+        {
+            var cancellationTokenType = _module.CommonCompilation.CommonGetWellKnownType(WellKnownType.System_Threading_CancellationToken);
+
+            _builder.EmitOpCode(ILOpCode.Ldsflda);
+            _builder.EmitToken(
+                _module.GetModuleCancellationToken(_module.Translate(cancellationTokenType, syntax, _diagnostics.DiagnosticBag), syntax, _diagnostics.DiagnosticBag),
+                syntax,
+                _diagnostics.DiagnosticBag);
+
+            var throwMethod = (MethodSymbol)_module.Compilation.GetWellKnownTypeMember(WellKnownMember.System_Threading_CancellationToken__ThrowIfCancellationRequested);
+
+            // BoundThrowIfModuleCancellationRequested should not be created if the method doesn't exist.
+            Debug.Assert(throwMethod != null);
+
+            _builder.EmitOpCode(ILOpCode.Call, -1);
+            _builder.EmitToken(
+                _module.Translate(throwMethod, syntax, _diagnostics.DiagnosticBag),
+                syntax,
+                _diagnostics.DiagnosticBag);
+        }
+
+        private void EmitModuleCancellationTokenLoad(SyntaxNode syntax)
+        {
+            var cancellationTokenType = _module.CommonCompilation.CommonGetWellKnownType(WellKnownType.System_Threading_CancellationToken);
+
+            _builder.EmitOpCode(ILOpCode.Ldsfld);
+            _builder.EmitToken(
+                _module.GetModuleCancellationToken(_module.Translate(cancellationTokenType, syntax, _diagnostics.DiagnosticBag), syntax, _diagnostics.DiagnosticBag),
+                syntax,
+                _diagnostics.DiagnosticBag);
         }
 
         private void EmitModuleVersionIdStringLoad()

src/Compilers/CSharp/Portable/CodeGen/EmitStatement.cs

@@ -677,38 +677,63 @@ private void EmitBlock(BoundBlock block)
 
         private void EmitInstrumentedBlock(BoundBlockInstrumentation instrumentation, BoundBlock block)
         {
-            _builder.OpenLocalScope();
-            DefineLocal(instrumentation.Local, block.Syntax);
-
-            if (_emitPdbSequencePoints)
+            if (!instrumentation.Locals.IsEmpty)
             {
-                EmitHiddenSequencePoint();
+                _builder.OpenLocalScope();
+
+                foreach (var local in instrumentation.Locals)
+                {
+                    DefineLocal(local, block.Syntax);
+                }
             }
 
-            EmitStatement(instrumentation.Prologue);
+            if (instrumentation.Prologue != null)
+            {
+                if (_emitPdbSequencePoints)
+                {
+                    EmitHiddenSequencePoint();
+                }
+
+                EmitStatement(instrumentation.Prologue);
+            }
 
             _builder.AssertStackEmpty();
 
-            _builder.OpenLocalScope(ScopeType.TryCatchFinally);
+            if (instrumentation.Epilogue != null)
+            {
+                _builder.OpenLocalScope(ScopeType.TryCatchFinally);
 
-            _builder.OpenLocalScope(ScopeType.Try);
-            EmitUninstrumentedBlock(block);
-            _builder.CloseLocalScope(); // try
+                _builder.OpenLocalScope(ScopeType.Try);
+
+                EmitUninstrumentedBlock(block);
+                _builder.CloseLocalScope(); // try
 
-            _builder.OpenLocalScope(ScopeType.Finally);
+                _builder.OpenLocalScope(ScopeType.Finally);
 
-            if (_emitPdbSequencePoints)
+                if (_emitPdbSequencePoints)
+                {
+                    EmitHiddenSequencePoint();
+                }
+
+                EmitStatement(instrumentation.Epilogue);
+                _builder.CloseLocalScope(); // finally
+
+                _builder.CloseLocalScope(); // try-finally
+            }
+            else
             {
-                EmitHiddenSequencePoint();
+                EmitUninstrumentedBlock(block);
             }
 
-            EmitStatement(instrumentation.Epilogue);
-            _builder.CloseLocalScope(); // finally
-
-            _builder.CloseLocalScope(); // try-finally
+            if (!instrumentation.Locals.IsEmpty)
+            {
+                foreach (var local in instrumentation.Locals)
+                {
+                    FreeLocal(local);
+                }
 
-            FreeLocal(instrumentation.Local);
-            _builder.CloseLocalScope();
+                _builder.CloseLocalScope();
+            }
         }
 
         private void EmitUninstrumentedBlock(BoundBlock block)

src/Compilers/CSharp/Portable/CodeGen/Optimizer.cs

@@ -616,12 +616,13 @@ public override BoundNode VisitBlock(BoundBlock node)
 
             if (node.Instrumentation != null)
             {
-                DeclareLocal(node.Instrumentation.Local, stack: 0);
+                foreach (var local in node.Instrumentation.Locals)
+                    DeclareLocal(local, stack: 0);
             }
 
             // normally we would not allow stack locals
             // when evaluation stack is not empty.
-            DeclareLocals(node.Locals, 0);
+            DeclareLocals(node.Locals, stack: 0);
 
             return base.VisitBlock(node);
         }

src/Compilers/CSharp/Portable/Compiler/MethodCompiler.cs

@@ -1018,7 +1018,9 @@ private void CompileMethod(
                             ((methodSymbol.ContainingType.IsStructType() && !methodSymbol.IsImplicitConstructor) ||
                             methodSymbol is SynthesizedPrimaryConstructor ||
                             instrumentation.Kinds.Contains(InstrumentationKind.TestCoverage) ||
-                            instrumentation.Kinds.Contains(InstrumentationKindExtensions.LocalStateTracing)))
+                            instrumentation.Kinds.Contains(InstrumentationKindExtensions.LocalStateTracing) ||
+                            instrumentation.Kinds.Contains(InstrumentationKind.StackOverflowProbing) ||
+                            instrumentation.Kinds.Contains(InstrumentationKind.ModuleCancellation)))
                         {
                             if (methodSymbol.IsImplicitConstructor &&
                                 (instrumentation.Kinds.Contains(InstrumentationKind.TestCoverage) || instrumentation.Kinds.Contains(InstrumentationKindExtensions.LocalStateTracing)))

src/Compilers/CSharp/Portable/FlowAnalysis/AbstractFlowPass.cs

@@ -1662,6 +1662,11 @@ public override BoundNode VisitInstrumentationPayloadRoot(BoundInstrumentationPa
             return null;
         }
 
+        public override BoundNode VisitThrowIfModuleCancellationRequested(BoundThrowIfModuleCancellationRequested node)
+        {
+            return null;
+        }
+
         public override BoundNode VisitSourceDocumentIndex(BoundSourceDocumentIndex node)
         {
             return null;

src/Compilers/CSharp/Portable/FlowAnalysis/DefiniteAssignment.cs

@@ -2075,13 +2075,18 @@ void assignPatternVariablesAndMarkReadFields(BoundPattern pattern, bool definite
                 }
             }
         }
-
-        public override BoundNode VisitBlock(BoundBlock node)
+#nullable enable
+        public override BoundNode? VisitBlock(BoundBlock node)
         {
-            if (node.Instrumentation != null)
+            var instrumentation = node.Instrumentation;
+            if (instrumentation != null)
             {
-                DeclareVariable(node.Instrumentation.Local);
-                Visit(node.Instrumentation.Prologue);
+                DeclareVariables(instrumentation.Locals);
+
+                if (instrumentation.Prologue != null)
+                {
+                    Visit(instrumentation.Prologue);
+                }
             }
 
             DeclareVariables(node.Locals);
@@ -2100,14 +2105,14 @@ public override BoundNode VisitBlock(BoundBlock node)
             ReportUnusedVariables(node.Locals);
             ReportUnusedVariables(node.LocalFunctions);
 
-            if (node.Instrumentation != null)
+            if (instrumentation?.Epilogue != null)
             {
-                Visit(node.Instrumentation.Epilogue);
+                Visit(instrumentation.Epilogue);
             }
 
             return null;
         }
-
+#nullable disable
         private void VisitStatementsWithLocalFunctions(BoundBlock block)
         {
             if (!TrackingRegions && !block.LocalFunctions.IsDefaultOrEmpty)
@@ -2243,6 +2248,14 @@ private void DeclareVariables(ImmutableArray<LocalSymbol> locals)
             }
         }
 
+        private void DeclareVariables(OneOrMany<LocalSymbol> locals)
+        {
+            foreach (var symbol in locals)
+            {
+                DeclareVariable(symbol);
+            }
+        }
+
         private void DeclareVariable(LocalSymbol symbol)
         {
             var initiallyAssigned =