Reapply "Account for @pragma("vm:entry-point") creating additional "r…

…oot" libraries when partitioning the program into loading units."

Weaken assertion in gen_snapshot requiring all libraries to have a loading unit as there can still be unreachable libraries:
  - Google3 and Fuchsia will compile all the sources in a package to a single dill file, then present multiple input dill files to the AOT compilation. Since the set of libraries was derived from package membership instead of imports, many can be unreachable.
  - When the root library's main comes from an export, the frontend's representation will incorrectly report the library containing main as the root library and the true root library may be unreachable from it.

Instead, assert only that surviving compiled code is assigned a loading unit.

TEST=gallery
Bug: flutter/gallery#545
Bug: #49325
Bug: #41974
Bug: b/237016312
Change-Id: Ia52563a6f517308d041368be11dcc85270f19acc
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/249724
Commit-Queue: Ryan Macnak <rmacnak@google.com>
Reviewed-by: Alexander Markov <alexmarkov@google.com>

pkg/vm/lib/kernel_front_end.dart

@@ -510,7 +510,7 @@ Future runGlobalTransformations(
   // it does it will need the obfuscation prohibitions.
   obfuscationProhibitions.transformComponent(component, coreTypes, target);
 
-  deferred_loading.transformComponent(component);
+  deferred_loading.transformComponent(component, coreTypes, target);
 }
 
 /// Runs given [action] with [CompilerContext]. This is needed to

pkg/vm/lib/transformations/deferred_loading.dart

@@ -5,8 +5,12 @@
 library vm.transformations.deferred_loading;
 
 import 'package:kernel/ast.dart';
+import 'package:kernel/core_types.dart' show CoreTypes;
+import 'package:kernel/target/targets.dart' show Target;
+
 import '../dominators.dart';
 import '../metadata/loading_units.dart';
+import 'pragma.dart';
 
 class _LoadingUnitBuilder {
   late int id;
@@ -37,7 +41,42 @@ class _LibraryVertex extends Vertex<_LibraryVertex> {
   String toString() => "_LibraryVertex(${library.importUri})";
 }
 
-List<LoadingUnit> computeLoadingUnits(Component component) {
+class HasEntryPointVisitor extends RecursiveVisitor {
+  final PragmaAnnotationParser parser;
+  bool _hasEntryPoint = false;
+
+  HasEntryPointVisitor(this.parser);
+
+  visitAnnotations(List<Expression> annotations) {
+    for (var ann in annotations) {
+      ParsedPragma? pragma = parser.parsePragma(ann);
+      if (pragma is ParsedEntryPointPragma) {
+        _hasEntryPoint = true;
+        return;
+      }
+    }
+  }
+
+  @override
+  visitClass(Class klass) {
+    visitAnnotations(klass.annotations);
+    klass.visitChildren(this);
+  }
+
+  @override
+  defaultMember(Member node) {
+    visitAnnotations(node.annotations);
+  }
+
+  bool hasEntryPoint(Library lib) {
+    _hasEntryPoint = false;
+    visitLibrary(lib);
+    return _hasEntryPoint;
+  }
+}
+
+List<LoadingUnit> computeLoadingUnits(
+    Component component, HasEntryPointVisitor visitor) {
   // 1. Build the dominator tree for the library import graph.
   final map = <Library, _LibraryVertex>{};
   for (final lib in component.libraries) {
@@ -51,10 +90,15 @@ List<LoadingUnit> computeLoadingUnits(Component component) {
   }
   final root = map[component.mainMethod!.enclosingLibrary]!;
 
-  // Fake imports from root library to every core library so they end up in
-  // the same loading unit attributed to the user's root library.
+  // Fake imports from root library to every core library or library containing
+  // an entry point pragma so that they end up in the same loading unit
+  // attributed to the user's root library.
   for (final vertex in map.values) {
-    if (vertex.library.importUri.isScheme("dart")) {
+    if (vertex == root) {
+      continue;
+    }
+    if (vertex.library.importUri.isScheme("dart") ||
+        visitor.hasEntryPoint(vertex.library)) {
       root.successors.add(vertex);
       vertex.isLoadingRoot = false;
     }
@@ -127,8 +171,12 @@ List<LoadingUnit> computeLoadingUnits(Component component) {
   return loadingUnits.map((u) => u.asLoadingUnit()).toList();
 }
 
-Component transformComponent(Component component) {
-  final metadata = new LoadingUnitsMetadata(computeLoadingUnits(component));
+Component transformComponent(
+    Component component, CoreTypes coreTypes, Target target) {
+  final parser = ConstantPragmaAnnotationParser(coreTypes, target);
+  final visitor = HasEntryPointVisitor(parser);
+  final metadata =
+      new LoadingUnitsMetadata(computeLoadingUnits(component, visitor));
   final repo = new LoadingUnitsMetadataRepository();
   component.addMetadataRepository(repo);
   repo.mapping[component] = metadata;

pkg/vm/test/transformations/deferred_loading_test.dart

@@ -6,6 +6,7 @@ import 'dart:io';
 
 import 'package:kernel/target/targets.dart';
 import 'package:kernel/ast.dart';
+import 'package:kernel/core_types.dart';
 import 'package:kernel/kernel.dart';
 import 'package:test/test.dart';
 import 'package:vm/transformations/deferred_loading.dart'
@@ -24,7 +25,9 @@ runTestCase(Uri source) async {
   final reversed = component.libraries.reversed.toList();
   component.libraries.setAll(0, reversed);
 
-  component = transformComponent(component);
+  final coreTypes = CoreTypes(component);
+
+  component = transformComponent(component, coreTypes, target);
 
   // Remove core libraries so the expected output isn't enormous and broken by
   // core libraries changes.

runtime/vm/app_snapshot.cc

@@ -6916,10 +6916,8 @@ bool Serializer::InCurrentLoadingUnitOrRoot(ObjectPtr obj) {
 
   intptr_t unit_id = heap_->GetLoadingUnit(obj);
   if (unit_id == WeakTable::kNoValue) {
-    // Not found in early assignment. Conservatively choose the root.
-    // TODO(41974): Are these always type testing stubs?
-    unit_id = LoadingUnit::kRootId;
-    heap_->SetLoadingUnit(obj, unit_id);
+    FATAL("Missing loading unit assignment: %s\n",
+          Object::Handle(obj).ToCString());
   }
   return unit_id == LoadingUnit::kRootId || unit_id == current_loading_unit_id_;
 }

runtime/vm/heap/heap.h

@@ -219,9 +219,9 @@ class Heap {
   }
   void ResetObjectIdTable();
 
-  void SetLoadingUnit(ObjectPtr raw_obj, intptr_t object_id) {
+  void SetLoadingUnit(ObjectPtr raw_obj, intptr_t unit_id) {
     ASSERT(Thread::Current()->IsMutatorThread());
-    SetWeakEntry(raw_obj, kLoadingUnits, object_id);
+    SetWeakEntry(raw_obj, kLoadingUnits, unit_id);
   }
   intptr_t GetLoadingUnit(ObjectPtr raw_obj) const {
     ASSERT(Thread::Current()->IsMutatorThread());

runtime/vm/program_visitor.cc

@@ -1429,30 +1429,50 @@ void ProgramVisitor::Dedup(Thread* thread) {
 }
 
 #if defined(DART_PRECOMPILER)
-class AssignLoadingUnitsCodeVisitor : public CodeVisitor {
+class AssignLoadingUnitsCodeVisitor : public ObjectVisitor {
  public:
   explicit AssignLoadingUnitsCodeVisitor(Zone* zone)
       : heap_(Thread::Current()->heap()),
+        code_(Code::Handle(zone)),
         func_(Function::Handle(zone)),
         cls_(Class::Handle(zone)),
         lib_(Library::Handle(zone)),
         unit_(LoadingUnit::Handle(zone)),
         obj_(Object::Handle(zone)) {}
 
+  void VisitObject(ObjectPtr obj) {
+    if (obj->IsCode()) {
+      code_ ^= obj;
+      VisitCode(code_);
+    }
+  }
+
   void VisitCode(const Code& code) {
     intptr_t id;
     if (code.IsFunctionCode()) {
       func_ ^= code.function();
-      cls_ = func_.Owner();
+      obj_ = func_.Owner();
+      cls_ ^= obj_.ptr();
       lib_ = cls_.library();
-      unit_ = lib_.loading_unit();
-      id = unit_.id();
+      if (lib_.IsNull()) {
+        // E.g., dynamic.
+        id = LoadingUnit::kRootId;
+      } else {
+        unit_ = lib_.loading_unit();
+        if (unit_.IsNull()) {
+          return;  // Assignment remains LoadingUnit::kIllegalId
+        }
+        id = unit_.id();
+      }
     } else if (code.IsAllocationStubCode()) {
       cls_ ^= code.owner();
       lib_ = cls_.library();
       unit_ = lib_.loading_unit();
+      if (unit_.IsNull()) {
+        return;  // Assignment remains LoadingUnit::kIllegalId
+      }
       id = unit_.id();
-    } else if (code.IsStubCode()) {
+    } else if (code.IsTypeTestStubCode() || code.IsStubCode()) {
       id = LoadingUnit::kRootId;
     } else {
       UNREACHABLE();
@@ -1484,6 +1504,7 @@ class AssignLoadingUnitsCodeVisitor : public CodeVisitor {
 
  private:
   Heap* heap_;
+  Code& code_;
   Function& func_;
   Class& cls_;
   Library& lib_;
@@ -1493,31 +1514,16 @@ class AssignLoadingUnitsCodeVisitor : public CodeVisitor {
 
 void ProgramVisitor::AssignUnits(Thread* thread) {
   StackZone stack_zone(thread);
-  Zone* zone = stack_zone.GetZone();
-
-  // VM stubs.
-  Instructions& inst = Instructions::Handle(zone);
-  Code& code = Code::Handle(zone);
-  for (intptr_t i = 0; i < StubCode::NumEntries(); i++) {
-    inst = StubCode::EntryAt(i).instructions();
-    thread->heap()->SetLoadingUnit(inst.ptr(), LoadingUnit::kRootId);
-  }
+  Heap* heap = thread->heap();
 
-  // Isolate stubs.
-  ObjectStore* object_store = thread->isolate_group()->object_store();
-  ObjectPtr* from = object_store->from();
-  ObjectPtr* to = object_store->to_snapshot(Snapshot::kFullAOT);
-  for (ObjectPtr* p = from; p <= to; p++) {
-    if ((*p)->IsHeapObject() && (*p)->IsCode()) {
-      code ^= *p;
-      inst = code.instructions();
-      thread->heap()->SetLoadingUnit(inst.ptr(), LoadingUnit::kRootId);
-    }
-  }
+  // Oddballs.
+  heap->SetLoadingUnit(Object::null(), LoadingUnit::kRootId);
+  heap->SetLoadingUnit(Object::empty_object_pool().ptr(), LoadingUnit::kRootId);
 
-  // Function code / allocation stubs.
-  AssignLoadingUnitsCodeVisitor visitor(zone);
-  WalkProgram(zone, thread->isolate_group(), &visitor);
+  AssignLoadingUnitsCodeVisitor visitor(thread->zone());
+  HeapIterationScope iter(thread);
+  iter.IterateVMIsolateObjects(&visitor);
+  iter.IterateObjects(&visitor);
 }
 
 class ProgramHashVisitor : public CodeVisitor {

tests/language/deferred/exported_main_lib.dart

@@ -0,0 +1,7 @@
+// Copyright (c) 2022, 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.
+
+main() {
+  print("Usurper to the role of root library!");
+}

tests/language/deferred/exported_main_test.dart

@@ -0,0 +1,14 @@
+// Copyright (c) 2022, 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.
+
+export "exported_main_lib.dart";
+
+@pragma("vm:entry-point")
+foo() {
+  // The frontend will not recognize that this is the root library. Though this
+  // library is not reachable from what the frontend considers the root library,
+  // the entry point pragma will still cause this function to compiled by
+  // gen_snaption, so it is important that this library is assigned to a
+  // loading unit.
+}

tests/language/deferred/loading_unit_root_has_only_export_lib1.dart

@@ -0,0 +1,5 @@
+// Copyright (c) 2022, 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.
+
+export "loading_unit_root_has_only_export_lib2.dart";

tests/language/deferred/loading_unit_root_has_only_export_lib2.dart

@@ -0,0 +1,5 @@
+// Copyright (c) 2022, 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.
+
+foo() => "in lib2";

tests/language/deferred/loading_unit_root_has_only_export_test.dart

@@ -0,0 +1,11 @@
+// Copyright (c) 2022, 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:expect/expect.dart";
+import "loading_unit_root_has_only_export_lib1.dart" deferred as lib;
+
+main() async {
+  await lib.loadLibrary();
+  Expect.equals(lib.foo(), "in lib2");
+}

tests/language_2/deferred/exported_main_lib.dart

@@ -0,0 +1,9 @@
+// Copyright (c) 2022, 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.
+
+// @dart = 2.9
+
+main() {
+  print("Usurper to the role of root library!");
+}

tests/language_2/deferred/exported_main_test.dart

@@ -0,0 +1,16 @@
+// Copyright (c) 2022, 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.
+
+// @dart = 2.9
+
+export "exported_main_lib.dart";
+
+@pragma("vm:entry-point")
+foo() {
+  // The frontend will not recognize that this is the root library. Though this
+  // library is not reachable from what the frontend considers the root library,
+  // the entry point pragma will still cause this function to compiled by
+  // gen_snaption, so it is important that this library is assigned to a
+  // loading unit.
+}

tests/language_2/deferred/loading_unit_root_has_only_export_lib1.dart

@@ -0,0 +1,7 @@
+// Copyright (c) 2022, 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.
+
+// @dart = 2.9
+
+export "loading_unit_root_has_only_export_lib2.dart";

tests/language_2/deferred/loading_unit_root_has_only_export_lib2.dart

@@ -0,0 +1,7 @@
+// Copyright (c) 2022, 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.
+
+// @dart = 2.9
+
+foo() => "in lib2";

tests/language_2/deferred/loading_unit_root_has_only_export_test.dart

@@ -0,0 +1,13 @@
+// Copyright (c) 2022, 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.
+
+// @dart = 2.9
+
+import "package:expect/expect.dart";
+import "loading_unit_root_has_only_export_lib1.dart" deferred as lib;
+
+main() async {
+  await lib.loadLibrary();
+  Expect.equals(lib.foo(), "in lib2");
+}