VM: [Kernel] Fix remaining issues with kernel-based async/await implementation

pkg/kernel/lib/transformations/continuation.dart

 // Copyright (c) 2016, the Dart project authors.  Please see the AUTHORS file
 // for details. All rights reserved. Use of this source code is governed by a
 // BSD-style license that can be found in the LICENSE file.
 
 library kernel.transformations.continuation;
 
 import 'dart:math' as math;
 
 import '../ast.dart';
 import '../visitor.dart';
@@ skipping 32 matching lines @@
         return new AsyncFunctionRewriter(helper, node).rewrite();
       case AsyncMarker.AsyncStar:
         return new AsyncStarFunctionRewriter(helper, node).rewrite();
     }
   }
 }
 
 abstract class ContinuationRewriterBase extends RecursiveContinuationRewriter {
   final FunctionNode enclosingFunction;
 
-  int currentTryDepth; // Nesting depth for try-blocks.
+  int currentTryDepth = 0; // Nesting depth for try-blocks.
   int currentCatchDepth = 0; // Nesting depth for catch-blocks.
-  int capturedTryDepth = 0; // Deepest yield point within a try-block.
-  int capturedCatchDepth = 0; // Deepest yield point within a catch-block.
 
-  ContinuationRewriterBase(HelperNodes helper, this.enclosingFunction,
-      {this.currentTryDepth: 0})
+  // For both, try-catch and try-finally, the VM needs to have saved-try-context
+  // and exception/stacktrace available (both need the correct context when
+  // entering catch/finally blocks and finally needs exception/stacktrace for
+  // rethrow).
+  int capturedDepth = 0; // Deepest yield point within try/catch-blocks.
+
+  ContinuationRewriterBase(HelperNodes helper, this.enclosingFunction)
       : super(helper);
 
   Statement createContinuationPoint([Expression value]) {
     if (value == null) value = new NullLiteral();
-    capturedTryDepth = math.max(capturedTryDepth, currentTryDepth);
-    capturedCatchDepth = math.max(capturedCatchDepth, currentCatchDepth);
+    capturedDepth =
+        math.max(capturedDepth, currentTryDepth + currentCatchDepth);
     return new YieldStatement(value, isNative: true);
   }
 
   TreeNode visitTryCatch(TryCatch node) {
     if (node.body != null) {
       currentTryDepth++;
       node.body = node.body.accept(this);
       node.body?.parent = node;
       currentTryDepth--;
     }
 
-    currentCatchDepth++;
+    ++currentCatchDepth;
     transformList(node.catches, this, node);
-    currentCatchDepth--;
+    --currentCatchDepth;
     return node;
   }
 
   TreeNode visitTryFinally(TryFinally node) {
     if (node.body != null) {
       currentTryDepth++;
       node.body = node.body.accept(this);
       node.body?.parent = node;
       currentTryDepth--;
     }
     if (node.finalizer != null) {
       node.finalizer = node.finalizer.accept(this);
       node.finalizer?.parent = node;
     }
     return node;
   }
 
   Iterable<VariableDeclaration> createCapturedTryVariables() =>
-      new Iterable.generate(capturedTryDepth,
+      new Iterable.generate(capturedDepth,
           (depth) => new VariableDeclaration(":saved_try_context_var${depth}"));
 
   Iterable<VariableDeclaration> createCapturedCatchVariables() =>
-      new Iterable.generate(capturedCatchDepth).expand((depth) => [
+      new Iterable.generate(capturedDepth).expand((depth) => [
             new VariableDeclaration(":exception${depth}"),
             new VariableDeclaration(":stack_trace${depth}"),
           ]);
 
   List<VariableDeclaration> variableDeclarations() =>
       [asyncJumpVariable, asyncContextVariable]
         ..addAll(createCapturedTryVariables())
         ..addAll(createCapturedCatchVariables());
 }
 
@@ skipping 69 matching lines @@
   final VariableDeclaration thenContinuationVariable =
       new VariableDeclaration(":async_op_then");
   final VariableDeclaration catchErrorContinuationVariable =
       new VariableDeclaration(":async_op_error");
 
   LabeledStatement labeledBody;
 
   ExpressionLifter expressionRewriter;
 
   AsyncRewriterBase(helper, enclosingFunction)
-      // Body is wrapped in the try-catch so initial currentTryDepth is 1.
-      : super(helper, enclosingFunction, currentTryDepth: 1) {}
+      : super(helper, enclosingFunction) {}
 
   void setupAsyncContinuations(List<Statement> statements) {
     expressionRewriter = new ExpressionLifter(this);
 
     // var :async_op_then;
     statements.add(thenContinuationVariable);
 
     // var :async_op_error;
     statements.add(catchErrorContinuationVariable);
 
@@ skipping 36 matching lines @@
     final boundCatchErrorClosure = new StaticInvocation(
         helper.asyncErrorWrapper,
         new Arguments(<Expression>[new VariableGet(nestedClosureVariable)]));
     final catchErrorClosureVariableAssign = new ExpressionStatement(
         new VariableSet(
             catchErrorContinuationVariable, boundCatchErrorClosure));
     statements.add(catchErrorClosureVariableAssign);
   }
 
   Statement buildWrappedBody() {
+    ++currentCatchDepth;
     labeledBody = new LabeledStatement(null);
     labeledBody.body = visitDelimited(enclosingFunction.body)
       ..parent = labeledBody;
 
     var exceptionVariable = new VariableDeclaration(":exception");
     var stackTraceVariable = new VariableDeclaration(":stack_trace");
 
-    return new TryCatch(buildReturn(labeledBody), <Catch>[
+    var body = buildReturn(labeledBody);
+    --currentCatchDepth;
+    return new TryCatch(body, <Catch>[
       new Catch(
           exceptionVariable,
           new Block(<Statement>[
             buildCatchBody(exceptionVariable, stackTraceVariable)
           ]),
           stackTrace: stackTraceVariable)
     ]);
   }
 
   Statement buildCatchBody(
@@ skipping 409 matching lines @@
     var returnStatement = new ReturnStatement(
         new PropertyGet(completerGet, new Name('stream', helper.asyncLibrary)));
     statements.add(returnStatement);
 
     enclosingFunction.body = new Block(statements);
     enclosingFunction.body.parent = enclosingFunction;
     enclosingFunction.asyncMarker = AsyncMarker.Sync;
     return enclosingFunction;
   }
 
+  Statement buildWrappedBody() {
+    ++currentTryDepth;
+    Statement body = super.buildWrappedBody();
+    --currentTryDepth;
+    var tryFinally =
+        new TryFinally(body, new Block(<Statement>[buildFinallyBody()]));
+    return tryFinally;
+  }
+
   Statement buildCatchBody(exceptionVariable, stackTraceVariable) {
     return new ExpressionStatement(new MethodInvocation(
         new VariableGet(controllerVariable),
-        new Name("completeError", helper.asyncLibrary),
+        new Name("addError", helper.asyncLibrary),
         new Arguments(<Expression>[
           new VariableGet(exceptionVariable),
           new VariableGet(stackTraceVariable)
         ])));
   }
 
+  Statement buildFinallyBody() {
+    return new ExpressionStatement(new MethodInvocation(
+        new VariableGet(controllerVariable),
+        new Name("close", helper.asyncLibrary),
+        new Arguments(<Expression>[])));
+  }
+
   Statement buildReturn(Statement body) {
     // Async* functions cannot return a value.  The returns from the function
     // have been translated into breaks from the labeled body.
     return new Block(<Statement>[
       body,
-      new ExpressionStatement(new MethodInvocation(
-          new VariableGet(controllerVariable),
-          new Name("close", helper.asyncLibrary),
-          new Arguments(<Expression>[]))),
-      new ReturnStatement()..fileOffset = enclosingFunction.fileEndOffset
+      new ReturnStatement()..fileOffset = enclosingFunction.fileEndOffset,
     ]);
   }
 
   TreeNode visitYieldStatement(YieldStatement stmt) {
     Expression expr = expressionRewriter.rewrite(stmt.expression, statements);
 
     var addExpression = new MethodInvocation(
         new VariableGet(controllerVariable),
         new Name(stmt.isYieldStar ? 'addStream' : 'add', helper.asyncLibrary),
         new Arguments(<Expression>[expr]))..fileOffset = stmt.fileOffset;
@@ skipping 190 matching lines @@
         findFactoryConstructor(completerClass, 'sync'),
         findConstructor(syncIterableClass, ''),
         findConstructor(streamIteratorClass, ''),
         findFactoryConstructor(futureClass, 'microtask'),
         findConstructor(streamControllerClass, ''),
         findProcedure(asyncLibrary, '_asyncThenWrapperHelper'),
         findProcedure(asyncLibrary, '_asyncErrorWrapperHelper'),
         findProcedure(asyncLibrary, '_awaitHelper'));
   }
 }