Handle exactness in MinimizeRecGroups (#7555)

Treat rec groups differing only in the exactness of a reference as
different, but only when custom descriptors is enabled. When custom
descriptors is not enabled, exactness will be erased before the binary
is written, so if two minimized rec groups differed only in exactness,
they would in fact be written as the same rec group. This would change
the behavior of casts meant to differentiate between types in that rec
group, so it would be incorrect.

Author: tlively

scripts/test/fuzzing.py

@@ -127,6 +127,8 @@
     'type-merging-exact.wast',
     'type-refining-exact.wast',
     'type-refining-gufa-exact.wast',
+    'mimimize-rec-groups-exact.wast',
+    'mimimize-rec-groups-ignore-exact.wast',
     'public-exact.wast',
     # TODO: fuzzer support for custom descriptors
     'custom-descriptors.wast',

src/passes/MinimizeRecGroups.cpp

@@ -208,14 +208,14 @@ struct GroupClassInfo {
   static std::vector<std::vector<Index>> initSubtypeGraph(RecGroupInfo& info);
   GroupClassInfo(RecGroupInfo& info);
 
-  void advance() {
+  void advance(FeatureSet features) {
     ++orders;
     if (orders == orders.end()) {
-      advanceBrand();
+      advanceBrand(features);
     }
   }
 
-  void advanceBrand() {
+  void advanceBrand(FeatureSet features) {
     if (brand) {
       ++*brand;
     } else {
@@ -231,8 +231,8 @@ struct GroupClassInfo {
       }
     }
     // Make sure the brand is not the same as the real type.
-    if (singletonType &&
-        RecGroupShape({**brand}) == RecGroupShape({*singletonType})) {
+    if (singletonType && RecGroupShape({**brand}, features) ==
+                           RecGroupShape({*singletonType}, features)) {
       ++*brand;
     }
     // Start back at the initial permutation with the new brand.
@@ -370,9 +370,13 @@ struct MinimizeRecGroups : Pass {
   // whose shapes we need to check for uniqueness to avoid deep recursions.
   std::vector<Index> shapesToUpdate;
 
+  // The comparison of rec group shapes depends on the features.
+  FeatureSet features;
+
   void run(Module* module) override {
+    features = module->features;
     // There are no recursion groups to minimize if GC is not enabled.
-    if (!module->features.hasGC()) {
+    if (!features.hasGC()) {
       return;
     }
 
@@ -402,7 +406,7 @@ struct MinimizeRecGroups : Pass {
     for (auto group : publicGroups) {
       publicGroupTypes.emplace_back(group.begin(), group.end());
       [[maybe_unused]] auto [_, inserted] = groupShapeIndices.insert(
-        {RecGroupShape(publicGroupTypes.back()), PublicGroupIndex});
+        {RecGroupShape(publicGroupTypes.back(), features), PublicGroupIndex});
       assert(inserted);
     }
 
@@ -452,8 +456,8 @@ struct MinimizeRecGroups : Pass {
   }
 
   void updateShape(Index group) {
-    auto [it, inserted] =
-      groupShapeIndices.insert({RecGroupShape(groups[group].group), group});
+    auto [it, inserted] = groupShapeIndices.insert(
+      {RecGroupShape(groups[group].group, features), group});
     if (inserted) {
       // This shape was unique. We're done.
       return;
@@ -509,7 +513,7 @@ struct MinimizeRecGroups : Pass {
         // We have everything we need to generate the next permutation of this
         // group.
         auto& classInfo = *groups[groupRep].classInfo;
-        classInfo.advance();
+        classInfo.advance(features);
         classInfo.permute(groupInfo);
         shapesToUpdate.push_back(group);
         return;
@@ -538,7 +542,7 @@ struct MinimizeRecGroups : Pass {
 
       // Move to the next permutation after advancing the type brand to skip
       // further repeated shapes.
-      classInfo.advanceBrand();
+      classInfo.advanceBrand(features);
       classInfo.permute(groupInfo);
 
       shapesToUpdate.push_back(group);
@@ -556,7 +560,7 @@ struct MinimizeRecGroups : Pass {
     // conflict.
     if (groups[groupRep].classInfo && groups[otherRep].classInfo) {
       auto& classInfo = *groups[groupRep].classInfo;
-      classInfo.advance();
+      classInfo.advance(features);
       classInfo.permute(groupInfo);
       shapesToUpdate.push_back(group);
       return;
@@ -578,7 +582,7 @@ struct MinimizeRecGroups : Pass {
       // same shape. Advance `group` to the next permutation.
       otherInfo.classInfo = std::nullopt;
       otherInfo.permutation = groupInfo.permutation;
-      classInfo.advance();
+      classInfo.advance(features);
       classInfo.permute(groupInfo);
 
       shapesToUpdate.push_back(group);
@@ -600,7 +604,7 @@ struct MinimizeRecGroups : Pass {
       // permutation.
       groupInfo.classInfo = std::nullopt;
       groupInfo.permutation = otherInfo.permutation;
-      classInfo.advance();
+      classInfo.advance(features);
       classInfo.permute(groupInfo);
 
       shapesToUpdate.push_back(group);
@@ -754,9 +758,10 @@ struct MinimizeRecGroups : Pass {
     // shapes to lists of automorphically equivalent root types.
     std::map<ComparableRecGroupShape, std::vector<HeapType>> typeClasses;
     for (const auto& order : dfsOrders) {
-      ComparableRecGroupShape shape(order, [this](HeapType a, HeapType b) {
-        return this->typeIndices.at(a) < this->typeIndices.at(b);
-      });
+      ComparableRecGroupShape shape(
+        order, features, [this](HeapType a, HeapType b) {
+          return this->typeIndices.at(a) < this->typeIndices.at(b);
+        });
       typeClasses[shape].push_back(order[0]);
     }
 

src/tools/wasm-fuzz-types.cpp

@@ -542,7 +542,7 @@ void Fuzzer::checkRecGroupShapes() {
   };
 
   for (size_t i = 0; i < groups.size(); ++i) {
-    ComparableRecGroupShape shape(groups[i], less);
+    ComparableRecGroupShape shape(groups[i], FeatureSet::All, less);
     // A rec group should compare equal to itself.
     if (shape != shape) {
       Fatal() << "Rec group shape " << i << " not equal to itself";
@@ -556,7 +556,7 @@ void Fuzzer::checkRecGroupShapes() {
 
     // Check how it compares to other groups.
     for (size_t j = i + 1; j < groups.size(); ++j) {
-      ComparableRecGroupShape other(groups[j], less);
+      ComparableRecGroupShape other(groups[j], FeatureSet::All, less);
       bool isLess = shape < other;
       bool isEq = shape == other;
       bool isGreater = shape > other;
@@ -598,7 +598,7 @@ void Fuzzer::checkRecGroupShapes() {
 
       if (j + 1 < groups.size()) {
         // Check transitivity.
-        RecGroupShape third(groups[j + 1]);
+        RecGroupShape third(groups[j + 1], FeatureSet::All);
         if ((isLess && other <= third && shape >= third) ||
             (isEq && other == third && shape != third) ||
             (isGreater && other >= third && shape <= third)) {

src/wasm-type-shape.h

@@ -20,6 +20,7 @@
 #include <functional>
 #include <vector>
 
+#include "wasm-features.h"
 #include "wasm-type.h"
 
 namespace wasm {
@@ -35,7 +36,13 @@ namespace wasm {
 struct RecGroupShape {
   const std::vector<HeapType>& types;
 
-  RecGroupShape(const std::vector<HeapType>& types) : types(types) {}
+  // Depending on the feature set, some types may be generalized when they are
+  // written out. Take the features into account to ensure our comparisons
+  // account for the rec groups that will ultimately be written.
+  const FeatureSet features;
+
+  RecGroupShape(const std::vector<HeapType>& types, const FeatureSet features)
+    : types(types), features(features) {}
 
   bool operator==(const RecGroupShape& other) const;
   bool operator!=(const RecGroupShape& other) const {
@@ -51,8 +58,9 @@ struct ComparableRecGroupShape : RecGroupShape {
   std::function<bool(HeapType, HeapType)> less;
 
   ComparableRecGroupShape(const std::vector<HeapType>& types,
+                          FeatureSet features,
                           std::function<bool(HeapType, HeapType)> less)
-    : RecGroupShape(types), less(less) {}
+    : RecGroupShape(types, features), less(less) {}
 
   bool operator<(const RecGroupShape& other) const;
   bool operator>(const RecGroupShape& other) const;

src/wasm/wasm-type-shape.cpp

@@ -25,13 +25,16 @@ namespace {
 enum Comparison { EQ, LT, GT };
 
 template<typename CompareTypes> struct RecGroupComparator {
+  FeatureSet features;
   std::unordered_map<HeapType, Index> indicesA;
   std::unordered_map<HeapType, Index> indicesB;
   CompareTypes compareTypes;
 
   RecGroupComparator(CompareTypes compareTypes) : compareTypes(compareTypes) {}
 
   Comparison compare(const RecGroupShape& a, const RecGroupShape& b) {
+    assert(a.features == b.features);
+    features = a.features;
     if (a.types.size() != b.types.size()) {
       return a.types.size() < b.types.size() ? LT : GT;
     }
@@ -147,6 +150,11 @@ template<typename CompareTypes> struct RecGroupComparator {
       return compare(a.getTuple(), b.getTuple());
     }
     assert(a.isRef() && b.isRef());
+    // Only consider exactness if custom descriptors are enabled. Otherwise, it
+    // will be erased when the types are written, so we ignore it here, too.
+    if (features.hasCustomDescriptors() && a.isExact() != b.isExact()) {
+      return a.isExact() < b.isExact() ? LT : GT;
+    }
     if (a.isNullable() != b.isNullable()) {
       return a.isNullable() < b.isNullable() ? LT : GT;
     }
@@ -201,9 +209,11 @@ template<typename CompareTypes>
 RecGroupComparator(CompareTypes) -> RecGroupComparator<CompareTypes>;
 
 struct RecGroupHasher {
+  FeatureSet features;
   std::unordered_map<HeapType, Index> typeIndices;
 
   size_t hash(const RecGroupShape& shape) {
+    features = shape.features;
     for (auto type : shape.types) {
       typeIndices.insert({type, typeIndices.size()});
     }
@@ -285,6 +295,9 @@ struct RecGroupHasher {
       return digest;
     }
     assert(type.isRef());
+    if (features.hasCustomDescriptors()) {
+      wasm::rehash(digest, type.isExact());
+    }
     wasm::rehash(digest, type.isNullable());
     hash_combine(digest, hash(type.getHeapType()));
     return digest;

test/lit/passes/minimize-rec-groups-exact.wast

@@ -0,0 +1,19 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-opt %s -all --minimize-rec-groups -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $foo (struct))
+  (type $foo (struct))
+  ;; CHECK:      (type $exact (struct (field (ref (exact $foo)))))
+  (type $exact (struct (field (ref (exact $foo)))))
+  ;; CHECK:      (type $inexact (struct (field (ref $foo))))
+  (type $inexact (struct (field (ref $foo))))
+
+  ;; If we didn't differentiate between exact and inexact types, there would be
+  ;; an assertion failure on adding these public types to the set of public
+  ;; shapes.
+  ;; CHECK:      (import "" "exact" (global $exact (ref null $exact)))
+  (import "" "exact" (global $exact (ref null $exact)))
+  ;; CHECK:      (import "" "inexact" (global $inexact (ref null $inexact)))
+  (import "" "inexact" (global $inexact (ref null $inexact)))
+)

test/lit/passes/minimize-rec-groups-ignore-exact.wast

@@ -0,0 +1,53 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; Check that we take exactness into account correctly depending on the
+;; features. It differentiates shapes only when custom descriptors is enabled.
+
+;; RUN: wasm-opt %s -all                              --minimize-rec-groups -S -o - | filecheck %s
+;; RUN: wasm-opt %s -all --disable-custom-descriptors --minimize-rec-groups -S -o - | filecheck %s --check-prefix=NO_CD
+
+(module
+  (rec
+    (type $foo (struct))
+
+    ;; This SCC has only one distinct permutation.
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $b-inexact (struct (field (ref $a-inexact))))
+
+    ;; CHECK:       (type $a-inexact (struct (field (ref $b-inexact))))
+    ;; NO_CD:      (rec
+    ;; NO_CD-NEXT:  (type $b-inexact (struct (field (ref $a-inexact))))
+
+    ;; NO_CD:       (type $a-inexact (struct (field (ref $b-inexact))))
+    (type $a-inexact (struct (field (ref $b-inexact))))
+    (type $b-inexact (struct (field (ref $a-inexact))))
+
+    ;; This SCC is only different because of exactness. It needs a brand only if
+    ;; custom descriptors is disabled.
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $b-exact (struct (field (ref (exact $a-exact)))))
+
+    ;; CHECK:       (type $a-exact (struct (field (ref (exact $b-exact)))))
+    ;; NO_CD:      (rec
+    ;; NO_CD-NEXT:  (type $2 (struct))
+
+    ;; NO_CD:       (type $b-exact (struct (field (ref (exact $a-exact)))))
+
+    ;; NO_CD:       (type $a-exact (struct (field (ref (exact $b-exact)))))
+    (type $a-exact (struct (field (ref (exact $b-exact)))))
+    (type $b-exact (struct (field (ref (exact $a-exact)))))
+  )
+
+  ;; CHECK:      (global $a-inexact (ref null $a-inexact) (ref.null none))
+  ;; NO_CD:      (global $a-inexact (ref null $a-inexact) (ref.null none))
+  (global $a-inexact (ref null $a-inexact) (ref.null none))
+  ;; CHECK:      (global $b-inexact (ref null $b-inexact) (ref.null none))
+  ;; NO_CD:      (global $b-inexact (ref null $b-inexact) (ref.null none))
+  (global $b-inexact (ref null $b-inexact) (ref.null none))
+  ;; CHECK:      (global $a-exact (ref null $a-exact) (ref.null none))
+  ;; NO_CD:      (global $a-exact (ref null $a-exact) (ref.null none))
+  (global $a-exact (ref null $a-exact) (ref.null none))
+  ;; CHECK:      (global $b-exact (ref null $b-exact) (ref.null none))
+  ;; NO_CD:      (global $b-exact (ref null $b-exact) (ref.null none))
+  (global $b-exact (ref null $b-exact) (ref.null none))
+)