[vm/bytecode] Fix AST removal for package-split kernel files with bytecode

In package-split mode, bytecode generation is performed separately for
each package. Previously, dropping AST was done right after generating
bytecode. However, dropping AST for a package makes it impossible to do
constant evaluation in other packages which import the package with dropped
AST. This breaks bytecode generation for subsequent packages.

To work around this problem, in package-split mode AST is removed
temporary until dillp file is written. After that, removed AST is restored
back.

Change-Id: I3d8b6a8ad98f2fe88b57f7b6393bbbe87b046c21
Reviewed-on: https://dart-review.googlesource.com/c/89822
Auto-Submit: Alexander Markov <alexmarkov@google.com>
Reviewed-by: Zach Anderson <zra@google.com>
Commit-Queue: Alexander Markov <alexmarkov@google.com>

pkg/vm/lib/bytecode/ast_remover.dart

@@ -0,0 +1,74 @@
+// Copyright (c) 2019, 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 vm.bytecode.ast_remover;
+
+import 'package:kernel/ast.dart' hide MapEntry;
+import '../metadata/bytecode.dart';
+
+/// Drops kernel AST for members with bytecode.
+/// Can preserve removed AST and restore it if needed.
+class ASTRemover extends Transformer {
+  final BytecodeMetadataRepository metadata;
+  final droppedAST = <Member, dynamic>{};
+
+  ASTRemover(Component component)
+      : metadata = component.metadata[new BytecodeMetadataRepository().tag];
+
+  @override
+  TreeNode defaultMember(Member node) {
+    if (_hasBytecode(node)) {
+      if (node is Field) {
+        droppedAST[node] = node.initializer;
+        node.initializer = null;
+      } else if (node is Constructor) {
+        droppedAST[node] =
+            new _DroppedConstructor(node.initializers, node.function.body);
+        node.initializers = <Initializer>[];
+        node.function.body = null;
+      } else if (node.function != null) {
+        droppedAST[node] = node.function.body;
+        node.function.body = null;
+      }
+    }
+
+    // Instance field initializers do not form separate functions, and bytecode
+    // is not attached to instance fields (it is included into constructors).
+    // When VM reads a constructor from kernel, it also reads and translates
+    // instance field initializers. So, their ASTs can be dropped only if
+    // bytecode was generated for all generative constructors.
+    if (node is Field && !node.isStatic && node.initializer != null) {
+      if (node.enclosingClass.constructors.every(_hasBytecode)) {
+        droppedAST[node] = node.initializer;
+        node.initializer = null;
+      }
+    }
+
+    return node;
+  }
+
+  bool _hasBytecode(Member node) =>
+      metadata != null && metadata.mapping.containsKey(node);
+
+  void restoreAST() {
+    droppedAST.forEach((Member node, dynamic dropped) {
+      if (node is Field) {
+        node.initializer = dropped;
+      } else if (node is Constructor) {
+        _DroppedConstructor droppedConstructor = dropped;
+        node.initializers = droppedConstructor.initializers;
+        node.function.body = droppedConstructor.body;
+      } else {
+        node.function.body = dropped;
+      }
+    });
+  }
+}
+
+class _DroppedConstructor {
+  final List<Initializer> initializers;
+  final Statement body;
+
+  _DroppedConstructor(this.initializers, this.body);
+}

pkg/vm/lib/bytecode/gen_bytecode.dart

@@ -41,7 +41,6 @@ const String symbolForTypeCast = ' in type cast';
 
 void generateBytecode(
   Component component, {
-  bool dropAST: false,
   bool emitSourcePositions: false,
   bool omitAssertSourcePositions: false,
   bool useFutureBytecodeFormat: false,
@@ -71,12 +70,6 @@ void generateBytecode(
   for (var library in libraries) {
     bytecodeGenerator.visitLibrary(library);
   }
-  if (dropAST) {
-    final astRemover = new DropAST(component);
-    for (var library in libraries) {
-      astRemover.visitLibrary(library);
-    }
-  }
 }
 
 class BytecodeGenerator extends RecursiveVisitor<Null> {
@@ -3129,44 +3122,6 @@ class FindFreeTypeParametersVisitor extends DartTypeVisitor<bool> {
   }
 }
 
-// Drop kernel AST for members with bytecode.
-class DropAST extends Transformer {
-  BytecodeMetadataRepository metadata;
-
-  DropAST(Component component)
-      : metadata = component.metadata[new BytecodeMetadataRepository().tag];
-
-  @override
-  TreeNode defaultMember(Member node) {
-    if (_hasBytecode(node)) {
-      if (node is Field) {
-        node.initializer = null;
-      } else if (node is Constructor) {
-        node.initializers = <Initializer>[];
-        node.function.body = null;
-      } else if (node.function != null) {
-        node.function.body = null;
-      }
-    }
-
-    // Instance field initializers do not form separate functions, and bytecode
-    // is not attached to instance fields (it is included into constructors).
-    // When VM reads a constructor from kernel, it also reads and translates
-    // instance field initializers. So, their ASTs can be dropped only if
-    // bytecode was generated for all generative constructors.
-    if (node is Field && !node.isStatic && node.initializer != null) {
-      if (node.enclosingClass.constructors.every(_hasBytecode)) {
-        node.initializer = null;
-      }
-    }
-
-    return node;
-  }
-
-  bool _hasBytecode(Member node) =>
-      metadata != null && metadata.mapping.containsKey(node);
-}
-
 typedef void GenerateContinuation();
 
 class FinallyBlock {

pkg/vm/lib/kernel_front_end.dart

@@ -41,6 +41,7 @@ import 'package:kernel/target/targets.dart' show Target, TargetFlags, getTarget;
 import 'package:kernel/transformations/constants.dart' as constants;
 import 'package:kernel/vm/constants_native_effects.dart' as vm_constants;
 
+import 'bytecode/ast_remover.dart' show ASTRemover;
 import 'bytecode/gen_bytecode.dart' show generateBytecode;
 
 import 'constants_error_reporter.dart' show ForwardConstantEvaluationErrors;
@@ -288,11 +289,14 @@ Future<Component> compileToKernel(Uri source, CompilerOptions options,
   if (genBytecode && !errorDetector.hasCompilationErrors && component != null) {
     await runWithFrontEndCompilerContext(source, options, component, () {
       generateBytecode(component,
-          dropAST: dropAST,
           emitSourcePositions: emitBytecodeSourcePositions,
           useFutureBytecodeFormat: useFutureBytecodeFormat,
           environmentDefines: environmentDefines);
     });
+
+    if (dropAST) {
+      new ASTRemover(component).visitComponent(component);
+    }
   }
 
   // Restore error handler (in case 'options' are reused).
@@ -613,22 +617,32 @@ Future writeOutputSplitByPackages(
         component.mainMethod = null;
       }
 
+      ASTRemover astRemover;
       if (genBytecode) {
         final List<Library> libraries = component.libraries
             .where((lib) => packageFor(lib) == package)
             .toList();
         generateBytecode(component,
             libraries: libraries,
-            dropAST: dropAST,
             emitSourcePositions: emitBytecodeSourcePositions,
             useFutureBytecodeFormat: useFutureBytecodeFormat,
             environmentDefines: environmentDefines);
+
+        if (dropAST) {
+          astRemover = new ASTRemover(component);
+          for (var library in libraries) {
+            astRemover.visitLibrary(library);
+          }
+        }
       }
 
       final BinaryPrinter printer = new LimitedBinaryPrinter(sink,
           (lib) => packageFor(lib) == package, false /* excludeUriToSource */);
       printer.writeComponentFile(component);
       component.mainMethod = main;
+      if (genBytecode && dropAST) {
+        astRemover.restoreAST();
+      }
 
       await sink.close();
     }