Add support for await expressions

Change-Id: I36c1bad75a0f6981715e621c09df1befef38061a
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/106166
Commit-Queue: Brian Wilkerson <brianwilkerson@google.com>
Reviewed-by: Paul Berry <paulberry@google.com>
diff --git a/pkg/nnbd_migration/lib/src/graph_builder.dart b/pkg/nnbd_migration/lib/src/graph_builder.dart
index 1f32714..f87211d 100644
--- a/pkg/nnbd_migration/lib/src/graph_builder.dart
+++ b/pkg/nnbd_migration/lib/src/graph_builder.dart
@@ -180,7 +180,13 @@
 
   @override
   DecoratedType visitAwaitExpression(AwaitExpression node) {
-    throw new UnimplementedError('TODO(brianwilkerson)');
+    var expressionType = node.expression.accept(this);
+    // TODO(paulberry) Handle subclasses of Future.
+    if (expressionType.type.isDartAsyncFuture ||
+        expressionType.type.isDartAsyncFutureOr) {
+      expressionType = expressionType.typeArguments[0];
+    }
+    return expressionType;
   }
 
   @override
diff --git a/pkg/nnbd_migration/test/migration_visitor_test.dart b/pkg/nnbd_migration/test/migration_visitor_test.dart
index f60975d..ccc0c03 100644
--- a/pkg/nnbd_migration/test/migration_visitor_test.dart
+++ b/pkg/nnbd_migration/test/migration_visitor_test.dart
@@ -272,6 +272,40 @@
         checkExpression('c.s'), decoratedTypeAnnotation('C c').node);
   }
 
+  @failingTest
+  test_awaitExpression_future_nonNullable() async {
+    await analyze('''
+Future<void> f() async {
+  int x = await g();
+}
+Future<int> g() async => 3;
+''');
+
+    assertNoUpstreamNullability(decoratedTypeAnnotation('int').node);
+  }
+
+  @failingTest
+  test_awaitExpression_future_nullable() async {
+    await analyze('''
+Future<void> f() async {
+  int x = await g();
+}
+Future<int> g() async => null;
+''');
+
+    assertNoUpstreamNullability(decoratedTypeAnnotation('int').node);
+  }
+
+  test_awaitExpression_nonFuture() async {
+    await analyze('''
+Future<void> f() async {
+  int x = await 3;
+}
+''');
+
+    assertNoUpstreamNullability(decoratedTypeAnnotation('int').node);
+  }
+
   test_binaryExpression_add_left_check() async {
     await analyze('''
 int f(int i, int j) => i + j;