closure conversion: Support closures in initializers

pkg/kernel/lib/transformations/closure/converter.dart

Index: pkg/kernel/lib/transformations/closure/converter.dart
diff --git a/pkg/kernel/lib/transformations/closure/converter.dart b/pkg/kernel/lib/transformations/closure/converter.dart
index 91a36dfc5e17e4933a479a8f6bb807851dedf391..325d9e54abdc5fb050024fc2b00409feaa271f8d 100644
--- a/pkg/kernel/lib/transformations/closure/converter.dart
+++ b/pkg/kernel/lib/transformations/closure/converter.dart
@@ -27,6 +27,7 @@ import '../../ast.dart'
         Initializer,
         InvalidExpression,
         InvocationExpression,
+        Let,
         Library,
         LocalInitializer,
         Member,
@@ -67,6 +68,8 @@ import 'context.dart' show Context, NoContext;
 
 import 'info.dart' show ClosureInfo;
 
+import 'rewriter.dart' show AstRewriter, BlockRewriter, InitializerRewriter;
+
 class ClosureConverter extends Transformer {
   final CoreTypes coreTypes;
   final Class contextClass;
@@ -103,12 +106,10 @@ class ClosureConverter extends Transformer {
 
   FunctionNode currentFunction;
 
-  Block _currentBlock;
-
-  int _insertionIndex = 0;
-
   Context context;
 
+  AstRewriter rewriter;
+
   /// Maps original type variable (aka type parameter) to a hoisted type
   /// variable type.
   ///
@@ -150,20 +151,13 @@ class ClosureConverter extends Transformer {
     throw "No file uri for ${currentMember.runtimeType}";
   }
 
-  void insert(Statement statement) {
-    _currentBlock.statements.insert(_insertionIndex++, statement);
-    statement.parent = _currentBlock;
-  }
-
   TreeNode saveContext(TreeNode f()) {
-    Block savedBlock = _currentBlock;
-    int savedIndex = _insertionIndex;
+    AstRewriter old = rewriter;
     Context savedContext = context;
     try {
       return f();
     } finally {
-      _currentBlock = savedBlock;
-      _insertionIndex = savedIndex;
+      rewriter = old;
       context = savedContext;
     }
   }
@@ -196,26 +190,68 @@ class ClosureConverter extends Transformer {
     return node;
   }
 
+  void extendContextWith(VariableDeclaration parameter) {
+    context.extend(parameter, new VariableGet(parameter));
+  }
+
   TreeNode visitConstructor(Constructor node) {
     assert(isEmptyContext);
-
     currentMember = node;
-
+    // Transform initializers.
+    for (Initializer initializer in node.initializers) {
+      if (initializer is FieldInitializer) {
+        // Create a rewriter and a context for the initializer expression.
+        rewriter = new InitializerRewriter(initializer.value);
+        context = new NoContext(this);
+        // Save the expression to visit it in the extended context, since the
+        // rewriter will modify `initializer.value`.
+        Expression initializerExpression = initializer.value;
+        // Extend the context with all captured parameters of the constructor.
+        // TODO(karlklose): add a fine-grained analysis of captured parameters.
+        node.function.positionalParameters
+            .where(capturedVariables.contains)
+            .forEach(extendContextWith);
+        node.function.namedParameters
+            .where(capturedVariables.contains)
+            .forEach(extendContextWith);
+        // Transform the initializer expression.
+        var parent = initializerExpression.parent;
+        initializerExpression = initializerExpression.accept(this);
+        initializerExpression.parent = parent;
+        if (parent is Let) {
+          parent.body = initializerExpression;
+        } else if (parent is FieldInitializer) {
+          parent.value = initializerExpression;
+        } else {
+          throw "Found unexpected node '${node.runtimeType}, expected a 'Let' "
+              "or a 'FieldInitializer'.";
+        }
+      }
+    }
+    rewriter = null;
+    // Transform constructor body.
     FunctionNode function = node.function;
     if (function.body != null && function.body is! EmptyStatement) {
       setupContextForFunctionBody(function);
       VariableDeclaration self = thisAccess[currentMemberFunction];
-      // TODO(karlklose): transform initializers
       if (self != null) {
         context.extend(self, new ThisExpression());
       }
       node.function.accept(this);
       resetContext();
     }
-
     return node;
   }
 
+  AstRewriter makeRewriterForBody(FunctionNode function) {
+    Statement body = function.body;
+    if (body is! Block) {
+      body = new Block(<Statement>[body]);
+      function.body = function.body.parent = body;
+    }
+    return new BlockRewriter(body);
+  }
+
   Expression handleLocalFunction(FunctionNode function) {
     FunctionNode enclosingFunction = currentFunction;
     Map<TypeParameter, DartType> enclosingTypeSubstitution = typeSubstitution;
@@ -223,13 +259,7 @@ class ClosureConverter extends Transformer {
     Statement body = function.body;
     assert(body != null);
 
-    if (body is Block) {
-      _currentBlock = body;
-    } else {
-      _currentBlock = new Block(<Statement>[body]);
-      function.body = body.parent = _currentBlock;
-    }
-    _insertionIndex = 0;
+    rewriter = makeRewriterForBody(function);
 
     VariableDeclaration contextVariable = new VariableDeclaration(
         "#contextParameter",
@@ -275,9 +305,11 @@ class ClosureConverter extends Transformer {
     });
   }
 
-  TreeNode visitFunctionExpression(FunctionExpression node) => saveContext(() {
-        return handleLocalFunction(node.function);
-      });
+  TreeNode visitFunctionExpression(FunctionExpression node) {
+    return saveContext(() {
+      return handleLocalFunction(node.function);
+    });
+  }
 
   /// Add a new class to the current library that looks like this:
   ///
@@ -394,28 +426,21 @@ class ClosureConverter extends Transformer {
     assert(body != null);
     currentMemberFunction = function;
     // Ensure that the body is a block which becomes the current block.
-    if (body is Block) {
-      _currentBlock = body;
-    } else {
-      _currentBlock = new Block(<Statement>[body]);
-      function.body = body.parent = _currentBlock;
-    }
-    _insertionIndex = 0;
+    rewriter = makeRewriterForBody(function);
     // Start with no context.  This happens after setting up _currentBlock
     // so statements can be emitted into _currentBlock if necessary.
     context = new NoContext(this);
   }
 
   void resetContext() {
-    _currentBlock = null;
-    _insertionIndex = 0;
+    rewriter = null;
     context = null;
     currentMemberFunction = null;
     currentMember = null;
   }
 
   bool get isEmptyContext {
-    return _currentBlock == null && _insertionIndex == 0 && context == null;
+    return rewriter == null && context == null;
   }
 
   TreeNode visitLocalInitializer(LocalInitializer node) {
@@ -426,16 +451,14 @@ class ClosureConverter extends Transformer {
 
   TreeNode visitFunctionNode(FunctionNode node) {
     transformList(node.typeParameters, this, node);
-
-    void extend(VariableDeclaration parameter) {
-      context.extend(parameter, new VariableGet(parameter));
-    }
-
     // TODO: Can parameters contain initializers (e.g., for optional ones) that
     // need to be closure converted?
-    node.positionalParameters.where(capturedVariables.contains).forEach(extend);
-    node.namedParameters.where(capturedVariables.contains).forEach(extend);
-
+    node.positionalParameters
+        .where(capturedVariables.contains)
+        .forEach(extendContextWith);
+    node.namedParameters
+        .where(capturedVariables.contains)
+        .forEach(extendContextWith);
     assert(node.body != null);
     node.body = node.body.accept(this);
     node.body.parent = node;
@@ -444,25 +467,8 @@ class ClosureConverter extends Transformer {
 
   TreeNode visitBlock(Block node) {
     return saveContext(() {
-      if (_currentBlock != node) {
-        _currentBlock = node;
-        _insertionIndex = 0;
-      }
-
-      while (_insertionIndex < _currentBlock.statements.length) {
-        assert(_currentBlock == node);
-
-        var original = _currentBlock.statements[_insertionIndex];
-        var transformed = original.accept(this);
-        assert(_currentBlock.statements[_insertionIndex] == original);
-        if (transformed == null) {
-          _currentBlock.statements.removeAt(_insertionIndex);
-        } else {
-          _currentBlock.statements[_insertionIndex++] = transformed;
-          transformed.parent = _currentBlock;
-        }
-      }
-
+      BlockRewriter blockRewriter = rewriter = rewriter.forNestedBlock(node);
+      blockRewriter.transformStatements(node, this);
       return node;
     });
   }
@@ -556,9 +562,9 @@ class ClosureConverter extends Transformer {
         statements.add(node);
         node.variables.clear();
         node.updates.insert(0, cloneContext());
-        _currentBlock = new Block(statements);
-        _insertionIndex = 0;
-        return _currentBlock.accept(this);
+        Block block = new Block(statements);
+        rewriter = new BlockRewriter(block);
+        return block.accept(this);
       });
     }
     return super.visitForStatement(node);