Implement custom awaitable support

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

@@ -2587,4 +2587,11 @@
     <Field Name="InitializerExpression" Type="BoundObjectInitializerExpressionBase" />
   </Node>
 
+  <!-- Node only created during some lowering stage as a possible side effect in a BoundSequence. It effectively represents `if (condition) expr`, as an expression -->
+  <Node Name="BoundLoweredConditionalSideEffect" Base="BoundExpression">
+    <Field Name="Type" Type="TypeSymbol?" Override="true" Null="always" />
+    <Field Name="Condition" Type="BoundExpression" Null="disallow" />
+    <Field Name="SideEffect" Type="BoundExpression" Null="disallow" />
+  </Node>
+
 </Tree>

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

@@ -166,6 +166,9 @@ private LocalDefinition EmitAddress(BoundExpression expression, AddressKind addr
                     EmitExpression(expression, used: true);
                     return null;
 
+                case BoundKind.LoweredConditionalSideEffect:
+                    throw ExceptionUtilities.UnexpectedValue(expression.Kind);
+
                 default:
                     Debug.Assert(!HasHome(expression, addressKind));
                     return EmitAddressOfTempClone(expression);

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

@@ -296,6 +296,11 @@ private void EmitExpressionCore(BoundExpression expression, bool used)
                     }
                     break;
 
+                case BoundKind.LoweredConditionalSideEffect:
+                    Debug.Assert(!used);
+                    EmitLoweredConditionalSideEffect((BoundLoweredConditionalSideEffect)expression);
+                    break;
+
                 case BoundKind.ConditionalOperator:
                     EmitConditionalOperator((BoundConditionalOperator)expression, used);
                     break;
@@ -3849,6 +3854,23 @@ private void EmitConditionalOperator(BoundConditionalOperator expr, bool used)
             _builder.MarkLabel(doneLabel);
         }
 
+        /// <summary>
+        /// Emit code for a conditional side effect
+        /// </summary>
+        /// <remarks>
+        /// if (expr) sideeffect becomes
+        ///   if !expr goto AFTER
+        ///   sideeffect
+        ///   AFTER:
+        /// </remarks>
+        private void EmitLoweredConditionalSideEffect(BoundLoweredConditionalSideEffect boundLoweredIfSideEffect)
+        {
+            object afterLabel = new object();
+            EmitCondBranch(boundLoweredIfSideEffect.Condition, ref afterLabel, sense: false);
+            EmitExpression(boundLoweredIfSideEffect.SideEffect, used: false);
+            _builder.MarkLabel(afterLabel);
+        }
+
         /// <summary>
         /// Emit code for a null-coalescing operator.
         /// </summary>

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

@@ -1613,6 +1613,19 @@ public override BoundNode VisitLoweredConditionalAccess(BoundLoweredConditionalA
             return node.Update(receiver, node.HasValueMethodOpt, whenNotNull, whenNull, node.Id, node.ForceCopyOfNullableValueType, node.Type);
         }
 
+        public override BoundNode VisitLoweredConditionalSideEffect(BoundLoweredConditionalSideEffect node)
+        {
+            var origStack = StackDepth();
+            var condition = (BoundExpression)this.Visit(node.Condition);
+
+            var cookie = GetStackStateCookie();     // implicit branch here
+            SetStackDepth(origStack);  // side effect is evaluated with original stack
+            var sideEffect = (BoundExpression)this.Visit(node.SideEffect);
+            EnsureStackState(cookie);   // implicit label here
+
+            return node.Update(condition, sideEffect);
+        }
+
         public override BoundNode VisitComplexConditionalReceiver(BoundComplexConditionalReceiver node)
         {
             EnsureOnlyEvalStack();

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

@@ -2,7 +2,9 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
 
 namespace Microsoft.CodeAnalysis.CSharp;
@@ -26,11 +28,13 @@ public static BoundStatement Rewrite(
 
     private readonly CSharpCompilation _compilation;
     private readonly SyntheticBoundNodeFactory _factory;
+    private readonly Dictionary<BoundAwaitableValuePlaceholder, BoundExpression> _placeholderMap;
 
     private RuntimeAsyncRewriter(CSharpCompilation compilation, SyntheticBoundNodeFactory factory)
     {
         _compilation = compilation;
         _factory = factory;
+        _placeholderMap = [];
     }
 
     private NamedTypeSymbol Task
@@ -53,10 +57,11 @@ private NamedTypeSymbol ValueTaskT
         get => field ??= _compilation.GetWellKnownType(WellKnownType.System_Threading_Tasks_ValueTask_T);
     } = null!;
 
-    public BoundExpression VisitExpression(BoundExpression node)
+    [return: NotNullIfNotNull(nameof(node))]
+    public BoundExpression? VisitExpression(BoundExpression? node)
     {
         var result = Visit(node);
-        return (BoundExpression)result;
+        return (BoundExpression?)result;
     }
 
     public override BoundNode? VisitAwaitExpression(BoundAwaitExpression node)
@@ -88,8 +93,7 @@ public BoundExpression VisitExpression(BoundExpression node)
         }
         else
         {
-            // PROTOTYPE: when it's not a method with Task/TaskT/ValueTask/ValueTaskT returns, use the helpers
-            return base.VisitAwaitExpression(node);
+            return RewriteCustomAwaiterAwait(node);
         }
 
         // PROTOTYPE: Make sure that we report an error in initial binding if these are missing
@@ -112,4 +116,82 @@ public BoundExpression VisitExpression(BoundExpression node)
         // System.Runtime.CompilerServices.RuntimeHelpers.Await(awaitedExpression)
         return _factory.Call(receiver: null, awaitMethod, VisitExpression(node.Expression));
     }
+
+    private BoundExpression RewriteCustomAwaiterAwait(BoundAwaitExpression node)
+    {
+        // await expr
+        // becomes
+        // var _tmp = expr.GetAwaiter();
+        // if (!_tmp.IsCompleted)
+        //    UnsafeAwaitAwaiterFromRuntimeAsync(_tmp) OR AwaitAwaiterFromRuntimeAsync(_tmp);
+        // _tmp.GetResult();
+
+        // PROTOTYPE: await dynamic will need runtime checks, see AsyncMethodToStateMachine.GenerateAwaitOnCompletedDynamic
+
+        var expr = VisitExpression(node.Expression);
+
+        var awaitablePlaceholder = node.AwaitableInfo.AwaitableInstancePlaceholder;
+        if (awaitablePlaceholder is not null)
+        {
+            _placeholderMap.Add(awaitablePlaceholder, expr);
+        }
+
+        // expr.GetAwaiter()
+        var getAwaiter = VisitExpression(node.AwaitableInfo.GetAwaiter);
+        Debug.Assert(getAwaiter is not null);
+
+        if (awaitablePlaceholder is not null)
+        {
+            _placeholderMap.Remove(awaitablePlaceholder);
+        }
+
+        // var _tmp = expr.GetAwaiter();
+        var tmp = _factory.StoreToTemp(getAwaiter, out BoundAssignmentOperator store, kind: SynthesizedLocalKind.Awaiter);
+
+        // _tmp.IsCompleted
+        var isCompletedMethod = node.AwaitableInfo.IsCompleted!.GetMethod;
+        Debug.Assert(isCompletedMethod is not null);
+        var isCompletedCall = _factory.Call(tmp, isCompletedMethod);
+
+        // UnsafeAwaitAwaiterFromRuntimeAsync(_tmp) OR AwaitAwaiterFromRuntimeAsync(_tmp)
+        var discardedUseSiteInfo = CompoundUseSiteInfo<AssemblySymbol>.Discarded;
+        var useUnsafeAwait = _factory.Compilation.Conversions.ClassifyImplicitConversionFromType(
+            tmp.Type,
+            _factory.Compilation.GetWellKnownType(WellKnownType.System_Runtime_CompilerServices_ICriticalNotifyCompletion),
+            ref discardedUseSiteInfo).IsImplicit;
+
+        // PROTOTYPE: Make sure that we report an error in initial binding if these are missing
+        var awaitMethod = (MethodSymbol?)_compilation.GetWellKnownTypeMember(useUnsafeAwait
+            ? WellKnownMember.System_Runtime_CompilerServices_RuntimeHelpers__UnsafeAwaitAwaiterFromRuntimeAsync_TAwaiter
+            : WellKnownMember.System_Runtime_CompilerServices_RuntimeHelpers__AwaitAwaiterFromRuntimeAsync_TAwaiter);
+
+        Debug.Assert(awaitMethod is { Arity: 1 });
+
+        var awaitCall = _factory.Call(
+            receiver: null,
+            awaitMethod.Construct(tmp.Type),
+            tmp);
+
+        // if (!_tmp.IsCompleted) awaitCall
+        var ifNotCompleted = new BoundLoweredConditionalSideEffect(
+            node.Syntax,
+            condition: _factory.Not(isCompletedCall),
+            sideEffect: awaitCall);
+
+        // _tmp.GetResult()
+        var getResultMethod = node.AwaitableInfo.GetResult;
+        Debug.Assert(getResultMethod is not null);
+        var getResultCall = _factory.Call(tmp, getResultMethod);
+
+        // final sequence
+        return _factory.Sequence(
+            locals: [tmp.LocalSymbol],
+            sideEffects: [store, ifNotCompleted],
+            result: getResultCall);
+    }
+
+    public override BoundNode VisitAwaitableValuePlaceholder(BoundAwaitableValuePlaceholder node)
+    {
+        return _placeholderMap[node];
+    }
 }

src/Compilers/CSharp/Portable/Lowering/SpillSequenceSpiller.cs

@@ -1349,6 +1349,27 @@ public override BoundNode VisitLoweredConditionalAccess(BoundLoweredConditionalA
             }
         }
 
+        public override BoundNode VisitLoweredConditionalSideEffect(BoundLoweredConditionalSideEffect node)
+        {
+            // PROTOTYPE: No current path will actually hit this. Should we leave a Debug.Fail here so
+            // that when this is hit in the future, we know to add proper testing? Or just throw Unreachable?
+            BoundSpillSequenceBuilder conditionBuilder = null;
+            var condition = VisitExpression(ref conditionBuilder, node.Condition);
+
+            BoundSpillSequenceBuilder sideEffectBuilder = null;
+            var sideEffect = VisitExpression(ref sideEffectBuilder, node.SideEffect);
+
+            if (sideEffectBuilder == null)
+            {
+                return UpdateExpression(conditionBuilder, node.Update(condition, sideEffect));
+            }
+
+            conditionBuilder ??= new BoundSpillSequenceBuilder(sideEffectBuilder.Syntax);
+            conditionBuilder.AddStatement(_F.If(condition, UpdateStatement(sideEffectBuilder, _F.ExpressionStatement(sideEffect))));
+
+            return conditionBuilder.Update(_F.Default(node.Type));
+        }
+
         private sealed class ConditionalReceiverReplacer : BoundTreeRewriterWithStackGuardWithoutRecursionOnTheLeftOfBinaryOperator
         {
             private readonly BoundExpression _receiver;