[dart2wasm] Retain CFE inserted deferred loading nodes for TFA

TFA has handling of the `LoadLibrary` & `CheckLibraryIsLoaded` nodes. It
does logic based on them. For example it will ensure to that if one of
these nodes is retained in the AST that we also retain the deferred
library itself as well as the deferred library import.

For now dart2wasm lowered those two nodes during modular
transformations. This means TFA didn't see those nodes and weird things
can happen (e.g. a `checkLibraryLoad()` call without that library
existing anymore). Other interesting things that can happen: TFA
sees only one call to `checklibraryLoad` and constant propagates
the argument into the body & changes signature, ...

We could keep the existing behavior and make TFA aware of the dart2wasm
lowering (directly which is very hacky or via e.g. `Target`
indirection). Though this seems too complex.

So instead of making TFA aware of the lowerings, we move them to be
after TFA. Though we want precise TFA results if we don't use deferred
loading. So the approach is: If we actually need the runtime functions
we'll inject `@pragma('wasm:entry-point')` annotations before running
TFA and then lower those nodes after TFA.

This also means we avoid the ugly CFE nodes lowering followed by
lowering of the lowering (for load ids). Instead there's one place where
we lower those two nodes and we choose either to lower to "load id"
methods or the normal methods.

This is part of moving constants to deferred modules.

Issue #61727

Change-Id: Icebaf5a9495e00a8f85ecbf161b6ef891ea179ea
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/457700
Reviewed-by: Nate Biggs <[email protected]>
Commit-Queue: Martin Kustermann <[email protected]>

Author: mkustermann

pkg/dart2wasm/lib/code_generator.dart

@@ -54,13 +54,9 @@ abstract class AstCodeGenerator
     extends ExpressionVisitor1<w.ValueType, w.ValueType>
     with
         ExpressionVisitor1DefaultMixin<w.ValueType, w.ValueType>,
-        ExpressionVisitor1ExperimentExclusionMixin<w.ValueType,
-            w.ValueType>,
+        ExpressionVisitor1ExperimentExclusionMixin<w.ValueType, w.ValueType>,
         StatementVisitorExperimentExclusionMixin<void>
-    implements
-        InitializerVisitor<void>,
-        StatementVisitor<void>,
-        CodeGenerator {
+    implements InitializerVisitor<void>, StatementVisitor<void>, CodeGenerator {
   final Translator translator;
   final w.FunctionType functionType;
   final Member enclosingMember;

pkg/dart2wasm/lib/compile.dart

@@ -584,7 +584,7 @@ ModuleStrategy _createModuleStrategy(
     return DynamicSubmoduleStrategy(
         component, options, target, coreTypes, dynamicMainModuleUri!);
   }
-  return DefaultModuleStrategy(component, options);
+  return DefaultModuleStrategy(coreTypes, component, options);
 }
 
 // Patches `dart:_internal`s `mainTearOff{0,1,2}` getters.

pkg/dart2wasm/lib/deferred_loading.dart

@@ -18,6 +18,7 @@ import 'compiler_options.dart';
 import 'generate_wasm.dart';
 import 'modules.dart';
 import 'target.dart';
+import 'util.dart' show addPragma;
 
 /// The root of a deferred import subgraph.
 ///
@@ -91,15 +92,21 @@ class DeferredLoadingModuleStrategy extends ModuleStrategy {
   DeferredLoadingModuleStrategy(
       this.component, this.options, this.kernelTarget, this.coreTypes);
 
+  late final deferredLoweringTransformer = options.loadsIdsUri != null
+      ? _DeferredLoadingLoweringLoadIds(component, coreTypes)
+      : DeferredLoadingLoweringUris(coreTypes);
+
   @override
   void prepareComponent() {
-    if (options.loadsIdsUri != null) {
-      component.accept(_DeferredLoadingLoadIdTransformer(component, coreTypes));
-    }
+    // Since we've deferred module loading enabled, we mark the necessary
+    // runtime methods as entrypoints for TFA.
+    deferredLoweringTransformer.markRuntimeFunctionsAsEntrypoints();
   }
 
   @override
   Future<void> processComponentAfterTfa() async {
+    component.accept(deferredLoweringTransformer);
+
     final (libraryToRootSet, importTargetMap) = _buildLibraryToImports();
 
     final builder = ModuleMetadataBuilder(options);
@@ -229,6 +236,10 @@ class StressTestModuleStrategy extends ModuleStrategy {
   final WasmCompilerOptions options;
   late final ModuleOutputData moduleOutputData;
 
+  late final deferredLoweringTransformer = options.loadsIdsUri != null
+      ? _DeferredLoadingLoweringLoadIds(component, coreTypes)
+      : DeferredLoadingLoweringUris(coreTypes);
+
   /// We load all 'dart:*' libraries since just doing the deferred load of modules
   /// requires a significant portion of the SDK libraries.
   late final Set<Library> _testModeMainLibraries = {
@@ -283,13 +294,15 @@ class StressTestModuleStrategy extends ModuleStrategy {
     await_transformer.transformLibraries(
         [invokeMain.enclosingLibrary], classHierarchy, coreTypes);
 
-    if (options.loadsIdsUri != null) {
-      component.accept(_DeferredLoadingLoadIdTransformer(component, coreTypes));
-    }
+    // Since we've deferred module loading enabled, we mark the necessary
+    // runtime methods as entrypoints for TFA.
+    deferredLoweringTransformer.markRuntimeFunctionsAsEntrypoints();
   }
 
   @override
   Future<void> processComponentAfterTfa() async {
+    component.accept(deferredLoweringTransformer);
+
     final moduleBuilder = ModuleMetadataBuilder(options);
     final mainModule = moduleBuilder.buildModuleMetadata();
     final initLibraries = _testModeMainLibraries;
@@ -409,27 +422,90 @@ String _generateDeferredMapJson(Component component, Uri rootLibraryUri,
   return const JsonEncoder.withIndent('  ').convert(output);
 }
 
-class _DeferredLoadingLoadIdTransformer extends Transformer {
-  final Procedure? _loadLibrary;
-  final Procedure? _checkLibraryIsLoaded;
+abstract class DeferredLoadingLoweringBase extends Transformer {
+  final CoreTypes coreTypes;
+
+  DeferredLoadingLoweringBase(this.coreTypes);
+
+  void markRuntimeFunctionsAsEntrypoints();
+
+  void addEntryPointPragma(Procedure node) {
+    addPragma(node, 'wasm:entry-point', coreTypes, value: BoolConstant(true));
+  }
+}
+
+class DeferredLoadingLoweringUris extends DeferredLoadingLoweringBase {
+  final Procedure _loadLibrary;
+  final Procedure _checkLibraryIsLoaded;
+
+  DeferredLoadingLoweringUris(super.coreTypes)
+      : _loadLibrary = coreTypes.index
+            .getTopLevelProcedure('dart:_internal', 'loadLibrary'),
+        _checkLibraryIsLoaded = coreTypes.index
+            .getTopLevelProcedure('dart:_internal', 'checkLibraryIsLoaded');
+
+  @override
+  void markRuntimeFunctionsAsEntrypoints() {
+    addEntryPointPragma(_loadLibrary);
+    addEntryPointPragma(_checkLibraryIsLoaded);
+  }
+
+  @override
+  TreeNode visitLibrary(Library node) {
+    for (final dep in node.dependencies) {
+      if (dep.isDeferred) {
+        return super.visitLibrary(node);
+      }
+    }
+    // No need to transform libraries without deferred imports.
+    return node;
+  }
+
+  @override
+  TreeNode visitLoadLibrary(LoadLibrary node) {
+    final import = node.import;
+    return StaticInvocation(
+        _loadLibrary,
+        Arguments([
+          StringLiteral('${import.enclosingLibrary.importUri}'),
+          StringLiteral(import.name!)
+        ]));
+  }
+
+  @override
+  TreeNode visitCheckLibraryIsLoaded(CheckLibraryIsLoaded node) {
+    final import = node.import;
+    return StaticInvocation(
+        _checkLibraryIsLoaded,
+        Arguments([
+          StringLiteral('${import.enclosingLibrary.importUri}'),
+          StringLiteral(import.name!)
+        ]));
+  }
+}
+
+class _DeferredLoadingLoweringLoadIds extends DeferredLoadingLoweringBase {
   final Procedure _loadLibraryFromLoadId;
   final Procedure _checkLibraryIsLoadedFromLoadId;
+
   final LoadIdRepository _loadIdRepository = LoadIdRepository();
   Map<String, int> _libraryLoadIds = {};
   int _loadIdCounter = 1;
 
-  _DeferredLoadingLoadIdTransformer(Component component, CoreTypes coreTypes)
-      : _loadLibrary = coreTypes.index.tryGetProcedure(
-            'dart:_internal', LibraryIndex.topLevel, 'loadLibrary'),
-        _checkLibraryIsLoaded = coreTypes.index.tryGetProcedure(
-            'dart:_internal', LibraryIndex.topLevel, 'checkLibraryIsLoaded'),
-        _loadLibraryFromLoadId = coreTypes.index
+  _DeferredLoadingLoweringLoadIds(Component component, super.coreTypes)
+      : _loadLibraryFromLoadId = coreTypes.index
             .getTopLevelProcedure('dart:_internal', 'loadLibraryFromLoadId'),
         _checkLibraryIsLoadedFromLoadId = coreTypes.index.getTopLevelProcedure(
             'dart:_internal', 'checkLibraryIsLoadedFromLoadId') {
     component.addMetadataRepository(_loadIdRepository);
   }
 
+  @override
+  void markRuntimeFunctionsAsEntrypoints() {
+    addEntryPointPragma(_loadLibraryFromLoadId);
+    addEntryPointPragma(_checkLibraryIsLoadedFromLoadId);
+  }
+
   @override
   TreeNode visitLibrary(Library node) {
     // Assign a load ID to each deferred import.
@@ -446,17 +522,19 @@ class _DeferredLoadingLoadIdTransformer extends Transformer {
   }
 
   @override
-  TreeNode visitStaticInvocation(StaticInvocation node) {
-    if (node.target != _loadLibrary && node.target != _checkLibraryIsLoaded) {
-      return super.visitStaticInvocation(node);
-    }
-    final [_, importPrefix as StringLiteral] = node.arguments.positional;
-    final loadId = '${_libraryLoadIds[importPrefix.value]!}';
+  TreeNode visitLoadLibrary(LoadLibrary node) {
+    final import = node.import;
+    final loadId = '${_libraryLoadIds[import.name!]!}';
+    return StaticInvocation(
+        _loadLibraryFromLoadId, Arguments([StringLiteral(loadId)]));
+  }
+
+  @override
+  TreeNode visitCheckLibraryIsLoaded(CheckLibraryIsLoaded node) {
+    final import = node.import;
+    final loadId = '${_libraryLoadIds[import.name!]!}';
     return StaticInvocation(
-        node.target == _loadLibrary
-            ? _loadLibraryFromLoadId
-            : _checkLibraryIsLoadedFromLoadId,
-        Arguments([StringLiteral(loadId)]));
+        _checkLibraryIsLoadedFromLoadId, Arguments([StringLiteral(loadId)]));
   }
 }
 

pkg/dart2wasm/lib/dynamic_modules.dart

@@ -20,6 +20,7 @@ import 'class_info.dart';
 import 'code_generator.dart';
 import 'compiler_options.dart';
 import 'constants.dart' show maxArrayNewFixedLength;
+import 'deferred_loading.dart' show DeferredLoadingLoweringUris;
 import 'dispatch_table.dart';
 import 'dynamic_module_kernel_metadata.dart';
 import 'intrinsics.dart' show MemberIntrinsic;
@@ -117,6 +118,9 @@ class DynamicMainModuleStrategy extends ModuleStrategy with KernelNodes {
       this.dynamicInterfaceSpecificationBaseUri)
       : index = coreTypes.index;
 
+  late final deferredLoweringTransformer =
+      DeferredLoadingLoweringUris(coreTypes);
+
   @override
   void prepareComponent() {
     // Annotate the kernel with info from dynamic interface.
@@ -131,6 +135,8 @@ class DynamicMainModuleStrategy extends ModuleStrategy with KernelNodes {
 
     component.addMetadataRepository(DynamicModuleConstantRepository());
     component.addMetadataRepository(DynamicModuleGlobalIdRepository());
+
+    deferredLoweringTransformer.markRuntimeFunctionsAsEntrypoints();
   }
 
   @override
@@ -202,7 +208,9 @@ class DynamicMainModuleStrategy extends ModuleStrategy with KernelNodes {
   }
 
   @override
-  Future<void> processComponentAfterTfa() async {}
+  Future<void> processComponentAfterTfa() async {
+    component.accept(deferredLoweringTransformer);
+  }
 }
 
 class DynamicSubmoduleStrategy extends ModuleStrategy {

pkg/dart2wasm/lib/modules.dart

@@ -7,6 +7,7 @@ import 'package:kernel/core_types.dart';
 import 'package:path/path.dart' as path;
 
 import 'compiler_options.dart';
+import 'deferred_loading.dart' show DeferredLoadingLoweringUris;
 import 'target.dart';
 import 'util.dart';
 
@@ -101,10 +102,14 @@ class ModuleOutputData {
 
 /// Module strategy that puts all libraries into a single module.
 class DefaultModuleStrategy extends ModuleStrategy {
+  final CoreTypes coreTypes;
   final Component component;
   final WasmCompilerOptions options;
 
-  DefaultModuleStrategy(this.component, this.options);
+  DefaultModuleStrategy(this.coreTypes, this.component, this.options);
+
+  late final deferredLoweringTransformer =
+      DeferredLoadingLoweringUris(coreTypes);
 
   @override
   ModuleOutputData buildModuleOutputData() {
@@ -117,10 +122,14 @@ class DefaultModuleStrategy extends ModuleStrategy {
   }
 
   @override
-  void prepareComponent() {}
+  void prepareComponent() {
+    deferredLoweringTransformer.markRuntimeFunctionsAsEntrypoints();
+  }
 
   @override
-  Future<void> processComponentAfterTfa() async {}
+  Future<void> processComponentAfterTfa() async {
+    component.accept(deferredLoweringTransformer);
+  }
 }
 
 bool _hasWasmExportPragma(CoreTypes coreTypes, Member m) =>

pkg/dart2wasm/lib/transformers.dart

@@ -57,9 +57,6 @@ class _WasmTransformer extends Transformer {
   final Procedure _trySetStackTraceForwarder;
   final Procedure _trySetStackTrace;
 
-  final Procedure _loadLibrary;
-  final Procedure _checkLibraryIsLoaded;
-
   final List<_AsyncStarFrame> _asyncStarFrames = [];
   bool _enclosingIsAsyncStar = false;
 
@@ -164,10 +161,6 @@ class _WasmTransformer extends Transformer {
             .getTopLevelProcedure('dart:async', '_trySetStackTrace'),
         _trySetStackTrace = coreTypes.index
             .getProcedure('dart:core', 'Error', '_trySetStackTrace'),
-        _loadLibrary = coreTypes.index
-            .getTopLevelProcedure("dart:_internal", "loadLibrary"),
-        _checkLibraryIsLoaded = coreTypes.index
-            .getTopLevelProcedure("dart:_internal", "checkLibraryIsLoaded"),
         _factorySpecializer = FactorySpecializer(coreTypes),
         _pushPopWasmArrayTransformer = PushPopWasmArrayTransformer(coreTypes);
 
@@ -838,30 +831,6 @@ class _WasmTransformer extends Transformer {
     return node.receiver;
   }
 
-  @override
-  TreeNode visitLoadLibrary(LoadLibrary node) {
-    node.transformChildren(this);
-    final import = node.import;
-    return StaticInvocation(
-        _loadLibrary,
-        Arguments([
-          StringLiteral('${import.enclosingLibrary.importUri}'),
-          StringLiteral(import.name!)
-        ]));
-  }
-
-  @override
-  TreeNode visitCheckLibraryIsLoaded(CheckLibraryIsLoaded node) {
-    node.transformChildren(this);
-    final import = node.import;
-    return StaticInvocation(
-        _checkLibraryIsLoaded,
-        Arguments([
-          StringLiteral('${import.enclosingLibrary.importUri}'),
-          StringLiteral(import.name!)
-        ]));
-  }
-
   @override
   TreeNode visitThrow(Throw node) {
     node.transformChildren(this);