[dart2wasm] Modify wasm_builder to support multi-module deferred loading (2/2).

Golem results verifying no regressions from this change:
https://golem.corp.goog/Revision?repository=dart&revision=111273&patch=19368

This modifies the wasm_builder to emit types in minimally sized recursive groups.

The main motivation of these changes is to minimize the size of the type section when dart2wasm produces multiple modules. If we included every type in every module, for small modules the type section would be a very significant % of the total file size (in a small example it was ~98% for a deferred module).

The two main changes are:
1) Type tree-shaking for each module.

We detect which types are actually used in that module (via instructions, function signatures, etc.) and only include the detected types in the module's type section. This simplifies the compilation pipeline as we don't need to worry about assigning types to modules, we can just build the module as before and post-process the IR to collect the set of used types.

2) Create minimally-sized recursive groups.

In order for wasm type-checking to work across modules, equal types have to be in equal recursive groups. Tree-shaking therefore has to occur at the rec group level as opposed to the individual type level. So to make tree-shaking effective, we need the smallest possible rec groups.

We achieve this by creating a graph of the wasm types' dependencies and then calculating the set of strongly connected components for that graph. Each component represents a minimally-sized recursive group. The DAG formed by the components is the order to emit them so that definitions come before usages.

Importantly, by separating types into different rec groups, we are also changing the equivalence relationship between them. This can have a meaningful impact on binaryen's ability to optimize the module, two types that were distinguishable might not be anymore. To avoid this regression we group together any types that are structurally equivalent. This way binaryen will differentiate them as separate types as they were in the original Dart source.


Change-Id: I67acdd21a89ff2718e8bbd6360f342c150494a9a
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/378764
Reviewed-by: Ömer Ağacan <omersa@google.com>
Reviewed-by: Martin Kustermann <kustermann@google.com>
Commit-Queue: Nate Biggs <natebiggs@google.com>

pkg/dart2wasm/lib/functions.dart

@@ -54,14 +54,9 @@ class FunctionCollector {
         String module = importName.substring(0, dot);
         String name = importName.substring(dot + 1);
         if (member is Procedure) {
-          // Define the function type in a singular recursion group to enable it
-          // to be unified with function types defined in FFI modules or using
-          // `WebAssembly.Function`.
-          m.types.splitRecursionGroup();
           w.FunctionType ftype = _makeFunctionType(
               translator, member.reference, null,
               isImportOrExport: true);
-          m.types.splitRecursionGroup();
           _functions[member.reference] =
               m.functions.import(module, name, ftype, "$importName (import)");
         }
@@ -71,15 +66,8 @@ class FunctionCollector {
         translator.getPragma(member, "wasm:export", member.name.text);
     if (exportName != null) {
       if (member is Procedure) {
-        // Although we don't need type unification for the types of exported
-        // functions, we still place these types in singleton recursion groups,
-        // since Binaryen's `--closed-world` optimization mode requires all
-        // publicly exposed types to be defined in separate recursion groups
-        // from GC types.
-        m.types.splitRecursionGroup();
         _makeFunctionType(translator, member.reference, null,
             isImportOrExport: true);
-        m.types.splitRecursionGroup();
       }
       addExport(member.reference, exportName);
     }

pkg/wasm_builder/lib/src/builder/functions.dart

@@ -29,9 +29,19 @@ class FunctionsBuilder with Builder<ir.Functions> {
     }
   }
 
+  void collectUsedTypes(Set<ir.DefType> usedTypes) {
+    for (final f in _functionBuilders) {
+      usedTypes.add(f.type);
+      f.body.collectUsedTypes(usedTypes);
+    }
+    for (final f in _importedFunctions) {
+      usedTypes.add(f.type);
+    }
+  }
+
   /// Defines a new function in this module with the given function type.
   ///
-  /// The [DefinedFunction.body] must be completed (including the terminating
+  /// The [ir.DefinedFunction.body] must be completed (including the terminating
   /// `end`) before the module can be serialized.
   FunctionBuilder define(ir.FunctionType type, [String? name]) {
     final function =

pkg/wasm_builder/lib/src/builder/globals.dart

@@ -18,6 +18,18 @@ class GlobalsBuilder with Builder<ir.Globals> {
 
   GlobalsBuilder(this._module);
 
+  void collectUsedTypes(Set<ir.DefType> usedTypes) {
+    for (final global in _globalBuilders) {
+      final defType = global.type.type.containedDefType;
+      if (defType != null) usedTypes.add(defType);
+      global.initializer.collectUsedTypes(usedTypes);
+    }
+    for (final global in _importedGlobals) {
+      final defType = global.type.type.containedDefType;
+      if (defType != null) usedTypes.add(defType);
+    }
+  }
+
   /// Defines a new global variable in this module.
   GlobalBuilder define(ir.GlobalType type, [String? name]) {
     final global = GlobalBuilder(_module, ir.FinalizableIndex(), type, name);

pkg/wasm_builder/lib/src/builder/instructions.dart

@@ -165,6 +165,20 @@ class InstructionsBuilder with Builder<ir.Instructions> {
 
   bool get recordSourceMaps => _sourceMappings != null;
 
+  void collectUsedTypes(Set<ir.DefType> usedTypes) {
+    for (final local in locals) {
+      final localDefType = local.type.containedDefType;
+      if (localDefType != null) usedTypes.add(localDefType);
+    }
+    for (final instruction in _instructions) {
+      usedTypes.addAll(instruction.usedDefTypes);
+      for (final valueType in instruction.usedValueTypes) {
+        final type = valueType.containedDefType;
+        if (type != null) usedTypes.add(type);
+      }
+    }
+  }
+
   @override
   ir.Instructions forceBuild() => ir.Instructions(
       locals, _instructions, _stackTraces, _traceLines, _sourceMappings);

pkg/wasm_builder/lib/src/builder/module.dart

@@ -10,7 +10,7 @@ import 'builder.dart';
 class ModuleBuilder with Builder<ir.Module> {
   final Uri? sourceMapUrl;
   final List<int> watchPoints;
-  final types = TypesBuilder();
+  late final types = TypesBuilder(this);
   late final functions = FunctionsBuilder(this);
   final tables = TablesBuilder();
   final memories = MemoriesBuilder();

pkg/wasm_builder/lib/src/builder/tags.dart

@@ -13,10 +13,14 @@ class TagsBuilder with Builder<ir.Tags> {
 
   TagsBuilder();
 
-  Set<ir.DefType> usedTypes() => {
-        for (final tag in _defined) tag.type,
-        for (final tag in _imported) tag.type,
-      };
+  void collectUsedTypes(Set<ir.DefType> usedTypes) {
+    for (final tag in _defined) {
+      usedTypes.add(tag.type);
+    }
+    for (final tag in _imported) {
+      usedTypes.add(tag.type);
+    }
+  }
 
   /// Defines a new tag in the module.
   ir.Tag define(ir.FunctionType type) {

pkg/wasm_builder/lib/src/builder/types.dart

@@ -2,14 +2,142 @@
 // 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:collection/collection.dart';
 import '../ir/ir.dart' as ir;
 import 'builder.dart';
 
+/// Creates minimally sized rec groups given the types added to the builder.
+///
+/// The set of [ir.DefType]s defined in the program create a graph. Edges in
+/// this graph are either supertype relationships or usage within the structure
+/// of the type (e.g. a field in a struct or an input in a function).
+///
+/// Given this graph each rec group is a strongly connected component in the
+/// graph. That is, two types should be in the same rec group if they are part
+/// of a cycle in the graph.
+///
+/// The strongly connected components then form a DAG which defines the order
+/// of our rec groups. The root of the DAG is the top type and all other types
+/// follow from there.
+///
+/// Two Dart classes may map to structurally equivalent [ir.StructType]s. Having
+/// them in the same rec group disambiguates them for the wasm type system. In
+/// order to allow binaryen to optimize these class structs as independent types
+/// we put structurally equivalent struct types into the same recursive group as
+/// well.
+class _RecGroupBuilder {
+  late final List<List<ir.DefType>> _allRecursiveGroups =
+      _createAllRecursiveGroups();
+  final List<ir.DefType> _allDefinedTypes = [];
+
+  _RecGroupBuilder();
+
+  static bool _areStructurallyEqual(ir.DefType type1, ir.DefType type2) {
+    return type1.isStructuralSubtypeOf(type2) &&
+        type2.isStructuralSubtypeOf(type1);
+  }
+
+  /// Get the out-edges for [type] in the wasm type graph.
+  Set<ir.DefType> _edgesforType(ir.DefType type) {
+    final edges = <ir.DefType>{};
+    for (final constituentType in type.constituentTypes) {
+      if (constituentType is ir.RefType) {
+        final heapType = constituentType.heapType;
+        if (heapType is ir.DefType) {
+          edges.add(heapType);
+        }
+      }
+    }
+    final superType = type.superType;
+    if (superType != null) {
+      edges.add(superType);
+    }
+
+    bool isStructurallyEqual(ir.DefType t) => _areStructurallyEqual(type, t);
+
+    edges.addAll(_allDefinedTypes.where(isStructurallyEqual));
+    return edges;
+  }
+
+  /// Create minimal recursive groups.
+  ///
+  /// Some groups may contain only unused types. Those groups can be dropped by
+  /// callers of this function.
+  List<List<ir.DefType>> _createAllRecursiveGroups() {
+    Map<ir.DefType, Set<ir.DefType>> typeGraph = {};
+    for (ir.DefType type in _allDefinedTypes) {
+      typeGraph[type] = _edgesforType(type);
+    }
+    final components = stronglyConnectedComponents(typeGraph);
+    final groups = <List<ir.DefType>>[];
+    // Make sure to reverse the list since the components are returned with the
+    // leaves first.
+    for (final component in components.reversed) {
+      final group = <ir.DefType>[];
+      final added = <ir.DefType>{};
+      void addToGroup(ir.DefType type) {
+        if (!added.add(type)) return;
+        if (component.contains(type.superType)) {
+          // Supertypes must be added before their subtypes.
+          addToGroup(type.superType!);
+        }
+        group.add(type);
+      }
+
+      for (final type in component) {
+        addToGroup(type);
+      }
+      groups.add(group);
+    }
+    return groups;
+  }
+
+  void addDefinedType(ir.DefType type) {
+    _allDefinedTypes.add(type);
+  }
+
+  /// Create a filtered list of recursive groups in type hierarchy order.
+  ///
+  /// [directlyUsedTypes] should be the list of types used directly in the
+  /// module (i.e. types referenced from code, function definitions, etc.).
+  /// Indirectly used types are referenced transitively from directly used
+  /// types (i.e. fields of structs, function inputs, etc.).
+  ///
+  /// The returned list includes all rec groups that contain a directly or
+  /// indirectly used type.
+  List<List<ir.DefType>> createGroupsForModule(
+      Set<ir.DefType> directlyUsedTypes) {
+    final allUsedTypes = {...directlyUsedTypes};
+
+    void addUsedType(ir.DefType type) {
+      // Visit all the children of type and include them as well.
+      for (final edge in _edgesforType(type)) {
+        if (!allUsedTypes.add(edge)) continue;
+        addUsedType(edge);
+      }
+    }
+
+    for (final type in directlyUsedTypes) {
+      addUsedType(type);
+    }
+
+    final usedGroups = <List<ir.DefType>>[];
+    for (final group in _allRecursiveGroups) {
+      if (group.any((type) => allUsedTypes.contains(type))) {
+        usedGroups.add(group);
+      }
+    }
+    return usedGroups;
+  }
+}
+
 class TypesBuilder with Builder<ir.Types> {
-  final _recursionGroupSplits = <int>[];
-  final _functionTypeMap = <_FunctionTypeKey, ir.FunctionType>{};
-  final _defTypes = <ir.DefType>[];
-  int _nameCount = 0;
+  final ModuleBuilder _module;
+
+  late final Map<_FunctionTypeKey, ir.FunctionType> _functionTypeMap = {};
+  late final _RecGroupBuilder _recGroupBuilder = _RecGroupBuilder();
+
+  TypesBuilder(this._module);
 
   /// Add a new function type to the module.
   ///
@@ -27,9 +155,8 @@ class TypesBuilder with Builder<ir.Types> {
     final List<ir.ValueType> outputList = List.unmodifiable(outputs);
     final _FunctionTypeKey key = _FunctionTypeKey(inputList, outputList);
     return _functionTypeMap.putIfAbsent(key, () {
-      final type = ir.FunctionType(inputList, outputList, superType: superType)
-        ..index = _defTypes.length;
-      _defTypes.add(type);
+      final type = ir.FunctionType(inputList, outputList, superType: superType);
+      _recGroupBuilder.addDefinedType(type);
       return type;
     });
   }
@@ -40,10 +167,8 @@ class TypesBuilder with Builder<ir.Types> {
   /// struct types to be recursive.
   ir.StructType defineStruct(String name,
       {Iterable<ir.FieldType>? fields, ir.DefType? superType}) {
-    final type = ir.StructType(name, fields: fields, superType: superType)
-      ..index = _defTypes.length;
-    _defTypes.add(type);
-    _nameCount++;
+    final type = ir.StructType(name, fields: fields, superType: superType);
+    _recGroupBuilder.addDefinedType(type);
     return type;
   }
 
@@ -54,27 +179,27 @@ class TypesBuilder with Builder<ir.Types> {
   ir.ArrayType defineArray(String name,
       {ir.FieldType? elementType, ir.DefType? superType}) {
     final type =
-        ir.ArrayType(name, elementType: elementType, superType: superType)
-          ..index = _defTypes.length;
-    _defTypes.add(type);
-    _nameCount++;
+        ir.ArrayType(name, elementType: elementType, superType: superType);
+    _recGroupBuilder.addDefinedType(type);
     return type;
   }
 
-  /// Insert a recursion group split in the list of type definitions. Types can
-  /// only reference other types in the same or earlier recursion groups.
-  void splitRecursionGroup() {
-    int typeCount = _defTypes.length;
-    if (typeCount > 0 &&
-        (_recursionGroupSplits.isEmpty ||
-            _recursionGroupSplits.last != typeCount)) {
-      _recursionGroupSplits.add(typeCount);
-    }
+  Set<ir.DefType> _collectUsedTypes() {
+    final usedTypes = <ir.DefType>{};
+    _module.functions.collectUsedTypes(usedTypes);
+    _module.globals.collectUsedTypes(usedTypes);
+    _module.tags.collectUsedTypes(usedTypes);
+    return usedTypes;
   }
 
   @override
-  ir.Types forceBuild() =>
-      ir.Types(_defTypes, _recursionGroupSplits, _nameCount);
+  ir.Types forceBuild() {
+    final usedTypes = _collectUsedTypes();
+    final types = _recGroupBuilder.createGroupsForModule(usedTypes);
+    final nameCount =
+        types.fold(0, (p, e) => p + e.whereType<ir.DataType>().length);
+    return ir.Types(types, nameCount);
+  }
 }
 
 class _FunctionTypeKey {

pkg/wasm_builder/lib/src/ir/instruction.dart

@@ -4,9 +4,21 @@
 
 part of 'instructions.dart';
 
-abstract class Instruction implements Serializable {}
+abstract class Instruction implements Serializable {
+  /// The [ValueType] types referenced by this instruction. Used to determine
+  /// which types need to be included in the module. Unused types will not be
+  /// emitted in the wasm output.
+  List<ValueType> get usedValueTypes => const [];
 
-abstract class SingleByteInstruction implements Instruction {
+  /// The [DefType] types referenced by this instruction. Used to determine
+  /// which types need to be included in the module. Unused types will not be
+  /// emitted in the wasm output.
+  List<DefType> get usedDefTypes => const [];
+
+  const Instruction();
+}
+
+abstract class SingleByteInstruction extends Instruction {
   final int byte;
 
   const SingleByteInstruction(this.byte);
@@ -15,7 +27,7 @@ abstract class SingleByteInstruction implements Instruction {
   void serialize(Serializer s) => s.writeByte(byte);
 }
 
-abstract class MultiByteInstruction implements Instruction {
+abstract class MultiByteInstruction extends Instruction {
   final List<int> bytes;
 
   const MultiByteInstruction(this.bytes);
@@ -36,9 +48,12 @@ class BeginNoEffectBlock extends MultiByteInstruction {
   const BeginNoEffectBlock() : super(const [0x02, 0x40]);
 }
 
-class BeginOneOutputBlock implements Instruction {
+class BeginOneOutputBlock extends Instruction {
   final ValueType type;
 
+  @override
+  List<ValueType> get usedValueTypes => [type];
+
   BeginOneOutputBlock(this.type);
 
   @override
@@ -48,15 +63,18 @@ class BeginOneOutputBlock implements Instruction {
   }
 }
 
-class BeginFunctionBlock implements Instruction {
+class BeginFunctionBlock extends Instruction {
   final FunctionType type;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   BeginFunctionBlock(this.type);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x02);
-    s.writeSigned(type.index);
+    s.write(type);
   }
 }
 
@@ -64,9 +82,12 @@ class BeginNoEffectLoop extends MultiByteInstruction {
   const BeginNoEffectLoop() : super(const [0x03, 0x40]);
 }
 
-class BeginOneOutputLoop implements Instruction {
+class BeginOneOutputLoop extends Instruction {
   final ValueType type;
 
+  @override
+  List<ValueType> get usedValueTypes => [type];
+
   BeginOneOutputLoop(this.type);
 
   @override
@@ -76,15 +97,18 @@ class BeginOneOutputLoop implements Instruction {
   }
 }
 
-class BeginFunctionLoop implements Instruction {
+class BeginFunctionLoop extends Instruction {
   final FunctionType type;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   BeginFunctionLoop(this.type);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x03);
-    s.writeSigned(type.index);
+    s.write(type);
   }
 }
 
@@ -92,9 +116,12 @@ class BeginNoEffectIf extends MultiByteInstruction {
   const BeginNoEffectIf() : super(const [0x04, 0x40]);
 }
 
-class BeginOneOutputIf implements Instruction {
+class BeginOneOutputIf extends Instruction {
   final ValueType type;
 
+  @override
+  List<ValueType> get usedValueTypes => [type];
+
   BeginOneOutputIf(this.type);
 
   @override
@@ -104,15 +131,18 @@ class BeginOneOutputIf implements Instruction {
   }
 }
 
-class BeginFunctionIf implements Instruction {
+class BeginFunctionIf extends Instruction {
   final FunctionType type;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   BeginFunctionIf(this.type);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x04);
-    s.writeSigned(type.index);
+    s.write(type);
   }
 }
 
@@ -124,9 +154,12 @@ class BeginNoEffectTry extends MultiByteInstruction {
   const BeginNoEffectTry() : super(const [0x06, 0x40]);
 }
 
-class BeginOneOutputTry implements Instruction {
+class BeginOneOutputTry extends Instruction {
   final ValueType type;
 
+  @override
+  List<ValueType> get usedValueTypes => [type];
+
   BeginOneOutputTry(this.type);
 
   @override
@@ -136,19 +169,22 @@ class BeginOneOutputTry implements Instruction {
   }
 }
 
-class BeginFunctionTry implements Instruction {
+class BeginFunctionTry extends Instruction {
   final FunctionType type;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   BeginFunctionTry(this.type);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x06);
-    s.writeSigned(type.index);
+    s.write(type);
   }
 }
 
-class Catch implements Instruction {
+class Catch extends Instruction {
   final Tag tag;
 
   Catch(this.tag);
@@ -164,7 +200,7 @@ class CatchAll extends SingleByteInstruction {
   const CatchAll() : super(0x19);
 }
 
-class Throw implements Instruction {
+class Throw extends Instruction {
   final Tag tag;
 
   Throw(this.tag);
@@ -176,7 +212,7 @@ class Throw implements Instruction {
   }
 }
 
-class Rethrow implements Instruction {
+class Rethrow extends Instruction {
   final int labelIndex;
 
   Rethrow(this.labelIndex);
@@ -192,7 +228,7 @@ class End extends SingleByteInstruction {
   const End() : super(0x0B);
 }
 
-class Br implements Instruction {
+class Br extends Instruction {
   final int labelIndex;
 
   Br(this.labelIndex);
@@ -204,7 +240,7 @@ class Br implements Instruction {
   }
 }
 
-class BrIf implements Instruction {
+class BrIf extends Instruction {
   final int labelIndex;
 
   BrIf(this.labelIndex);
@@ -216,7 +252,7 @@ class BrIf implements Instruction {
   }
 }
 
-class BrTable implements Instruction {
+class BrTable extends Instruction {
   final List<int> labelIndices;
   final int defaultLabelIndex;
 
@@ -237,7 +273,7 @@ class Return extends SingleByteInstruction {
   const Return() : super(0x0F);
 }
 
-class Call implements Instruction {
+class Call extends Instruction {
   final BaseFunction function;
 
   Call(this.function);
@@ -249,29 +285,35 @@ class Call implements Instruction {
   }
 }
 
-class CallIndirect implements Instruction {
+class CallIndirect extends Instruction {
   final FunctionType type;
   final Table? table;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   CallIndirect(this.type, this.table);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x11);
-    s.writeUnsigned(type.index);
+    s.write(type);
     s.writeUnsigned(table?.index ?? 0);
   }
 }
 
-class CallRef implements Instruction {
+class CallRef extends Instruction {
   final FunctionType type;
 
+  @override
+  List<DefType> get usedDefTypes => [type];
+
   CallRef(this.type);
 
   @override
   void serialize(Serializer s) {
     s.writeByte(0x14);
-    s.writeUnsigned(type.index);
+    s.write(type);
   }
 }
 
@@ -279,9 +321,12 @@ class Drop extends SingleByteInstruction {
   const Drop() : super(0x1A);
 }
 
-class Select implements Instruction {
+class Select extends Instruction {
   final ValueType type;
 
+  @override
+  List<ValueType> get usedValueTypes => [type];
+
   Select(this.type);
 
   @override
@@ -296,7 +341,7 @@ class Select implements Instruction {
   }
 }
 
-class LocalGet implements Instruction {
+class LocalGet extends Instruction {
   final Local local;
 
   LocalGet(this.local);
@@ -308,7 +353,7 @@ class LocalGet implements Instruction {
   }
 }
 
-class LocalSet implements Instruction {
+class LocalSet extends Instruction {
   final Local local;
 
   LocalSet(this.local);
@@ -320,7 +365,7 @@ class LocalSet implements Instruction {
   }
 }
 
-class LocalTee implements Instruction {
+class LocalTee extends Instruction {
   final Local local;
 
   LocalTee(this.local);
@@ -332,7 +377,7 @@ class LocalTee implements Instruction {
   }
 }
 
-class GlobalGet implements Instruction {
+class GlobalGet extends Instruction {
   final Global global;
 
   GlobalGet(this.global);
@@ -344,7 +389,7 @@ class GlobalGet implements Instruction {
   }
 }
 
-class GlobalSet implements Instruction {
+class GlobalSet extends Instruction {
   final Global global;
 
   GlobalSet(this.global);
@@ -356,7 +401,7 @@ class GlobalSet implements Instruction {
   }
 }
 
-class TableSet implements Instruction {
+class TableSet extends Instruction {
   final Table table;
 
   TableSet(this.table);
@@ -368,7 +413,7 @@ class TableSet implements Instruction {
   }
 }
 
-class TableGet implements Instruction {
+class TableGet extends Instruction {
   final Table table;
 
   TableGet(this.table);
@@ -380,7 +425,7 @@ class TableGet implements Instruction {
   }
 }
 
-class TableSize implements Instruction {
+class TableSize extends Instruction {
   final Table table;
 
   TableSize(this.table);
@@ -412,7 +457,7 @@ class MemoryOffsetAlign implements Serializable {
   }
 }
 
-abstract class MemoryInstruction implements Instruction {
+abstract class MemoryInstruction extends Instruction {
   final MemoryOffsetAlign memory;
   final int encoding;
 
@@ -517,7 +562,7 @@ class I64Store32 extends MemoryInstruction {
   I64Store32(super.memory) : super(encoding: 0x3E);
 }
 
-class MemorySize implements Instruction {
+class MemorySize extends Instruction {
   final Memory memory;
 
   MemorySize(this.memory);
@@ -529,7 +574,7 @@ class MemorySize implements Instruction {
   }
 }
 
-class MemoryGrow implements Instruction {
+class MemoryGrow extends Instruction {
   final Memory memory;
 
   MemoryGrow(this.memory);
@@ -541,9 +586,15 @@ class MemoryGrow implements Instruction {
   }
 }
 
-class RefNull implements Instruction {
+class RefNull extends Instruction {
   final HeapType heapType;
 
+  @override
+  List<DefType> get usedDefTypes {
+    final type = heapType;
+    return type is DefType ? [type] : const [];
+  }
+
   RefNull(this.heapType);
 
   @override
@@ -557,7 +608,7 @@ class RefIsNull extends SingleByteInstruction {
   const RefIsNull() : super(0xD1);
 }
 
-class RefFunc implements Instruction {
+class RefFunc extends Instruction {
   final BaseFunction function;
 
   RefFunc(this.function);
@@ -573,7 +624,7 @@ class RefAsNonNull extends SingleByteInstruction {
   const RefAsNonNull() : super(0xD4);
 }
 
-class BrOnNull implements Instruction {
+class BrOnNull extends Instruction {
   final int labelIndex;
 
   BrOnNull(this.labelIndex);
@@ -589,7 +640,7 @@ class RefEq extends SingleByteInstruction {
   const RefEq() : super(0xD3);
 }
 
-class BrOnNonNull implements Instruction {
+class BrOnNonNull extends Instruction {
   final int labelIndex;
 
   BrOnNonNull(this.labelIndex);
@@ -601,131 +652,161 @@ class BrOnNonNull implements Instruction {
   }
 }
 
-class StructGet implements Instruction {
+class StructGet extends Instruction {
   final StructType structType;
   final int fieldIndex;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructGet(this.structType, this.fieldIndex);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x02]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
     s.writeUnsigned(fieldIndex);
   }
 }
 
-class StructGetS implements Instruction {
+class StructGetS extends Instruction {
   final StructType structType;
   final int fieldIndex;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructGetS(this.structType, this.fieldIndex);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x03]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
     s.writeUnsigned(fieldIndex);
   }
 }
 
-class StructGetU implements Instruction {
+class StructGetU extends Instruction {
   final StructType structType;
   final int fieldIndex;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructGetU(this.structType, this.fieldIndex);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x04]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
     s.writeUnsigned(fieldIndex);
   }
 }
 
-class StructSet implements Instruction {
+class StructSet extends Instruction {
   final StructType structType;
   final int fieldIndex;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructSet(this.structType, this.fieldIndex);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x05]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
     s.writeUnsigned(fieldIndex);
   }
 }
 
-class StructNew implements Instruction {
+class StructNew extends Instruction {
   final StructType structType;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructNew(this.structType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x00]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
   }
 }
 
-class StructNewDefault implements Instruction {
+class StructNewDefault extends Instruction {
   final StructType structType;
 
+  @override
+  List<DefType> get usedDefTypes => [structType];
+
   StructNewDefault(this.structType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x01]);
-    s.writeUnsigned(structType.index);
+    s.write(structType);
   }
 }
 
-class ArrayGet implements Instruction {
+class ArrayGet extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayGet(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x0b]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
-class ArrayGetS implements Instruction {
+class ArrayGetS extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayGetS(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x0c]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
-class ArrayGetU implements Instruction {
+class ArrayGetU extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayGetU(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x0d]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
-class ArraySet implements Instruction {
+class ArraySet extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArraySet(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x0E]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
@@ -733,81 +814,99 @@ class ArrayLen extends MultiByteInstruction {
   const ArrayLen() : super(const [0xFB, 0x0F]);
 }
 
-class ArrayNewFixed implements Instruction {
+class ArrayNewFixed extends Instruction {
   final ArrayType arrayType;
   final int length;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayNewFixed(this.arrayType, this.length);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x08]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
     s.writeUnsigned(length);
   }
 }
 
-class ArrayNew implements Instruction {
+class ArrayNew extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayNew(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x06]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
-class ArrayNewDefault implements Instruction {
+class ArrayNewDefault extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayNewDefault(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x07]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
-class ArrayNewData implements Instruction {
+class ArrayNewData extends Instruction {
   final ArrayType arrayType;
   final BaseDataSegment data;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayNewData(this.arrayType, this.data);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x09]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
     s.writeUnsigned(data.index);
   }
 }
 
-class ArrayCopy implements Instruction {
+class ArrayCopy extends Instruction {
   final ArrayType destArrayType;
   final ArrayType sourceArrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [destArrayType, sourceArrayType];
+
   ArrayCopy({required this.destArrayType, required this.sourceArrayType});
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x11]);
-    s.writeUnsigned(destArrayType.index);
-    s.writeUnsigned(sourceArrayType.index);
+    s.write(destArrayType);
+    s.write(sourceArrayType);
   }
 }
 
-class ArrayFill implements Instruction {
+class ArrayFill extends Instruction {
   final ArrayType arrayType;
 
+  @override
+  List<DefType> get usedDefTypes => [arrayType];
+
   ArrayFill(this.arrayType);
 
   @override
   void serialize(Serializer s) {
     s.writeBytes(const [0xFB, 0x10]);
-    s.writeUnsigned(arrayType.index);
+    s.write(arrayType);
   }
 }
 
@@ -823,7 +922,7 @@ class I31GetU extends MultiByteInstruction {
   const I31GetU() : super(const [0xFB, 0x1E]);
 }
 
-class RefTest implements Instruction {
+class RefTest extends Instruction {
   final RefType targetType;
 
   RefTest(this.targetType);
@@ -835,9 +934,12 @@ class RefTest implements Instruction {
   }
 }
 
-class RefCast implements Instruction {
+class RefCast extends Instruction {
   final RefType targetType;
 
+  @override
+  List<ValueType> get usedValueTypes => [targetType];
+
   RefCast(this.targetType);
 
   @override
@@ -847,11 +949,14 @@ class RefCast implements Instruction {
   }
 }
 
-class BrOnCast implements Instruction {
+class BrOnCast extends Instruction {
   final int labelIndex;
   final RefType inputType;
   final RefType targetType;
 
+  @override
+  List<ValueType> get usedValueTypes => [inputType, targetType];
+
   BrOnCast(this.labelIndex, this.inputType, this.targetType);
 
   @override
@@ -866,11 +971,14 @@ class BrOnCast implements Instruction {
   }
 }
 
-class BrOnCastFail implements Instruction {
+class BrOnCastFail extends Instruction {
   final int labelIndex;
   final RefType inputType;
   final RefType targetType;
 
+  @override
+  List<ValueType> get usedValueTypes => [inputType, targetType];
+
   BrOnCastFail(this.labelIndex, this.inputType, this.targetType);
 
   @override
@@ -893,7 +1001,7 @@ class ExternExternalize extends MultiByteInstruction {
   const ExternExternalize() : super(const [0xFB, 0x1B]);
 }
 
-class I32Const implements Instruction {
+class I32Const extends Instruction {
   final int value;
 
   I32Const(this.value);
@@ -905,7 +1013,7 @@ class I32Const implements Instruction {
   }
 }
 
-class I64Const implements Instruction {
+class I64Const extends Instruction {
   final int value;
 
   I64Const(this.value);
@@ -917,7 +1025,7 @@ class I64Const implements Instruction {
   }
 }
 
-class F32Const implements Instruction {
+class F32Const extends Instruction {
   final double value;
 
   F32Const(this.value);
@@ -929,7 +1037,7 @@ class F32Const implements Instruction {
   }
 }
 
-class F64Const implements Instruction {
+class F64Const extends Instruction {
   final double value;
 
   F64Const(this.value);

pkg/wasm_builder/lib/src/ir/module.dart

@@ -62,7 +62,7 @@ class Module implements Serializable {
         globals.namedCount > 0) {
       NameSection(
               <BaseFunction>[...functions.imported, ...functions.defined],
-              types.defined,
+              [for (final t in types.recursionGroups) ...t],
               <Global>[...globals.imported, ...globals.defined],
               watchPoints,
               functionNameCount: functions.namedCount,

pkg/wasm_builder/lib/src/ir/type.dart

@@ -45,6 +45,12 @@ abstract class ValueType implements StorageType {
 
   /// Whether this type is defaultable. Primitive types are always defaultable.
   bool get defaultable => true;
+
+  /// The heap [DefType] referenced by this type if any.
+  ///
+  /// Used by the type builder to determine the set of [DefType]s referenced in
+  /// a module.
+  DefType? get containedDefType => null;
 }
 
 enum NumTypeKind { i32, i64, f32, f64, v128 }
@@ -203,6 +209,12 @@ class RefType extends ValueType {
     return heapType.isSubtypeOf(other.heapType);
   }
 
+  @override
+  DefType? get containedDefType {
+    final type = heapType;
+    return type is DefType ? type : null;
+  }
+
   @override
   void serialize(Serializer s) {
     if (nullable != heapType.nullableByDefault) {

pkg/wasm_builder/lib/src/ir/types.dart

@@ -8,13 +8,10 @@ part 'type.dart';
 
 class Types {
   /// Types defined in this module.
-  final List<DefType> defined;
-
-  /// Recursion group splits.
-  final List<int> recursionGroupSplits;
+  final List<List<DefType>> recursionGroups;
 
   /// Name count.
   final int namedCount;
 
-  Types(this.defined, this.recursionGroupSplits, this.namedCount);
+  Types(this.recursionGroups, this.namedCount);
 }

pkg/wasm_builder/lib/src/serialize/sections.dart

@@ -35,34 +35,43 @@ class TypeSection extends Section {
 
   TypeSection(this.types, super.watchPoints);
 
-  List<ir.DefType> get defTypes => types.defined;
+  List<List<ir.DefType>> get recursionGroups => types.recursionGroups;
 
   @override
   int get id => 1;
 
   @override
-  bool get isNotEmpty => defTypes.isNotEmpty;
+  bool get isNotEmpty => recursionGroups.isNotEmpty;
 
   @override
   void serializeContents(Serializer s) {
-    s.writeUnsigned(types.recursionGroupSplits.length + 1);
+    s.writeUnsigned(types.recursionGroups.length);
     int typeIndex = 0;
-    for (int split
-        in types.recursionGroupSplits.followedBy([defTypes.length])) {
+
+    // Set all the indices first since types can be referenced before they are
+    // serialized.
+    for (final group in recursionGroups) {
+      assert(group.isNotEmpty, 'Empty groups are not allowed.');
+
+      for (final type in group) {
+        type.index = typeIndex++;
+      }
+    }
+    for (final group in recursionGroups) {
       s.writeByte(0x4E); // -0x32
-      s.writeUnsigned(split - typeIndex);
-      for (; typeIndex < split; typeIndex++) {
-        ir.DefType defType = defTypes[typeIndex];
-        assert(defType.superType == null || defType.superType!.index < split,
-            "Type '$defType' has a supertype in a later recursion group");
+      s.writeUnsigned(group.length);
+      for (final type in group) {
         assert(
-            defType.constituentTypes
+            type.superType == null || type.superType!.index <= group.last.index,
+            "Type '$type' has a supertype in a later recursion group");
+        assert(
+            type.constituentTypes
                 .whereType<ir.RefType>()
                 .map((t) => t.heapType)
                 .whereType<ir.DefType>()
-                .every((d) => d.index < split),
-            "Type '$defType' depends on a type in a later recursion group");
-        defType.serializeDefinition(s);
+                .every((d) => d.index <= group.last.index),
+            "Type '$type' depends on a type in a later recursion group");
+        type.serializeDefinition(s);
       }
     }
   }

pkg/wasm_builder/pubspec.yaml

@@ -6,6 +6,10 @@ description: Generate binary Wasm modules
 environment:
   sdk: '>=2.19.0 <3.0.0'
 
+# Use 'any' constraints here; we get our versions from the DEPS file.
+dependencies:
+  collection: any
+
 # Use 'any' constraints here; we get our versions from the DEPS file.
 dev_dependency:
   lints: any