Revert "Don't zone-register async callbacks for every await call in the VM."

runtime/vm/ast_transformer.cc

Index: runtime/vm/ast_transformer.cc
diff --git a/runtime/vm/ast_transformer.cc b/runtime/vm/ast_transformer.cc
index d8c01fded358f9c9046dae5c9e18e3cb4bff37f9..ac3f9f1b9132178de3c66fbec754f6bed1718397 100644
--- a/runtime/vm/ast_transformer.cc
+++ b/runtime/vm/ast_transformer.cc
@@ -112,18 +112,20 @@ void AwaitTransformer::VisitAwaitNode(AwaitNode* node) {
   // Await transformation:
   //
   //   :await_temp_var_X = <expr>;
+  //   :result_param = :await_temp_var_X;
+  //   if (:result_param is !Future) {
+  //     :result_param = Future.value(:result_param);
+  //   }
   //   AwaitMarker(kNewContinuationState);
-  //   :result_param = _awaitHelper(
-  //      :await_temp_var_X, :async_then_callback, :async_catch_error_callback);
+  //   :result_param = :result_param.then(:async_op);
+  //   _asyncCatchHelper(:result_param.catchError, :async_op);
   //   return;  // (return_type() == kContinuationTarget)
   //
   //   :saved_try_ctx_var = :await_saved_try_ctx_var_y;
   //   :await_temp_var_(X+1) = :result_param;
 
-  LocalVariable* async_then_callback = GetVariableInScope(
-      preamble_->scope(), Symbols::AsyncThenCallback());
-  LocalVariable* async_catch_error_callback = GetVariableInScope(
-      preamble_->scope(), Symbols::AsyncCatchErrorCallback());
+  LocalVariable* async_op = GetVariableInScope(
+      preamble_->scope(), Symbols::AsyncOperation());
   LocalVariable* result_param = GetVariableInScope(
       preamble_->scope(), Symbols::AsyncOperationParam());
   LocalVariable* error_param = GetVariableInScope(
@@ -132,34 +134,78 @@ void AwaitTransformer::VisitAwaitNode(AwaitNode* node) {
       preamble_->scope(), Symbols::AsyncOperationStackTraceParam());
 
   AstNode* transformed_expr = Transform(node->expr());
-  LocalVariable* await_temp = AddToPreambleNewTempVar(transformed_expr);
+  preamble_->Add(new(Z) StoreLocalNode(
+      Scanner::kNoSourcePos, result_param, transformed_expr));
+
+  LoadLocalNode* load_result_param = new(Z) LoadLocalNode(
+      Scanner::kNoSourcePos, result_param);
+
+  const Class& future_cls =
+      Class::ZoneHandle(Z, thread()->isolate()->object_store()->future_class());
+  ASSERT(!future_cls.IsNull());
+  const AbstractType& future_type =
+      AbstractType::ZoneHandle(Z, future_cls.RareType());
+  ASSERT(!future_type.IsNull());
+
+  LocalScope* is_not_future_scope = ChainNewScope(preamble_->scope());
+  SequenceNode* is_not_future_branch =
+      new (Z) SequenceNode(Scanner::kNoSourcePos, is_not_future_scope);
+
+  // if (:result_param is !Future) {
+  //   :result_param = Future.value(:result_param);
+  // }
+  const Function& value_ctor = Function::ZoneHandle(
+      Z, future_cls.LookupFunction(Symbols::FutureValue()));
+  ASSERT(!value_ctor.IsNull());
+  ArgumentListNode* ctor_args = new (Z) ArgumentListNode(Scanner::kNoSourcePos);
+  ctor_args->Add(new (Z) LoadLocalNode(Scanner::kNoSourcePos, result_param));
+  ConstructorCallNode* ctor_call =
+      new (Z) ConstructorCallNode(Scanner::kNoSourcePos,
+                                  TypeArguments::ZoneHandle(Z),
+                                  value_ctor,
+                                  ctor_args);
+  is_not_future_branch->Add(new (Z) StoreLocalNode(
+      Scanner::kNoSourcePos, result_param, ctor_call));
+  AstNode* is_not_future_test = new (Z) ComparisonNode(
+      Scanner::kNoSourcePos,
+      Token::kISNOT,
+      load_result_param,
+      new (Z) TypeNode(Scanner::kNoSourcePos, future_type));
+  preamble_->Add(new(Z) IfNode(Scanner::kNoSourcePos,
+                               is_not_future_test,
+                               is_not_future_branch,
+                               NULL));
 
   AwaitMarkerNode* await_marker = new (Z) AwaitMarkerNode();
   await_marker->set_scope(preamble_->scope());
   preamble_->Add(await_marker);
+  ArgumentListNode* args = new(Z) ArgumentListNode(Scanner::kNoSourcePos);
 
-  // :result_param = _awaitHelper(
-  //      :await_temp, :async_then_callback, :async_catch_error_callback)
-  const Library& async_lib = Library::Handle(Library::AsyncLibrary());
-  const Function& async_await_helper = Function::ZoneHandle(
-      Z, async_lib.LookupFunctionAllowPrivate(Symbols::AsyncAwaitHelper()));
-  ASSERT(!async_await_helper.IsNull());
-  ArgumentListNode* async_await_helper_args = new (Z) ArgumentListNode(
+  args->Add(new(Z) LoadLocalNode(Scanner::kNoSourcePos, async_op));
+  preamble_->Add(new (Z) StoreLocalNode(
+      Scanner::kNoSourcePos,
+      result_param,
+      new(Z) InstanceCallNode(node->token_pos(),
+                              load_result_param,
+                              Symbols::FutureThen(),
+                              args)));
+  const Library& core_lib = Library::Handle(Library::CoreLibrary());
+  const Function& async_catch_helper = Function::ZoneHandle(
+      Z, core_lib.LookupFunctionAllowPrivate(Symbols::AsyncCatchHelper()));
+  ASSERT(!async_catch_helper.IsNull());
+  ArgumentListNode* catch_helper_args = new (Z) ArgumentListNode(
       Scanner::kNoSourcePos);
-  async_await_helper_args->Add(
-      new(Z) LoadLocalNode(Scanner::kNoSourcePos, await_temp));
-  async_await_helper_args->Add(
-      new(Z) LoadLocalNode(Scanner::kNoSourcePos, async_then_callback));
-  async_await_helper_args->Add(
-      new(Z) LoadLocalNode(Scanner::kNoSourcePos, async_catch_error_callback));
-  StaticCallNode* await_helper_call = new (Z) StaticCallNode(
-      node->token_pos(),
-      async_await_helper,
-      async_await_helper_args);
-
-  preamble_->Add(new(Z) StoreLocalNode(
-      Scanner::kNoSourcePos, result_param, await_helper_call));
-
+  InstanceGetterNode* catch_error_getter = new (Z) InstanceGetterNode(
+      Scanner::kNoSourcePos,
+      load_result_param,
+      Symbols::FutureCatchError());
+  catch_helper_args->Add(catch_error_getter);
+  catch_helper_args->Add(new (Z) LoadLocalNode(
+      Scanner::kNoSourcePos, async_op));
+  preamble_->Add(new (Z) StaticCallNode(
+      Scanner::kNoSourcePos,
+      async_catch_helper,
+      catch_helper_args));
   ReturnNode* continuation_return = new(Z) ReturnNode(Scanner::kNoSourcePos);
   continuation_return->set_return_type(ReturnNode::kContinuationTarget);
   preamble_->Add(continuation_return);