[vm] First stab at a protobuf aware treeshaker

Relies on https://dart-review.googlesource.com/c/sdk/+/96782
Change-Id: I820e884f96923f9eb9a856faef6442f6ed645fe1
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/96162
Commit-Queue: Sigurd Meldgaard <sigurdm@google.com>
Reviewed-by: Martin Kustermann <kustermann@google.com>
diff --git a/pkg/vm/bin/protobuf_aware_treeshaker.dart b/pkg/vm/bin/protobuf_aware_treeshaker.dart
new file mode 100644
index 0000000..b971b13
--- /dev/null
+++ b/pkg/vm/bin/protobuf_aware_treeshaker.dart
@@ -0,0 +1,169 @@
+#!/usr/bin/env dart
+// 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.
+
+/// This program will take a .dill file and do a protobuf aware tree-shaking.
+///
+/// All fields of GeneratedMessage subclasses that are not accessed with their
+/// getter or setter will have their metadata removed from the class definition.
+///
+/// Then a general treeshaking will be run, and
+/// all GeneratedMessage subclasses that are never used directly will be
+/// removed.
+///
+/// The processed program will have observable differences: The tree-shaken
+/// fields will be parsed as unknown fields.
+/// The toString method will treat the unknown fields as missing.
+///
+/// Using the `GeneratedMessage.info_` field to reflect on fields will have
+/// unpredictable behavior.
+///
+/// Constants are evaluated, this is mainly to enable detecting
+/// `@pragma('vm:entry-point')`.
+library vm.protobuf_aware_treeshaker;
+
+import 'dart:async';
+import 'dart:io';
+import 'dart:typed_data';
+
+import 'package:args/args.dart';
+import 'package:kernel/kernel.dart';
+import 'package:kernel/binary/limited_ast_to_binary.dart';
+import 'package:kernel/target/targets.dart' show TargetFlags, getTarget;
+import 'package:meta/meta.dart';
+import 'package:vm/target/install.dart' show installAdditionalTargets;
+import 'package:vm/transformations/protobuf_aware_treeshaker/transformer.dart'
+    as treeshaker;
+
+ArgResults parseArgs(List<String> args) {
+  ArgParser argParser = ArgParser()
+    ..addOption('platform',
+        valueHelp: "path/to/vm_platform.dill",
+        help: 'A platform.dill file to append to the input. If not given, no '
+            'platform.dill will be appended.')
+    ..addOption('target',
+        allowed: ['dart_runner', 'flutter', 'flutter-runner', 'vm'],
+        defaultsTo: 'vm',
+        help: 'A platform.dill file to append to the input. If not given, no '
+            'platform.dill will be appended.')
+    ..addFlag('write-txt',
+        help: 'Also write the result in kernel-text format as <out.dill>.txt',
+        defaultsTo: false)
+    ..addFlag('remove-core-libs',
+        help:
+            'If set, the resulting dill file will not include `dart:` libraries',
+        defaultsTo: false)
+    ..addMultiOption('define',
+        abbr: 'D',
+        help: 'Perform constant evaluation with this environment define set.',
+        valueHelp: 'variable=value')
+    ..addFlag('remove-source',
+        help: 'Removes source code from the emitted dill', defaultsTo: false)
+    ..addFlag('enable-asserts',
+        help: 'Enables asserts in the emitted dill', defaultsTo: false)
+    ..addFlag('verbose',
+        help: 'Write to stdout about what classes and fields where remeoved')
+    ..addFlag('help', help: 'Prints this help', negatable: false);
+
+  ArgResults argResults;
+  try {
+    argResults = argParser.parse(args);
+  } on FormatException catch (e) {
+    print(e.message);
+  }
+  if (argResults == null || argResults['help'] || argResults.rest.length != 2) {
+    String script = 'protobuf_aware_treeshaker.dart';
+    print(
+        'A tool for removing protobuf messages types that are never referred by a program');
+    print('Usage: $script [args] <input.dill> <output.dill>');
+
+    print(argParser.usage);
+    exit(-1);
+  }
+
+  return argResults;
+}
+
+Future main(List<String> args) async {
+  ArgResults argResults = parseArgs(args);
+
+  final input = argResults.rest[0];
+  final output = argResults.rest[1];
+
+  final Map<String, String> environment = Map.fromIterable(
+      argResults['define'].map((x) => x.split('=')),
+      key: (x) => x[0],
+      value: (x) => x[1]);
+
+  var bytes = File(input).readAsBytesSync();
+  final platformFile = argResults['platform'];
+  if (platformFile != null) {
+    bytes = concatenate(File(platformFile).readAsBytesSync(), bytes);
+  }
+  final component = loadComponentFromBytes(bytes);
+
+  installAdditionalTargets();
+
+  treeshaker.TransformationInfo info = treeshaker.transformComponent(
+      component, environment, getTarget(argResults['target'], TargetFlags()),
+      collectInfo: argResults['verbose'],
+      enableAsserts: argResults['enable-asserts']);
+
+  if (argResults['verbose']) {
+    for (String fieldName in info.removedMessageFields) {
+      print('Removed $fieldName');
+    }
+    for (Class removedClass in info.removedMessageClasses) {
+      print('Removed $removedClass');
+    }
+  }
+
+  await writeComponent(component, output,
+      removeCoreLibs: argResults['remove-core-libs'],
+      removeSource: argResults['remove-source']);
+  if (argResults['write-txt']) {
+    writeComponentToText(component, path: output + '.txt');
+  }
+}
+
+Uint8List concatenate(Uint8List a, Uint8List b) {
+  final bytes = Uint8List(a.length + b.length);
+  bytes.setRange(0, a.length, a);
+  bytes.setRange(a.length, a.length + b.length, b);
+  return bytes;
+}
+
+Future writeComponent(Component component, String filename,
+    {@required bool removeCoreLibs, @required bool removeSource}) async {
+  if (removeSource) {
+    component.uriToSource.clear();
+  }
+
+  for (final lib in component.libraries) {
+    lib.dependencies.clear();
+    lib.additionalExports.clear();
+    lib.parts.clear();
+  }
+
+  final sink = File(filename).openWrite();
+  final printer = LimitedBinaryPrinter(sink, (lib) {
+    if (removeCoreLibs && isCoreLibrary(lib)) return false;
+    if (isLibEmpty(lib)) return false;
+    return true;
+  }, /*excludeUriToSource=*/ removeSource);
+
+  printer.writeComponentFile(component);
+  await sink.close();
+}
+
+bool isLibEmpty(Library lib) {
+  return lib.classes.isEmpty &&
+      lib.procedures.isEmpty &&
+      lib.fields.isEmpty &&
+      lib.typedefs.isEmpty;
+}
+
+bool isCoreLibrary(Library library) {
+  return library.importUri.scheme == 'dart';
+}
diff --git a/pkg/vm/lib/transformations/protobuf_aware_treeshaker/transformer.dart b/pkg/vm/lib/transformations/protobuf_aware_treeshaker/transformer.dart
new file mode 100644
index 0000000..fba5844
--- /dev/null
+++ b/pkg/vm/lib/transformations/protobuf_aware_treeshaker/transformer.dart
@@ -0,0 +1,270 @@
+// 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.
+
+import 'package:kernel/kernel.dart';
+import 'package:kernel/target/targets.dart';
+import 'package:kernel/core_types.dart';
+import 'package:kernel/transformations/constants.dart' as constants;
+import 'package:kernel/vm/constants_native_effects.dart' as vm_constants;
+import 'package:meta/meta.dart';
+import 'package:vm/transformations/type_flow/transformer.dart' as globalTypeFlow
+    show transformComponent;
+import 'package:vm/transformations/no_dynamic_invocations_annotator.dart'
+    show Selector;
+
+class TransformationInfo {
+  final List<String> removedMessageFields = <String>[];
+  final List<Class> removedMessageClasses = <Class>[];
+}
+
+TransformationInfo transformComponent(
+    Component component, Map<String, String> environment, Target target,
+    {@required bool collectInfo, @required bool enableAsserts}) {
+  final coreTypes = new CoreTypes(component);
+  component.computeCanonicalNames();
+
+  // Evaluate constants to ensure @pragma("vm:entry-point") is seen by the
+  // type-flow analysis.
+  final vmConstants = new vm_constants.VmConstantsBackend(coreTypes);
+  constants.transformComponent(component, vmConstants, environment, null,
+      keepFields: true,
+      evaluateAnnotations: true,
+      enableAsserts: enableAsserts);
+
+  TransformationInfo info = collectInfo ? TransformationInfo() : null;
+
+  _treeshakeProtos(target, component, coreTypes, info);
+  return info;
+}
+
+void _treeshakeProtos(Target target, Component component, CoreTypes coreTypes,
+    TransformationInfo info) {
+  globalTypeFlow.transformComponent(target, coreTypes, component);
+  final collector = _removeUnusedProtoReferences(component, coreTypes, info);
+  if (collector == null) {
+    return;
+  }
+  globalTypeFlow.transformComponent(target, coreTypes, component);
+  if (info != null) {
+    for (Class gmSubclass in collector.gmSubclasses) {
+      if (!gmSubclass.enclosingLibrary.classes.contains(gmSubclass)) {
+        info.removedMessageClasses.add(gmSubclass);
+      }
+    }
+  }
+  // Remove metadata added by the typeflow analysis.
+  component.metadata.clear();
+}
+
+InfoCollector _removeUnusedProtoReferences(
+    Component component, CoreTypes coreTypes, TransformationInfo info) {
+  final protobufUri = Uri.parse('package:protobuf/protobuf.dart');
+  final protobufLibs =
+      component.libraries.where((lib) => lib.importUri == protobufUri);
+  if (protobufLibs.isEmpty) {
+    return null;
+  }
+  final protobufLib = protobufLibs.single;
+
+  final gmClass = protobufLib.classes
+      .where((klass) => klass.name == 'GeneratedMessage')
+      .single;
+  final collector = InfoCollector(gmClass);
+
+  final biClass =
+      protobufLib.classes.where((klass) => klass.name == 'BuilderInfo').single;
+  final addMethod =
+      biClass.members.singleWhere((Member member) => member.name.name == 'add');
+
+  component.accept(collector);
+
+  _UnusedFieldMetadataPruner(
+          biClass, addMethod, collector.dynamicSelectors, coreTypes, info)
+      .removeMetadataForUnusedFields(
+    collector.gmSubclasses,
+    collector.gmSubclassesInvokedMethods,
+    coreTypes,
+    info,
+  );
+
+  return collector;
+}
+
+/// For protobuf fields which are not accessed, prune away its metadata.
+class _UnusedFieldMetadataPruner extends TreeVisitor<void> {
+  // All of those methods have the dart field name as second positional
+  // parameter.
+  // Method names are defined in:
+  // https://github.com/dart-lang/protobuf/blob/master/protobuf/lib/src/protobuf/builder_info.dart
+  // The code is generated by:
+  // https://github.com/dart-lang/protobuf/blob/master/protoc_plugin/lib/protobuf_field.dart.
+  static final fieldAddingMethods = Set<String>.from(const <String>[
+    'a',
+    'm',
+    'pp',
+    'pc',
+    'e',
+    'pc',
+    'aOS',
+    'aOB',
+  ]);
+
+  final Class builderInfoClass;
+  Class visitedClass;
+  final names = Set<String>();
+
+  final dynamicNames = Set<String>();
+  final CoreTypes coreTypes;
+  final TransformationInfo info;
+  final Member addMethod;
+
+  _UnusedFieldMetadataPruner(this.builderInfoClass, this.addMethod,
+      Set<Selector> dynamicSelectors, this.coreTypes, this.info) {
+    dynamicNames.addAll(dynamicSelectors.map((sel) => sel.target.name));
+  }
+
+  /// If a proto message field is never accessed (neither read nor written to),
+  /// remove its corresponding metadata in the construction of the Message._i
+  /// field (i.e. the BuilderInfo metadata).
+  void removeMetadataForUnusedFields(
+      Set<Class> gmSubclasses,
+      Map<Class, Set<Selector>> invokedMethods,
+      CoreTypes coreTypes,
+      TransformationInfo info) {
+    for (final klass in gmSubclasses) {
+      final selectors = invokedMethods[klass] ?? Set<Selector>();
+      final builderInfoFields = klass.fields.where((f) => f.name.name == '_i');
+      if (builderInfoFields.isEmpty) {
+        continue;
+      }
+      final builderInfoField = builderInfoFields.single;
+      _pruneBuilderInfoField(builderInfoField, selectors, klass);
+    }
+  }
+
+  void _pruneBuilderInfoField(
+      Field field, Set<Selector> selectors, Class gmSubclass) {
+    names.clear();
+    names.addAll(selectors.map((sel) => sel.target.name));
+    visitedClass = gmSubclass;
+    field.initializer.accept(this);
+  }
+
+  @override
+  visitLet(Let node) {
+    final initializer = node.variable.initializer;
+    if (initializer is MethodInvocation &&
+        initializer.interfaceTarget?.enclosingClass == builderInfoClass &&
+        fieldAddingMethods.contains(initializer.name.name)) {
+      final fieldName =
+          (initializer.arguments.positional[1] as StringLiteral).value;
+      final ucase = fieldName[0].toUpperCase() + fieldName.substring(1);
+      // The name of the related `clear` method.
+      final clearName = 'clear${ucase}';
+      // The name of the related `has` method.
+      final hasName = 'has${ucase}';
+
+      bool nameIsUsed(String name) =>
+          dynamicNames.contains(name) || names.contains(name);
+
+      if (!(nameIsUsed(fieldName) ||
+          nameIsUsed(clearName) ||
+          nameIsUsed(hasName))) {
+        if (info != null) {
+          info.removedMessageFields.add("${visitedClass.name}.$fieldName");
+        }
+
+        // Replace the field metadata method with a dummy call to
+        // `BuilderInfo.add`. This is to preserve the index calculations when
+        // removing a field.
+        // Change the tag-number to 0. Otherwise the decoder will get confused.
+        initializer.interfaceTarget = addMethod;
+        initializer.name = addMethod.name;
+        initializer.arguments.replaceWith(
+          Arguments(
+            <Expression>[
+              IntLiteral(0), // tagNumber
+              NullLiteral(), // name
+              NullLiteral(), // fieldType
+              NullLiteral(), // defaultOrMaker
+              NullLiteral(), // subBuilder
+              NullLiteral(), // valueOf
+              NullLiteral(), // enumValues
+            ],
+            types: <DartType>[InterfaceType(coreTypes.nullClass)],
+          ),
+        );
+      }
+    }
+    node.body.accept(this);
+  }
+}
+
+/// Finds all subclasses of [GeneratedMessage] and all methods invoked on them
+/// (potentially in a dynamic call).
+class InfoCollector extends RecursiveVisitor<void> {
+  final dynamicSelectors = Set<Selector>();
+  final Class generatedMessageClass;
+  final gmSubclasses = Set<Class>();
+  final gmSubclassesInvokedMethods = Map<Class, Set<Selector>>();
+
+  InfoCollector(this.generatedMessageClass);
+
+  @override
+  visitClass(Class klass) {
+    if (isGeneratedMethodSubclass(klass)) {
+      gmSubclasses.add(klass);
+    }
+    return super.visitClass(klass);
+  }
+
+  @override
+  visitMethodInvocation(MethodInvocation node) {
+    if (node.interfaceTarget == null) {
+      dynamicSelectors.add(Selector.doInvoke(node.name));
+    }
+
+    final targetClass = node.interfaceTarget?.enclosingClass;
+    if (isGeneratedMethodSubclass(targetClass)) {
+      addInvokedMethod(targetClass, Selector.doInvoke(node.name));
+    }
+    super.visitMethodInvocation(node);
+  }
+
+  @override
+  visitPropertyGet(PropertyGet node) {
+    if (node.interfaceTarget == null) {
+      dynamicSelectors.add(Selector.doGet(node.name));
+    }
+
+    final targetClass = node.interfaceTarget?.enclosingClass;
+    if (isGeneratedMethodSubclass(targetClass)) {
+      addInvokedMethod(targetClass, Selector.doGet(node.name));
+    }
+    super.visitPropertyGet(node);
+  }
+
+  @override
+  visitPropertySet(PropertySet node) {
+    if (node.interfaceTarget == null) {
+      dynamicSelectors.add(Selector.doSet(node.name));
+    }
+
+    final targetClass = node.interfaceTarget?.enclosingClass;
+    if (isGeneratedMethodSubclass(targetClass)) {
+      addInvokedMethod(targetClass, Selector.doSet(node.name));
+    }
+    super.visitPropertySet(node);
+  }
+
+  bool isGeneratedMethodSubclass(Class klass) {
+    return klass?.superclass == generatedMessageClass;
+  }
+
+  void addInvokedMethod(Class klass, Selector selector) {
+    final selectors =
+        gmSubclassesInvokedMethods.putIfAbsent(klass, () => Set<Selector>());
+    selectors.add(selector);
+  }
+}
diff --git a/pkg/vm/test/transformations/protobuf_aware_treeshaker/treeshaker_test.dart b/pkg/vm/test/transformations/protobuf_aware_treeshaker/treeshaker_test.dart
new file mode 100644
index 0000000..66d662e
--- /dev/null
+++ b/pkg/vm/test/transformations/protobuf_aware_treeshaker/treeshaker_test.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.
+
+import 'dart:io';
+
+import 'package:kernel/target/targets.dart';
+import 'package:kernel/ast.dart';
+import 'package:kernel/kernel.dart';
+import 'package:kernel/binary/limited_ast_to_binary.dart';
+import 'package:path/path.dart' as path;
+import 'package:test/test.dart';
+
+import 'package:vm/transformations/protobuf_aware_treeshaker/transformer.dart'
+    as treeshaker;
+
+import '../../common_test_utils.dart';
+
+final String pkgVmDir = Platform.script.resolve('../../..').toFilePath();
+
+runTestCase(Uri source) async {
+  final target = new TestingVmTarget(new TargetFlags());
+  Component component =
+      await compileTestCaseToKernelProgram(source, target: target);
+
+  List<Class> messageClasses = component.libraries
+      .expand(
+        (lib) => lib.classes.where((klass) =>
+            klass.superclass != null &&
+            klass.superclass.name == "GeneratedMessage"),
+      )
+      .toList();
+
+  treeshaker.transformComponent(component, {}, TestingVmTarget(TargetFlags()),
+      collectInfo: true, enableAsserts: false);
+
+  for (Class messageClass in messageClasses) {
+    expect(messageClass.enclosingLibrary.classes.contains(messageClass),
+        messageClass.name.endsWith('Keep'));
+  }
+
+  final systemTempDir = Directory.systemTemp;
+  final file =
+      new File('${systemTempDir.path}/${source.pathSegments.last}.dill');
+  try {
+    final sink = file.openWrite();
+    final printer =
+        LimitedBinaryPrinter(sink, (lib) => true, /*excludeUriToSource=*/ true);
+
+    printer.writeComponentFile(component);
+    await sink.close();
+
+    ProcessResult result =
+        Process.runSync(Platform.resolvedExecutable, [file.path]);
+    expect(result.exitCode, 0);
+  } finally {
+    if (file.existsSync()) {
+      file.deleteSync();
+    }
+  }
+}
+
+main() async {
+  final testCases = Directory(path.join(
+    pkgVmDir,
+    'testcases',
+    'transformations',
+    'protobuf_aware_treeshaker',
+    'lib',
+  )).listSync().where((f) => f.path.endsWith('_test.dart'));
+  for (final entry in testCases) {
+    test(entry.path, () => runTestCase(entry.uri));
+  }
+}
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/compile_protos.sh b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/compile_protos.sh
new file mode 100755
index 0000000..b4793ee
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/compile_protos.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+# 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.
+
+# Running this script requires having protoc_plugin installed in your path.
+
+rm -rf lib/generated
+mkdir lib/generated
+
+# Directory of the script
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+GENERATED_DIR=$DIR/lib/generated
+
+protoc --dart_out=$GENERATED_DIR -I$DIR/protos $DIR/protos/*.proto
+rm $GENERATED_DIR/*.pbenum.dart $GENERATED_DIR/*.pbjson.dart $GENERATED_DIR/*.pbserver.dart
+
+dartfmt -w $DIR/lib/generated
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/create_test.dart b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/create_test.dart
new file mode 100644
index 0000000..1912459
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/create_test.dart
@@ -0,0 +1,21 @@
+// 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.
+
+import 'package:test/test.dart';
+
+import 'generated/foo.pb.dart';
+
+main() {
+  FooKeep foo = FooKeep()
+    ..barKeep = (BarKeep()..aKeep = 5)
+    ..mapKeep['foo'] = (BarKeep()..aKeep = 2)
+    ..aKeep = 43;
+  test('retrieving values', () {
+    expect(foo.barKeep.aKeep, 5);
+    expect(foo.mapKeep['foo'].aKeep, 2);
+    expect(foo.hasHasKeep(), false);
+    expect(foo.aKeep, 43);
+    foo.clearClearKeep();
+  });
+}
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/decode_test.dart b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/decode_test.dart
new file mode 100644
index 0000000..d33289f
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/decode_test.dart
@@ -0,0 +1,28 @@
+// 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.
+
+import 'package:test/test.dart';
+
+import 'generated/foo.pb.dart';
+
+List<int> buffer = <int>[
+  10, 4, 8, 5, 16, //
+  4, 26, 9, 10, 3,
+  102, 111, 111, 18, 2,
+  8, 42, 34, 9, 10,
+  3, 122, 111, 112, 18,
+  2, 8, 3, 40, 43,
+  50, 0, 58, 0,
+];
+
+main() {
+  FooKeep foo = FooKeep.fromBuffer(buffer);
+  test('Kept values are restored correctly', () {
+    expect(foo.mapKeep['foo'].aKeep, 42);
+    expect(foo.barKeep.aKeep, 5);
+    expect(foo.aKeep, 43);
+    expect(foo.hasHasKeep(), true);
+    foo.clearClearKeep();
+  });
+}
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/encode_all_fields.dart b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/encode_all_fields.dart
new file mode 100644
index 0000000..de756cc
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/encode_all_fields.dart
@@ -0,0 +1,26 @@
+// 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.
+
+import 'dart:math';
+
+import 'generated/foo.pb.dart';
+
+main() {
+  FooKeep foo = FooKeep()
+    ..barKeep = (BarKeep()
+      ..aKeep = 5
+      ..bDrop = 4)
+    ..mapKeep['foo'] = (BarKeep()..aKeep = 42)
+    ..mapDrop['zop'] = (ZopDrop()..aDrop = 3)
+    ..aKeep = 43
+    ..hasKeep = HasKeep()
+    ..clearKeep = ClearKeep();
+  final buffer = foo.writeToBuffer();
+  print('List<int> buffer = <int>[');
+  for (int i = 0; i < buffer.length; i += 5) {
+    final numbers = buffer.sublist(i, min(buffer.length, i + 5)).join(', ');
+    print('  $numbers,${i == 0 ? ' //' : ''}');
+  }
+  print('];');
+}
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/generated/foo.pb.dart b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/generated/foo.pb.dart
new file mode 100644
index 0000000..4063773
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/lib/generated/foo.pb.dart
@@ -0,0 +1,253 @@
+///
+//  Generated code. Do not modify.
+//  source: foo.proto
+///
+// ignore_for_file: non_constant_identifier_names,library_prefixes,unused_import
+
+// ignore: UNUSED_SHOWN_NAME
+import 'dart:core' show int, bool, double, String, List, Map, override;
+
+import 'package:protobuf/protobuf.dart' as $pb;
+
+class FooKeep extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('FooKeep')
+    ..a<BarKeep>(
+        1, 'barKeep', $pb.PbFieldType.OM, BarKeep.getDefault, BarKeep.create)
+    ..a<BarKeep>(
+        2, 'barDrop', $pb.PbFieldType.OM, BarKeep.getDefault, BarKeep.create)
+    ..m<String, BarKeep>(3, 'mapKeep', 'FooKeep.MapKeepEntry',
+        $pb.PbFieldType.OS, $pb.PbFieldType.OM, BarKeep.create, null, null)
+    ..m<String, ZopDrop>(4, 'mapDrop', 'FooKeep.MapDropEntry',
+        $pb.PbFieldType.OS, $pb.PbFieldType.OM, ZopDrop.create, null, null)
+    ..a<int>(5, 'aKeep', $pb.PbFieldType.O3)
+    ..a<HasKeep>(
+        6, 'hasKeep', $pb.PbFieldType.OM, HasKeep.getDefault, HasKeep.create)
+    ..a<ClearKeep>(7, 'clearKeep', $pb.PbFieldType.OM, ClearKeep.getDefault,
+        ClearKeep.create)
+    ..hasRequiredFields = false;
+
+  FooKeep() : super();
+  FooKeep.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  FooKeep.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  FooKeep clone() => new FooKeep()..mergeFromMessage(this);
+  FooKeep copyWith(void Function(FooKeep) updates) =>
+      super.copyWith((message) => updates(message as FooKeep));
+  $pb.BuilderInfo get info_ => _i;
+  static FooKeep create() => new FooKeep();
+  FooKeep createEmptyInstance() => create();
+  static $pb.PbList<FooKeep> createRepeated() => new $pb.PbList<FooKeep>();
+  static FooKeep getDefault() => _defaultInstance ??= create()..freeze();
+  static FooKeep _defaultInstance;
+
+  BarKeep get barKeep => $_getN(0);
+  set barKeep(BarKeep v) {
+    setField(1, v);
+  }
+
+  bool hasBarKeep() => $_has(0);
+  void clearBarKeep() => clearField(1);
+
+  BarKeep get barDrop => $_getN(1);
+  set barDrop(BarKeep v) {
+    setField(2, v);
+  }
+
+  bool hasBarDrop() => $_has(1);
+  void clearBarDrop() => clearField(2);
+
+  Map<String, BarKeep> get mapKeep => $_getMap(2);
+
+  Map<String, ZopDrop> get mapDrop => $_getMap(3);
+
+  int get aKeep => $_get(4, 0);
+  set aKeep(int v) {
+    $_setSignedInt32(4, v);
+  }
+
+  bool hasAKeep() => $_has(4);
+  void clearAKeep() => clearField(5);
+
+  HasKeep get hasKeep => $_getN(5);
+  set hasKeep(HasKeep v) {
+    setField(6, v);
+  }
+
+  bool hasHasKeep() => $_has(5);
+  void clearHasKeep() => clearField(6);
+
+  ClearKeep get clearKeep => $_getN(6);
+  set clearKeep(ClearKeep v) {
+    setField(7, v);
+  }
+
+  bool hasClearKeep() => $_has(6);
+  void clearClearKeep() => clearField(7);
+}
+
+class BarKeep extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('BarKeep')
+    ..a<int>(1, 'aKeep', $pb.PbFieldType.O3)
+    ..a<int>(2, 'bDrop', $pb.PbFieldType.O3)
+    ..hasRequiredFields = false;
+
+  BarKeep() : super();
+  BarKeep.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  BarKeep.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  BarKeep clone() => new BarKeep()..mergeFromMessage(this);
+  BarKeep copyWith(void Function(BarKeep) updates) =>
+      super.copyWith((message) => updates(message as BarKeep));
+  $pb.BuilderInfo get info_ => _i;
+  static BarKeep create() => new BarKeep();
+  BarKeep createEmptyInstance() => create();
+  static $pb.PbList<BarKeep> createRepeated() => new $pb.PbList<BarKeep>();
+  static BarKeep getDefault() => _defaultInstance ??= create()..freeze();
+  static BarKeep _defaultInstance;
+
+  int get aKeep => $_get(0, 0);
+  set aKeep(int v) {
+    $_setSignedInt32(0, v);
+  }
+
+  bool hasAKeep() => $_has(0);
+  void clearAKeep() => clearField(1);
+
+  int get bDrop => $_get(1, 0);
+  set bDrop(int v) {
+    $_setSignedInt32(1, v);
+  }
+
+  bool hasBDrop() => $_has(1);
+  void clearBDrop() => clearField(2);
+}
+
+class HasKeep extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('HasKeep')
+    ..a<int>(1, 'aDrop', $pb.PbFieldType.O3)
+    ..hasRequiredFields = false;
+
+  HasKeep() : super();
+  HasKeep.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  HasKeep.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  HasKeep clone() => new HasKeep()..mergeFromMessage(this);
+  HasKeep copyWith(void Function(HasKeep) updates) =>
+      super.copyWith((message) => updates(message as HasKeep));
+  $pb.BuilderInfo get info_ => _i;
+  static HasKeep create() => new HasKeep();
+  HasKeep createEmptyInstance() => create();
+  static $pb.PbList<HasKeep> createRepeated() => new $pb.PbList<HasKeep>();
+  static HasKeep getDefault() => _defaultInstance ??= create()..freeze();
+  static HasKeep _defaultInstance;
+
+  int get aDrop => $_get(0, 0);
+  set aDrop(int v) {
+    $_setSignedInt32(0, v);
+  }
+
+  bool hasADrop() => $_has(0);
+  void clearADrop() => clearField(1);
+}
+
+class ClearKeep extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('ClearKeep')
+    ..a<int>(1, 'aDrop', $pb.PbFieldType.O3)
+    ..hasRequiredFields = false;
+
+  ClearKeep() : super();
+  ClearKeep.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  ClearKeep.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  ClearKeep clone() => new ClearKeep()..mergeFromMessage(this);
+  ClearKeep copyWith(void Function(ClearKeep) updates) =>
+      super.copyWith((message) => updates(message as ClearKeep));
+  $pb.BuilderInfo get info_ => _i;
+  static ClearKeep create() => new ClearKeep();
+  ClearKeep createEmptyInstance() => create();
+  static $pb.PbList<ClearKeep> createRepeated() => new $pb.PbList<ClearKeep>();
+  static ClearKeep getDefault() => _defaultInstance ??= create()..freeze();
+  static ClearKeep _defaultInstance;
+
+  int get aDrop => $_get(0, 0);
+  set aDrop(int v) {
+    $_setSignedInt32(0, v);
+  }
+
+  bool hasADrop() => $_has(0);
+  void clearADrop() => clearField(1);
+}
+
+class ZopDrop extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('ZopDrop')
+    ..a<int>(1, 'aDrop', $pb.PbFieldType.O3)
+    ..hasRequiredFields = false;
+
+  ZopDrop() : super();
+  ZopDrop.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  ZopDrop.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  ZopDrop clone() => new ZopDrop()..mergeFromMessage(this);
+  ZopDrop copyWith(void Function(ZopDrop) updates) =>
+      super.copyWith((message) => updates(message as ZopDrop));
+  $pb.BuilderInfo get info_ => _i;
+  static ZopDrop create() => new ZopDrop();
+  ZopDrop createEmptyInstance() => create();
+  static $pb.PbList<ZopDrop> createRepeated() => new $pb.PbList<ZopDrop>();
+  static ZopDrop getDefault() => _defaultInstance ??= create()..freeze();
+  static ZopDrop _defaultInstance;
+
+  int get aDrop => $_get(0, 0);
+  set aDrop(int v) {
+    $_setSignedInt32(0, v);
+  }
+
+  bool hasADrop() => $_has(0);
+  void clearADrop() => clearField(1);
+}
+
+class MobDrop extends $pb.GeneratedMessage {
+  static final $pb.BuilderInfo _i = new $pb.BuilderInfo('MobDrop')
+    ..a<int>(1, 'aDrop', $pb.PbFieldType.O3)
+    ..hasRequiredFields = false;
+
+  MobDrop() : super();
+  MobDrop.fromBuffer(List<int> i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromBuffer(i, r);
+  MobDrop.fromJson(String i,
+      [$pb.ExtensionRegistry r = $pb.ExtensionRegistry.EMPTY])
+      : super.fromJson(i, r);
+  MobDrop clone() => new MobDrop()..mergeFromMessage(this);
+  MobDrop copyWith(void Function(MobDrop) updates) =>
+      super.copyWith((message) => updates(message as MobDrop));
+  $pb.BuilderInfo get info_ => _i;
+  static MobDrop create() => new MobDrop();
+  MobDrop createEmptyInstance() => create();
+  static $pb.PbList<MobDrop> createRepeated() => new $pb.PbList<MobDrop>();
+  static MobDrop getDefault() => _defaultInstance ??= create()..freeze();
+  static MobDrop _defaultInstance;
+
+  int get aDrop => $_get(0, 0);
+  set aDrop(int v) {
+    $_setSignedInt32(0, v);
+  }
+
+  bool hasADrop() => $_has(0);
+  void clearADrop() => clearField(1);
+}
diff --git a/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/protos/foo.proto b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/protos/foo.proto
new file mode 100644
index 0000000..0a14433
--- /dev/null
+++ b/pkg/vm/testcases/transformations/protobuf_aware_treeshaker/protos/foo.proto
@@ -0,0 +1,41 @@
+// 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.
+
+syntax = "proto3";
+
+message FooKeep {
+  BarKeep barKeep = 1;
+  BarKeep barDrop = 2;
+  map<string, BarKeep> mapKeep = 3;
+  map<string, ZopDrop> mapDrop = 4;
+  int32 aKeep = 5;
+  HasKeep hasKeep = 6;
+  ClearKeep clearKeep = 7;
+}
+
+message BarKeep {
+  int32 aKeep = 1;
+  int32 bDrop = 2;
+}
+
+message HasKeep {
+  int32 aDrop = 1;
+}
+
+message ClearKeep {
+  int32 aDrop = 1;
+}
+
+message ZopDrop {
+  int32 aDrop = 1;
+}
+
+message MobDrop {
+  int32 aDrop = 1;
+}
+
+message A {
+  B unused = 1;
+  C used = 2;
+}
\ No newline at end of file