[dart2wasm] Implement instantiation closure equality

Changes:

- Instantiation context structs now have a common supertype:
  `#InstantiationClosureContextBase`. This supertype is used to check if
  a closure is an instantiation.

- Generic closure vtables now have one more entry for comparing types in
  instantiation contexts.

- All instantiations with same number of types use the same context
  struct type and the vtable entry for comparing types in contexts.

- `_Closure._equals` checks if closures are instantiations (with a type
  test against `#InstantiationClosureContextBase`), and if they are,
  gets the contexts and calls the function in the vtable for comparing
  the types in the contexts.

Fixes #51030

Change-Id: I680041480963381337dcfa1fbf6c3a77ca564205
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/285902
Commit-Queue: Ömer Ağacan <omersa@google.com>
Reviewed-by: Aske Simon Christensen <askesc@google.com>

pkg/dart2wasm/lib/class_info.dart

@@ -29,7 +29,8 @@ class FieldIndex {
   static const closureVtable = 3;
   static const closureRuntimeType = 4;
   static const vtableDynamicCallEntry = 0;
-  static const vtableInstantiationFunction = 1;
+  static const vtableInstantiationTypeComparisonFunction = 1;
+  static const vtableInstantiationFunction = 2;
   static const instantiationContextInner = 0;
   static const instantiationContextTypeArgumentsBase = 1;
   static const typeIsDeclaredNullable = 2;

pkg/dart2wasm/lib/closures.dart

@@ -71,6 +71,14 @@ class ClosureRepresentation {
       _instantiationFunctionThunk!();
   w.DefinedFunction Function()? _instantiationFunctionThunk;
 
+  /// The function that takes instantiation context of this generic closure and
+  /// another instantiation context (both as `ref
+  /// #InstantiationClosureContextBase`) and compares types in the contexts.
+  /// This function is used to implement function equality of instantiations.
+  late final w.DefinedFunction instantiationTypeComparisonFunction =
+      _instantiationTypeComparisonFunctionThunk!();
+  w.DefinedFunction Function()? _instantiationTypeComparisonFunctionThunk;
+
   /// The signature of the function that instantiates this generic closure.
   w.FunctionType get instantiationFunctionType {
     assert(isGeneric);
@@ -155,8 +163,17 @@ class ClosureLayouter extends RecursiveVisitor {
   // The member currently being visited while collecting function signatures.
   Member? currentMember;
 
+  // For non-generic closures. The entries are:
+  // 0: Dynamic call entry
+  // 1-...: Entries for calling the closure
   static const int vtableBaseIndexNonGeneric = 1;
-  static const int vtableBaseIndexGeneric = 2;
+
+  // For generic closures. The entries are:
+  // 0: Dynamic call entry
+  // 1: Instantiation type comparison function
+  // 2: Instantiation function
+  // 3-...: Entries for calling the closure
+  static const int vtableBaseIndexGeneric = 3;
 
   // Base struct for vtables without the dynamic call entry added. Referenced
   // by [closureBaseStruct] instead of the fully initialized version
@@ -164,13 +181,42 @@ class ClosureLayouter extends RecursiveVisitor {
   late final w.StructType _vtableBaseStructBare =
       m.addStructType("#VtableBase");
 
-  // Base struct for vtables.
+  /// Base struct for instantiation closure contexts. Type tests against this
+  /// type is used in `_Closure._equals` to check if a closure is an
+  /// instantiation.
+  late final w.StructType instantiationContextBaseStruct =
+      m.addStructType("#InstantiationClosureContextBase", fields: [
+    w.FieldType(w.RefType.def(closureBaseStruct, nullable: false),
+        mutable: false),
+  ]);
+
+  /// Base struct for non-generic closure vtables.
   late final w.StructType vtableBaseStruct = _vtableBaseStructBare
     ..fields.add(w.FieldType(
         w.RefType.def(translator.dynamicCallVtableEntryFunctionType,
             nullable: false),
         mutable: false));
 
+  /// Base struct for generic closure vtables.
+  late final w.StructType genericVtableBaseStruct = m.addStructType(
+      "#GenericVtableBase",
+      fields: vtableBaseStruct.fields.toList()
+        ..add(w.FieldType(
+            w.RefType.def(instantiationClosureTypeComparisonFunctionType,
+                nullable: false),
+            mutable: false)),
+      superType: vtableBaseStruct);
+
+  /// Type of [ClosureRepresentation.instantiationTypeComparisonFunction].
+  late final w.FunctionType instantiationClosureTypeComparisonFunctionType =
+      m.addFunctionType(
+    [
+      w.RefType.def(instantiationContextBaseStruct, nullable: false),
+      w.RefType.def(instantiationContextBaseStruct, nullable: false)
+    ],
+    [w.NumType.i32], // bool
+  );
+
   // Base struct for closures.
   late final w.StructType closureBaseStruct = _makeClosureStruct(
       "#ClosureBase", _vtableBaseStructBare, translator.closureInfo.struct);
@@ -180,6 +226,25 @@ class ClosureLayouter extends RecursiveVisitor {
   late final w.RefType functionTypeType =
       translator.classInfo[translator.functionTypeClass]!.nonNullableType;
 
+  final Map<int, w.StructType> _instantiationContextBaseStructs = {};
+
+  w.StructType _getInstantiationContextBaseStruct(int numTypes) =>
+      _instantiationContextBaseStructs.putIfAbsent(
+          numTypes,
+          () => m.addStructType("#InstantiationClosureContextBase-$numTypes",
+              fields: [
+                w.FieldType(w.RefType.def(closureBaseStruct, nullable: false),
+                    mutable: false),
+                ...List.filled(numTypes, w.FieldType(typeType, mutable: false))
+              ],
+              superType: instantiationContextBaseStruct));
+
+  final Map<int, w.DefinedFunction> _instantiationTypeComparisonFunctions = {};
+
+  w.DefinedFunction _getInstantiationTypeComparisonFunction(int numTypes) =>
+      _instantiationTypeComparisonFunctions.putIfAbsent(
+          numTypes, () => _createInstantiationTypeComparisonFunction(numTypes));
+
   w.StructType _makeClosureStruct(
       String name, w.StructType vtableStruct, w.StructType superType) {
     // A closure contains:
@@ -281,7 +346,8 @@ class ClosureLayouter extends RecursiveVisitor {
     List<String> nameTags = ["$typeCount", "$positionalCount", ...names];
     String vtableName = ["#Vtable", ...nameTags].join("-");
     String closureName = ["#Closure", ...nameTags].join("-");
-    w.StructType parentVtableStruct = parent?.vtableStruct ?? vtableBaseStruct;
+    w.StructType parentVtableStruct = parent?.vtableStruct ??
+        (typeCount == 0 ? vtableBaseStruct : genericVtableBaseStruct);
     w.StructType vtableStruct = m.addStructType(vtableName,
         fields: parentVtableStruct.fields, superType: parentVtableStruct);
     w.StructType closureStruct = _makeClosureStruct(
@@ -322,7 +388,7 @@ class ClosureLayouter extends RecursiveVisitor {
                 mutable: false),
             ...List.filled(typeCount, w.FieldType(typeType, mutable: false))
           ],
-          superType: parent?.instantiationContextStruct);
+          superType: _getInstantiationContextBaseStruct(typeCount));
     }
 
     // Add vtable fields for additional entry points relative to the parent.
@@ -412,6 +478,9 @@ class ClosureLayouter extends RecursiveVisitor {
             closureStruct,
             instantiationFunctionName);
       };
+
+      representation._instantiationTypeComparisonFunctionThunk =
+          () => _getInstantiationTypeComparisonFunction(typeCount);
     }
 
     return representation;
@@ -622,6 +691,67 @@ class ClosureLayouter extends RecursiveVisitor {
     return instantiationFunction;
   }
 
+  w.DefinedFunction _createInstantiationTypeComparisonFunction(int numTypes) {
+    final function = m.addFunction(
+        instantiationClosureTypeComparisonFunctionType,
+        "#InstantiationTypeComparison-$numTypes");
+
+    final w.Instructions b = function.body;
+
+    final contextStructType = _getInstantiationContextBaseStruct(numTypes);
+    final contextRefType = w.RefType.def(contextStructType, nullable: false);
+
+    final thisContext = function.locals[0];
+    final otherContext = function.locals[1];
+
+    final thisContextLocal = function.addLocal(contextRefType);
+    final otherContextLocal = function.addLocal(contextRefType);
+
+    // Call site (`_Closure._equals`) checks that closures are instantiations
+    // of the same function, so we can assume they have the right instantiation
+    // context types.
+    b.local_get(otherContext);
+    b.ref_cast(contextRefType);
+    b.local_set(otherContextLocal);
+
+    b.local_get(thisContext);
+    b.ref_cast(contextRefType);
+    b.local_set(thisContextLocal);
+
+    for (int i = 0; i < numTypes; i += 1) {
+      final typeFieldIdx = FieldIndex.instantiationContextTypeArgumentsBase + i;
+      b.local_get(thisContextLocal);
+      b.struct_get(contextStructType, typeFieldIdx);
+      b.local_get(otherContextLocal);
+      b.struct_get(contextStructType, typeFieldIdx);
+
+      // Virtual call to `Object.==`
+      final selector = translator.dispatchTable
+          .selectorForTarget(translator.coreTypes.objectEquals.reference);
+      final selectorOffset = selector.offset!;
+      b.local_get(thisContextLocal);
+      b.struct_get(contextStructType, typeFieldIdx);
+      b.struct_get(translator.topInfo.struct, FieldIndex.classId);
+      if (selectorOffset != 0) {
+        b.i32_const(selectorOffset);
+        b.i32_add();
+      }
+      b.call_indirect(selector.signature, translator.dispatchTable.wasmTable);
+      b.if_();
+    }
+
+    b.i32_const(1); // true
+    b.return_();
+
+    for (int i = 0; i < numTypes; i += 1) {
+      b.end();
+    }
+
+    b.i32_const(0); // false
+    b.end(); // end of function
+    return function;
+  }
+
   ClosureRepresentationsForParameterCount _representationsForCounts(
       int typeCount, int positionalCount) {
     while (representations.length <= typeCount) {

pkg/dart2wasm/lib/intrinsics.dart

@@ -1543,6 +1543,15 @@ class Intrinsifier {
       //   bool _equals(f1, f2) {
       //     if (identical(f1, f2) return true;
       //     if (f1.vtable == f2.vtable) {
+      //       if (f1 and f2 are instantiations) {
+      //         if (f1.context.inner.vtable == f2.context.inner.vtable) {
+      //           if (identical(f1.context.inner.context,
+      //                         f2.context.inner.context)) {
+      //             return typesEqual(f1.context, f2.context);
+      //           }
+      //         }
+      //         return false;
+      //       }
       //       if (v1.context is #Top && v2.context is #Top) {
       //         return identical(v1.context, v2.context);
       //       }
@@ -1586,6 +1595,76 @@ class Intrinsifier {
 
       b.if_(); // fun1.vtable == fun2.vtable
 
+      // Check if closures are instantiations. Since they have the same vtable
+      // it's enough to check just one of them.
+      final instantiationContextBase = w.RefType(
+          translator.closureLayouter.instantiationContextBaseStruct,
+          nullable: false);
+      final instantiationBlock =
+          b.block([], [w.RefType.struct(nullable: false)]);
+
+      b.local_get(fun1);
+      b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+      b.br_on_cast_fail(instantiationContextBase, instantiationBlock);
+
+      // Closures are instantiations. Compare inner function vtables to check
+      // that instantiations are for the same generic function.
+      final vtable1Local = codeGen.function.addLocal(w.RefType.def(
+          translator.closureLayouter.vtableBaseStruct,
+          nullable: false));
+
+      void getInstantiationContextInner(w.Local fun) {
+        b.local_get(fun);
+        // instantiation.context
+        b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+        b.ref_cast(instantiationContextBase);
+        // instantiation.context.inner
+        b.struct_get(translator.closureLayouter.instantiationContextBaseStruct,
+            FieldIndex.instantiationContextInner);
+      }
+
+      getInstantiationContextInner(fun1);
+      b.struct_get(closureBaseStruct, FieldIndex.closureVtable);
+      b.local_tee(vtable1Local);
+
+      getInstantiationContextInner(fun2);
+      b.struct_get(closureBaseStruct, FieldIndex.closureVtable);
+
+      b.ref_eq();
+      b.if_(); // fun1.context.inner.vtable == fun2.context.inner.vtable
+
+      // Closures are instantiations of the same function, compare inner
+      // contexts and types
+      getInstantiationContextInner(fun1);
+      b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+      getInstantiationContextInner(fun2);
+      b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+      b.ref_eq();
+      b.if_(); // fun1.context.inner.context == fun2.context.inner.context
+      // Inner contexts are equal, compare types
+      b.local_get(fun1);
+      b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+      b.ref_cast(instantiationContextBase);
+      b.local_get(fun2);
+      b.struct_get(closureBaseStruct, FieldIndex.closureContext);
+      b.ref_cast(instantiationContextBase);
+      b.local_get(vtable1Local);
+      b.ref_cast(w.RefType.def(
+          translator.closureLayouter.genericVtableBaseStruct,
+          nullable: false));
+      b.struct_get(translator.closureLayouter.genericVtableBaseStruct,
+          FieldIndex.vtableInstantiationTypeComparisonFunction);
+      b.call_ref(translator
+          .closureLayouter.instantiationClosureTypeComparisonFunctionType);
+      b.return_();
+      b.end(); // fun1.context.inner.context == fun2.context.inner.context
+      b.end(); // fun1.context.inner.vtable == fun2.context.inner.vtable
+      // Contexts or inner vtables are not equal
+      b.i32_const(0); // false
+      b.return_();
+      b.end(); // instantiationBlock
+      b.drop();
+
       // Compare context references. If context of a function has the top type
       // then the function is an instance tear-off. Otherwise it's a closure.
       final contextCheckFail = b.block([], [w.RefType.struct(nullable: false)]);

pkg/dart2wasm/lib/translator.dart

@@ -748,6 +748,7 @@ class Translator with KernelNodes {
     final dynamicCallEntry = makeDynamicCallEntry();
     ib.ref_func(dynamicCallEntry);
     if (representation.isGeneric) {
+      ib.ref_func(representation.instantiationTypeComparisonFunction);
       ib.ref_func(representation.instantiationFunction);
     }
     for (int posArgCount = 0; posArgCount <= positionalCount; posArgCount++) {