Add support for nullable with runtime async (#82516)

The nullable walker was missing analysis of the runtime async helper,
which caused several warnings to be missing. This adds support for those
missing warnings.

Relates to test plan #75960

Author: 333fred

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

@@ -224,6 +224,7 @@ protected BoundNode VisitExpressionOrPatternWithStackGuard(ref int recursionDept
             return result;
         }
 
+        [DebuggerStepThrough]
         protected virtual void EnsureSufficientExecutionStack(int recursionDepth)
         {
             StackGuard.EnsureSufficientExecutionStack(recursionDepth);

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

@@ -374,6 +374,7 @@ protected override BoundNode VisitExpressionOrPatternWithoutStackGuard(BoundNode
             return base.Visit(node);
         }
 
+        [DebuggerStepThrough]
         protected override bool ConvertInsufficientExecutionStackExceptionToCancelledByStackGuardException()
         {
             return false; // just let the original exception bubble up.

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

@@ -529,6 +529,7 @@ public string GetDebuggerDisplay()
         // For purpose of nullability analysis, awaits create pending branches, so async usings and foreachs do too
         public sealed override bool AwaitUsingAndForeachAddsPendingBranch => true;
 
+        [DebuggerStepThrough]
         protected override void EnsureSufficientExecutionStack(int recursionDepth)
         {
             if (recursionDepth > StackGuard.MaxUncheckedRecursionDepth &&
@@ -542,6 +543,7 @@ protected override void EnsureSufficientExecutionStack(int recursionDepth)
             base.EnsureSufficientExecutionStack(recursionDepth);
         }
 
+        [DebuggerStepThrough]
         protected override bool ConvertInsufficientExecutionStackExceptionToCancelledByStackGuardException()
         {
             return true;
@@ -595,7 +597,6 @@ private static void AssertPlaceholderAllowedWithoutRegistration(BoundValuePlaceh
                 case BoundKind.InterpolatedStringHandlerPlaceholder:
                 case BoundKind.InterpolatedStringArgumentPlaceholder:
                 case BoundKind.ObjectOrCollectionValuePlaceholder:
-                case BoundKind.AwaitableValuePlaceholder:
                     return;
 
                 case BoundKind.ImplicitIndexerValuePlaceholder:
@@ -3598,16 +3599,43 @@ private void VisitLocalFunctionUse(LocalFunctionSymbol symbol)
         public override BoundNode? VisitUsingStatement(BoundUsingStatement node)
         {
             DeclareLocals(node.Locals);
-            Visit(node.AwaitOpt);
+            VisitAwaitableInfoForUsing(node.AwaitOpt, node.PatternDisposeInfoOpt?.Method);
             return base.VisitUsingStatement(node);
         }
 
         public override BoundNode? VisitUsingLocalDeclarations(BoundUsingLocalDeclarations node)
         {
-            Visit(node.AwaitOpt);
+            VisitAwaitableInfoForUsing(node.AwaitOpt, node.PatternDisposeInfoOpt?.Method);
             return base.VisitUsingLocalDeclarations(node);
         }
 
+        private void VisitAwaitableInfoForUsing(BoundAwaitableInfo? awaitInfo, MethodSymbol? patternDisposeMethod)
+        {
+            // Placeholder can be null in error scenarios. In such cases, nothing is filled in and errors would
+            // have been reported already. We can just return.
+            if (awaitInfo is not { AwaitableInstancePlaceholder: { } placeholder })
+            {
+                return;
+            }
+
+            VisitResult placeholderResult;
+            if (patternDisposeMethod is not null)
+            {
+                placeholderResult = new VisitResult(GetReturnTypeWithState(patternDisposeMethod), patternDisposeMethod.ReturnTypeWithAnnotations);
+            }
+            else
+            {
+                // IAsyncDisposable.DisposeAsync returns ValueTask, which is a struct; we know it can never be `ValueTask?`, as that is a different
+                // type that would not match the descriptor. Any other scenario either has an error already, or will report an
+                // error during emit when we fail to find IAsyncDisposable.DisposeAsync.
+                placeholderResult = new VisitResult(placeholder.Type, NullableAnnotation.NotAnnotated, NullableFlowState.NotNull);
+            }
+
+            AddPlaceholderReplacement(placeholder, placeholder, placeholderResult);
+            Visit(awaitInfo);
+            RemovePlaceholderReplacement(placeholder);
+        }
+
         public override BoundNode? VisitFixedStatement(BoundFixedStatement node)
         {
             DeclareLocals(node.Locals);
@@ -12069,23 +12097,13 @@ private void VisitForEachExpression(
                 // Analyze `await DisposeAsync()`
                 if (enumeratorInfoOpt is { NeedsDisposal: true, DisposeAwaitableInfo: BoundAwaitableInfo awaitDisposalInfo })
                 {
-                    var disposalPlaceholder = awaitDisposalInfo.AwaitableInstancePlaceholder;
-                    bool addedPlaceholder = false;
-                    if (enumeratorInfoOpt.PatternDisposeInfo is { Method: var originalDisposeMethod }) // no statically known Dispose method if doing a runtime check
+                    var patternDisposeMethod = enumeratorInfoOpt.PatternDisposeInfo?.Method;
+                    if (patternDisposeMethod is not null)
                     {
-                        Debug.Assert(disposalPlaceholder is not null);
-                        var disposeAsyncMethod = (MethodSymbol)AsMemberOfType(reinferredGetEnumeratorMethod.ReturnType, originalDisposeMethod);
-                        var result = new VisitResult(GetReturnTypeWithState(disposeAsyncMethod), disposeAsyncMethod.ReturnTypeWithAnnotations);
-                        AddPlaceholderReplacement(disposalPlaceholder, disposalPlaceholder, result);
-                        addedPlaceholder = true;
+                        patternDisposeMethod = (MethodSymbol)AsMemberOfType(reinferredGetEnumeratorMethod.ReturnType, patternDisposeMethod);
                     }
 
-                    Visit(awaitDisposalInfo);
-
-                    if (addedPlaceholder)
-                    {
-                        RemovePlaceholderReplacement(disposalPlaceholder!);
-                    }
+                    VisitAwaitableInfoForUsing(awaitDisposalInfo, patternDisposeMethod);
                 }
             }
 
@@ -12571,7 +12589,12 @@ protected override void AfterLeftChildOfBinaryLogicalOperatorHasBeenVisited(Boun
             Visit(awaitableInfo);
             RemovePlaceholderReplacement(placeholder);
 
-            if (node.Type.IsValueType || node.HasErrors || awaitableInfo.GetResult is null)
+            if (awaitableInfo is { GetResult: null, RuntimeAsyncAwaitCall: not null })
+            {
+                // This is AsyncHelpers.Await. We can directly use `_visitResult`
+                SetResult(node, _visitResult, updateAnalyzedNullability: true, isLvalue: false);
+            }
+            else if (node.Type.IsValueType || node.HasErrors || awaitableInfo.GetResult is null)
             {
                 SetNotNullResult(node);
             }
@@ -13212,8 +13235,6 @@ private void VisitThrow(BoundExpression? expr)
 
         public override BoundNode? VisitAwaitableValuePlaceholder(BoundAwaitableValuePlaceholder node)
         {
-            // These placeholders don't always follow proper placeholder discipline yet
-            AssertPlaceholderAllowedWithoutRegistration(node);
             VisitPlaceholderWithReplacement(node);
             return null;
         }
@@ -13235,8 +13256,55 @@ private void VisitPlaceholderWithReplacement(BoundValuePlaceholderBase node)
 
         public override BoundNode? VisitAwaitableInfo(BoundAwaitableInfo node)
         {
-            Visit(node.AwaitableInstancePlaceholder);
-            Visit(node.GetAwaiter);
+            if (node is { GetAwaiter: null, GetResult: null, RuntimeAsyncAwaitCall: null, RuntimeAsyncAwaitCallPlaceholder: null })
+            {
+                if (node.AwaitableInstancePlaceholder is not null)
+                {
+                    // Ensure that the placeholder has an analyzed nullability for the rewriter
+                    SetNotNullResult(node.AwaitableInstancePlaceholder);
+                }
+
+                SetInvalidResult();
+
+                // Error scenario
+                return null;
+            }
+
+            Debug.Assert(node.AwaitableInstancePlaceholder is not null);
+
+            if (node.RuntimeAsyncAwaitCall is not null)
+            {
+                Debug.Assert(node.RuntimeAsyncAwaitCallPlaceholder is not null);
+
+                if (node.GetAwaiter is null)
+                {
+                    // This is the simple case of just `AsyncHelpers.Await(expr)`. AwaitableInstancePlaceholder and RuntimeAsyncAwaitCallPlaceholder
+                    // refer to the same value.
+                    Debug.Assert(node.GetResult is null);
+                    VisitPlaceholderWithReplacement(node.AwaitableInstancePlaceholder);
+                    AddPlaceholderReplacement(node.RuntimeAsyncAwaitCallPlaceholder, node.AwaitableInstancePlaceholder, _visitResult);
+                    Visit(node.RuntimeAsyncAwaitCall);
+                    RemovePlaceholderReplacement(node.RuntimeAsyncAwaitCallPlaceholder);
+                }
+                else
+                {
+                    // More complicated case of `AsyncHelpers.AwaitAwaiter/UnsafeAwaitAwaiter`. We need to visit GetAwaiter first to get the nullability
+                    // result for the awaiter, which we then save. It is used both here, as the state of the RuntimeAsyncAwaitCallPlaceholder replacement,
+                    // and also for the caller to be able to re-infer GetResult with nullability.
+                    Visit(node.GetAwaiter);
+                    var getAwaiterResult = _visitResult;
+                    AddPlaceholderReplacement(node.RuntimeAsyncAwaitCallPlaceholder, node.GetAwaiter, getAwaiterResult);
+                    Visit(node.RuntimeAsyncAwaitCall);
+                    RemovePlaceholderReplacement(node.RuntimeAsyncAwaitCallPlaceholder);
+                    _visitResult = getAwaiterResult;
+                }
+            }
+            else
+            {
+                // Not runtime async
+                Visit(node.GetAwaiter);
+            }
+
             return null;
         }
 

src/Compilers/CSharp/Test/Emit/CodeGen/CodeGenAsyncTests.cs

@@ -10883,5 +10883,188 @@ .maxstack  1
 }
 """);
         }
+
+        [Fact]
+        public void RuntimeAsyncNullable_WithTaskReturningMethod()
+        {
+            var source = """
+using System.Threading.Tasks;
+
+#nullable enable
+
+await M(null);
+
+async Task? M(object o) { }
+""";
+
+            var comp = CreateRuntimeAsyncCompilation(source);
+            var verifier = CompileAndVerify(comp, verify: Verification.Fails with
+            {
+                ILVerifyMessage = $"""
+                    {ReturnValueMissing("<Main>$", "0xb")}
+                    {ReturnValueMissing("<<Main>$>g__M|0_0", "0x0")}
+                    """
+            });
+            verifier.VerifyDiagnostics(
+                // (5,7): warning CS8604: Possible null reference argument for parameter 'task' in 'void AsyncHelpers.Await(Task task)'.
+                // await M(null);
+                Diagnostic(ErrorCode.WRN_NullReferenceArgument, "M(null)").WithArguments("task", "void AsyncHelpers.Await(Task task)").WithLocation(5, 7),
+                // (5,9): warning CS8625: Cannot convert null literal to non-nullable reference type.
+                // await M(null);
+                Diagnostic(ErrorCode.WRN_NullAsNonNullable, "null").WithLocation(5, 9)
+            );
+        }
+
+        [Fact]
+        public void RuntimeAsyncNullable_WithTaskLikeReturningMethod()
+        {
+            var corlib = """
+                namespace System
+                {
+                    public class Attribute {}
+                    public enum AttributeTargets {}
+                    public class AttributeUsageAttribute : Attribute
+                    {
+                        public AttributeUsageAttribute(AttributeTargets validOn) {}
+                        public bool AllowMultiple { get; set; }
+                        public bool Inherited { get; set; }
+                    }
+                    public struct Boolean {}
+                    public struct Byte {}
+                    public abstract class Enum {}
+                    public class Exception {}
+                    public struct Int32 {}
+                    public class Object {}
+                    public class String {}
+                    public class ValueType {}
+                    public class Void {}
+
+                    namespace Threading.Tasks
+                    {
+                        public class Task
+                        {
+                            public Runtime.CompilerServices.TaskAwaiter GetAwaiter() => throw null!;
+                            public static Task CompletedTask => throw null!;
+                        }
+                    }
+
+                    namespace Runtime.CompilerServices
+                    {
+                        public interface INotifyCompletion {}
+                        public interface ICriticalNotifyCompletion : INotifyCompletion {}
+                        public class TaskAwaiter : ICriticalNotifyCompletion
+                        {
+                            public bool IsCompleted => false;
+                            public void GetResult() {}
+                        }
+
+                        public static class AsyncHelpers
+                        {
+                            public static void UnsafeAwaitAwaiter<TAwaiter>(TAwaiter awaiter) where TAwaiter : notnull, INotifyCompletion {}
+                        }
+                    }
+                }
+                """;
+
+            var source = """
+#nullable enable
+
+await M(null);
+
+MyTask M(object o) => throw null!;
+
+public class MyTask
+{
+    public System.Runtime.CompilerServices.TaskAwaiter? GetAwaiter() => throw null!;
+}
+""";
+
+            var comp = CreateEmptyCompilation([source, corlib], parseOptions: WithRuntimeAsync(TestOptions.RegularPreview), assemblyName: "TestAssembly");
+            var verifier = CompileAndVerify(comp, verify: Verification.Skipped);
+            verifier.VerifyDiagnostics(
+                // warning CS8021: No value for RuntimeMetadataVersion found. No assembly containing System.Object was found nor was a value for RuntimeMetadataVersion specified through options.
+                Diagnostic(ErrorCode.WRN_NoRuntimeMetadataVersion).WithLocation(1, 1),
+                // (3,7): warning CS8714: The type 'System.Runtime.CompilerServices.TaskAwaiter?' cannot be used as type parameter 'TAwaiter' in the generic type or method 'AsyncHelpers.UnsafeAwaitAwaiter<TAwaiter>(TAwaiter)'. Nullability of type argument 'System.Runtime.CompilerServices.TaskAwaiter?' doesn't match 'notnull' constraint.
+                // await M(null);
+                Diagnostic(ErrorCode.WRN_NullabilityMismatchInTypeParameterNotNullConstraint, "M").WithArguments("System.Runtime.CompilerServices.AsyncHelpers.UnsafeAwaitAwaiter<TAwaiter>(TAwaiter)", "TAwaiter", "System.Runtime.CompilerServices.TaskAwaiter?").WithLocation(3, 7),
+                // (3,9): warning CS8625: Cannot convert null literal to non-nullable reference type.
+                // await M(null);
+                Diagnostic(ErrorCode.WRN_NullAsNonNullable, "null").WithLocation(3, 9)
+            );
+        }
+
+        [Theory]
+        [InlineData("Task")]
+        [InlineData("ValueTask")]
+        public void RuntimeAsyncAwaitNullableFlow_WithTaskReturningMethod(string taskType)
+        {
+            var source = $$"""
+                using System.Threading.Tasks;
+
+                #nullable enable
+
+                _ = (await GetNullableResultAsync()).ToString();
+                _ = (await GetNullableResultAsync().ConfigureAwait(true)).ToString();
+                var x = await GetNullableResultAsync();
+                _ = x.ToString(); // 1
+                x = await GetNullableResultAsync().ConfigureAwait(false);
+                _ = x.ToString(); // 2
+
+                {{taskType}}<string?> GetNullableResultAsync() => default!;
+                """;
+
+            var comp = CreateRuntimeAsyncCompilation(source);
+            comp.VerifyDiagnostics(
+                // (5,6): warning CS8602: Dereference of a possibly null reference.
+                // _ = (await GetNullableResultAsync()).ToString();
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "await GetNullableResultAsync()").WithLocation(5, 6),
+                // (6,6): warning CS8602: Dereference of a possibly null reference.
+                // _ = (await GetNullableResultAsync().ConfigureAwait(true)).ToString();
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "await GetNullableResultAsync().ConfigureAwait(true)").WithLocation(6, 6),
+                // (8,5): warning CS8602: Dereference of a possibly null reference.
+                // _ = x.ToString(); // 1
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "x").WithLocation(8, 5),
+                // (10,5): warning CS8602: Dereference of a possibly null reference.
+                // _ = x.ToString(); // 2
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "x").WithLocation(10, 5)
+            );
+        }
+
+        [Fact]
+        public void RuntimeAsyncAwaitNullableFlow_WithTaskLikeReturningMethod()
+        {
+            var source = """
+                #nullable enable
+
+                _ = (await GetNullableResultTaskLike()).ToString();
+                var x = await GetNullableResultTaskLike();
+                _ = x.ToString();
+
+                MyTask GetNullableResultTaskLike() => throw null!;
+
+                public class MyTask
+                {
+                    public MyAwaiter GetAwaiter() => throw null!;
+                }
+
+                public class MyAwaiter : System.Runtime.CompilerServices.ICriticalNotifyCompletion
+                {
+                    public bool IsCompleted => false;
+                    public string? GetResult() => null;
+                    public void OnCompleted(System.Action continuation) { }
+                    public void UnsafeOnCompleted(System.Action continuation) { }
+                }
+                """;
+
+            var comp = CreateRuntimeAsyncCompilation(source);
+            comp.VerifyDiagnostics(
+                // (3,6): warning CS8602: Dereference of a possibly null reference.
+                // _ = (await GetNullableResultTaskLike()).ToString();
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "await GetNullableResultTaskLike()").WithLocation(3, 6),
+                // (5,5): warning CS8602: Dereference of a possibly null reference.
+                // _ = x.ToString();
+                Diagnostic(ErrorCode.WRN_NullReferenceReceiver, "x").WithLocation(5, 5)
+            );
+        }
     }
 }