closure conversion: Support closures in initializers

pkg/kernel/lib/transformations/closure/converter.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.closure.converter;
 
 import '../../ast.dart'
     show
         Arguments,
         Block,
         Catch,
         Class,
         Constructor,
         ConstructorInvocation,
         DartType,
         EmptyStatement,
         Expression,
         ExpressionStatement,
         Field,
         FieldInitializer,
         ForInStatement,
         ForStatement,
         FunctionDeclaration,
         FunctionExpression,
         FunctionNode,
         InferredValue,
         Initializer,
         InvalidExpression,
         InvocationExpression,
+        Let,
         Library,
         LocalInitializer,
         Member,
         MethodInvocation,
         Name,
         NamedExpression,
         NullLiteral,
         Procedure,
         ProcedureKind,
         PropertyGet,
@@ skipping 20 matching lines @@
 import '../../core_types.dart' show CoreTypes;
 
 import '../../type_algebra.dart' show substitute;
 
 import 'clone_without_body.dart' show CloneWithoutBody;
 
 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;
   final Set<VariableDeclaration> capturedVariables;
   final Map<FunctionNode, Set<TypeParameter>> capturedTypeVariables;
   final Map<FunctionNode, VariableDeclaration> thisAccess;
   final Map<FunctionNode, String> localNames;
 
   /// Records place-holders for cloning contexts. See [visitForStatement].
   final Set<InvalidExpression> contextClonePlaceHolders =
@@ skipping 16 matching lines @@
   Library currentLibrary;
 
   Class currentClass;
 
   Member currentMember;
 
   FunctionNode currentMemberFunction;
 
   FunctionNode currentFunction;
 
-  Block _currentBlock;
+  Context context;
 
-  int _insertionIndex = 0;
-
-  Context context;
+  AstRewriter rewriter;
 
   /// Maps original type variable (aka type parameter) to a hoisted type
   /// variable type.
   ///
   /// For example, consider:
   ///
   ///     class C<T> {
   ///       f() => (x) => x is T;
   ///     }
   ///
@@ skipping 22 matching lines @@
     return currentFunction == null || currentMemberFunction == currentFunction;
   }
 
   String get currentFileUri {
     if (currentMember is Constructor) return currentClass.fileUri;
     if (currentMember is Field) return (currentMember as Field).fileUri;
     if (currentMember is Procedure) return (currentMember as Procedure).fileUri;
     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;
     }
   }
 
   TreeNode visitLibrary(Library node) {
     assert(newLibraryMembers.isEmpty);
     if (node == contextClass.enclosingLibrary) return node;
 
     currentLibrary = node;
     node = super.visitLibrary(node);
@@ skipping 12 matching lines @@
   TreeNode visitClass(Class node) {
     assert(newClassMembers.isEmpty);
     currentClass = node;
     node = super.visitClass(node);
     newClassMembers.forEach(node.addMember);
     newClassMembers.clear();
     currentClass = null;
     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;
+  }
 
-    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;
     currentFunction = function;
     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",
         type: contextClass.rawType,
         isFinal: true);
     Context parent = context;
     context = context.toNestedContext(new VariableAccessor(contextVariable));
 
     Set<TypeParameter> captured = capturedTypeVariables[currentFunction];
     if (captured != null) {
@@ skipping 25 matching lines @@
         parent.update(node.variable, expression);
         return null;
       } else {
         node.variable.initializer = expression;
         expression.parent = node.variable;
         return node.variable;
       }
     });
   }
 
-  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:
   ///
   ///     class Closure#0 extends core::Object implements core::Function {
   ///       field _in::Context context;
   ///       constructor •(final _in::Context #t1) → dynamic
   ///         : self::Closure 0::context = #t1
   ///         ;
   ///       method call(/* The parameters of [function] */) → dynamic {
   ///         /// #t2 is [contextVariable].
@@ skipping 96 matching lines @@
     }
 
     return node;
   }
 
   void setupContextForFunctionBody(FunctionNode function) {
     Statement body = function.body;
     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) {
     assert(!capturedVariables.contains(node.variable));
     node.transformChildren(this);
     return node;
   }
 
   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;
     return node;
   }
 
   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;
     });
   }
 
   TreeNode visitVariableDeclaration(VariableDeclaration node) {
     node.transformChildren(this);
 
     if (!capturedVariables.contains(node)) return node;
     context.extend(node, node.initializer ?? new NullLiteral());
 
@@ skipping 73 matching lines @@
       // `clone-context` is a place-holder that will later be replaced by an
       // expression that clones the current closure context (see
       // [visitInvalidExpression]).
       return saveContext(() {
         context = context.toNestedContext();
         List<Statement> statements = <Statement>[];
         statements.addAll(node.variables);
         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);
   }
 
   TreeNode visitForInStatement(ForInStatement node) {
     if (capturedVariables.contains(node.variable)) {
       // In Dart, loop variables are new variables on each iteration of the
       // loop. This is only observable when the loop variable is captured by a
       // closure, so we need to transform the for-in loop when `node.variable`
@@ skipping 237 matching lines @@
   // TODO(ahe): Remove this method when we don't generate closure classes
   // anymore.
   void addClosureClassNote(Class closureClass) {
     closureClass.addMember(new Field(new Name("note"),
         type: coreTypes.stringClass.rawType,
         initializer: new StringLiteral(
             "This is temporary. The VM doesn't need closure classes."),
         fileUri: currentFileUri));
   }
 }