Reject rec groups that will collide after writing (#8144)

We allow the IR to contain types that are more refined than the enabled
feature set would normally allow. For example, the IR can contain exact
references even when custom descriptors are not enabled or typed
function references even when GC is not enabled. Using these more
refined types in the IR can make the optimizer more powerful. When we
write binaries, we generalize these refined types according to what is
allowed by the enabled feature set.

Generalizing the types in the binary writer makes it possible that two
rec groups that have different structures in the IR will end up having
the same structure once the binary is written out. This makes types that
were meant to be separate identical, possibly changing the observable
behavior of casts inadvertently.

To prevent that from happening, we must ensure that types that are meant
to be separate in the IR will also be separate after binary writing. We
already handle this when globally rewriting types (as of #8139), but
that is not enough to prevent the problem from occurring when the
original input already has rec groups that would collide after writing.
In general we have to allow overly-refined types to appear in the input
so we can test optimizations that take advantage of them.

Since we generally allow refined types in the input, there's nothing
stopping the fuzzer from randomly generating inputs and feature sets
that produce colliding rec groups. In such a case even a simple round
trip would change program behavior.

Avoid this problem by failing to build types when the TypeBuilder
contains distinct rec groups that would collide after binary writing.
Check for this condition by maintaining a UniqueRecGroups set while
types are being built.

Add an `insertOrGet` method to UniqueRecGroups to better support the use
case where conflicts need to be detected but not necessarily resolved.
Add `asWrittenGivenFeatures` methods to Type and HeapType, and use them
both from the binary writer and from wasm-type-shape to ensure that the
shape comparisons actually reflect the behavior of the binary writer.

Author: tlively

src/wasm-type-shape.h

@@ -19,7 +19,7 @@
 
 #include <functional>
 #include <list>
-#include <unordered_set>
+#include <unordered_map>
 #include <vector>
 
 #include "wasm-features.h"
@@ -162,17 +162,22 @@ struct BrandTypeIterator {
 // with an extra brand type at the end of the group that differentiates it from
 // previous group.
 struct UniqueRecGroups {
-  std::list<std::vector<HeapType>> groups;
-  std::unordered_set<RecGroupShape> shapes;
+  using Groups = std::list<std::vector<HeapType>>;
+  Groups groups;
+  std::unordered_map<RecGroupShape, Groups::iterator> shapes;
 
   FeatureSet features;
 
   UniqueRecGroups(FeatureSet features) : features(features) {}
 
   // Insert a rec group. If it is already unique, return the original types.
-  // Otherwise rebuild the group  make it unique and return the rebuilt types,
+  // Otherwise rebuild the group to make it unique and return the rebuilt types,
   // including the brand.
   const std::vector<HeapType>& insert(std::vector<HeapType> group);
+
+  // If the group is unique, insert it and return the types. Otherwise, return
+  // the types that already have this shape.
+  const std::vector<HeapType>& insertOrGet(std::vector<HeapType> group);
 };
 
 } // namespace wasm

src/wasm-type.h

@@ -256,6 +256,11 @@ class HeapType {
   // Returns the feature set required to use this type.
   FeatureSet getFeatures() const;
 
+  // We support more precise types in the IR than the enabled feature set would
+  // suggest. Get the generalized version of the type that will be written by
+  // the binary writer given the feature set.
+  inline HeapType asWrittenGivenFeatures(FeatureSet feats) const;
+
   // Helper allowing the value of `print(...)` to be sent to an ostream. Stores
   // a `TypeID` because `Type` is incomplete at this point and using a reference
   // makes it less convenient to use.
@@ -275,6 +280,16 @@ class HeapType {
   std::string toString() const;
 };
 
+HeapType HeapType::asWrittenGivenFeatures(FeatureSet feats) const {
+  // Without GC, only top types like func and extern are supported. The
+  // exception is string, since stringref can be enabled without GC and we still
+  // expect to write stringref types in that case.
+  if (!feats.hasGC() && *this != HeapType::string) {
+    return getTop();
+  }
+  return *this;
+}
+
 class Type {
   // The `id` uniquely represents each type, so type equality is just a
   // comparison of the ids. The basic types are packed at the bottom of the
@@ -411,7 +426,7 @@ class Type {
     return isRef() && getHeapType().isContinuation();
   }
   bool isDefaultable() const;
-  bool isCastable();
+  bool isCastable() const;
 
   // TODO: Allow this only for reference types.
   Nullability getNullability() const {
@@ -442,6 +457,11 @@ class Type {
     return !isExact() || feats.hasCustomDescriptors() ? *this : with(Inexact);
   }
 
+  // We support more precise types in the IR than the enabled feature set would
+  // suggest. Get the generalized version of the type that will be written by
+  // the binary writer given the feature set.
+  inline Type asWrittenGivenFeatures(FeatureSet feats) const;
+
 private:
   template<bool (Type::*pred)() const> bool hasPredicate() {
     for (const auto& type : *this) {
@@ -570,6 +590,20 @@ class Type {
   const Type& operator[](size_t i) const { return *Iterator{{this, i}}; }
 };
 
+Type Type::asWrittenGivenFeatures(FeatureSet feats) const {
+  if (!isRef()) {
+    return *this;
+  }
+  auto type = with(getHeapType().asWrittenGivenFeatures(feats));
+  if (!feats.hasGC()) {
+    type = type.with(Nullable);
+  }
+  if (!feats.hasCustomDescriptors()) {
+    type = type.with(Inexact);
+  }
+  return type;
+}
+
 namespace HeapTypes {
 
 constexpr HeapType ext = HeapType::ext;
@@ -870,6 +904,9 @@ struct TypeBuilder {
     InvalidUnsharedDescribes,
     // The custom descriptors feature is missing.
     RequiresCustomDescriptors,
+    // Two rec groups with different shapes would have the same shapes after
+    // the binary writer generalizes refined types that use disabled features.
+    RecGroupCollision,
   };
 
   struct Error {

src/wasm/wasm-binary.cpp

@@ -1798,23 +1798,8 @@ void WasmBinaryWriter::writeInlineBuffer(const char* data, size_t size) {
 }
 
 void WasmBinaryWriter::writeType(Type type) {
+  type = type.asWrittenGivenFeatures(wasm->features);
   if (type.isRef()) {
-    // The only reference types allowed without GC are funcref, externref, and
-    // exnref. We internally use more refined versions of those types, but we
-    // cannot emit those without GC.
-    if (!wasm->features.hasGC()) {
-      auto ht = type.getHeapType();
-      if (ht.isMaybeShared(HeapType::string)) {
-        // Do not overgeneralize stringref to anyref. We have tests that when a
-        // stringref is expected, we actually get a stringref. If we see a
-        // string, the stringref feature must be enabled.
-        type = Type(HeapTypes::string.getBasic(ht.getShared()), Nullable);
-      } else {
-        // Only the top type (func, extern, exn) is available, and only the
-        // nullable version.
-        type = Type(type.getHeapType().getTop(), Nullable);
-      }
-    }
     auto heapType = type.getHeapType();
     if (type.isNullable() && heapType.isBasic() && !heapType.isShared()) {
       switch (heapType.getBasic(Unshared)) {
@@ -1902,15 +1887,10 @@ void WasmBinaryWriter::writeType(Type type) {
 }
 
 void WasmBinaryWriter::writeHeapType(HeapType type, Exactness exactness) {
-  // ref.null always has a bottom heap type in Binaryen IR, but those types are
-  // only actually valid with GC. Otherwise, emit the corresponding valid top
-  // types instead.
+  type = type.asWrittenGivenFeatures(wasm->features);
   if (!wasm->features.hasCustomDescriptors()) {
     exactness = Inexact;
   }
-  if (!wasm->features.hasGC()) {
-    type = type.getTop();
-  }
   assert(!type.isBasic() || exactness == Inexact);
   if (exactness == Exact) {
     o << uint8_t(BinaryConsts::EncodedType::Exact);

src/wasm/wasm-type-shape.cpp

@@ -136,6 +136,10 @@ template<typename CompareTypes> struct RecGroupComparator {
   }
 
   Comparison compare(Type a, Type b) {
+    // Compare types as they will eventually be written out, not as they are in
+    // the IR.
+    a = a.asWrittenGivenFeatures(features);
+    b = b.asWrittenGivenFeatures(features);
     if (a.isBasic() != b.isBasic()) {
       return b.isBasic() < a.isBasic() ? LT : GT;
     }
@@ -152,14 +156,12 @@ 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;
     }
+    if (a.isExact() != b.isExact()) {
+      return a.isExact() < b.isExact() ? LT : GT;
+    }
     return compare(a.getHeapType(), b.getHeapType());
   }
 
@@ -294,6 +296,9 @@ struct RecGroupHasher {
   }
 
   size_t hash(Type type) {
+    // Hash types as they will eventually be written out, not as they are in the
+    // IR.
+    type = type.asWrittenGivenFeatures(features);
     size_t digest = wasm::hash(type.isBasic());
     if (type.isBasic()) {
       wasm::rehash(digest, type.getBasic());
@@ -305,10 +310,8 @@ struct RecGroupHasher {
       return digest;
     }
     assert(type.isRef());
-    if (features.hasCustomDescriptors()) {
-      wasm::rehash(digest, type.isExact());
-    }
     wasm::rehash(digest, type.isNullable());
+    wasm::rehash(digest, type.isExact());
     hash_combine(digest, hash(type.getHeapType()));
     return digest;
   }
@@ -372,15 +375,16 @@ bool ComparableRecGroupShape::operator>(const RecGroupShape& other) const {
 
 const std::vector<HeapType>&
 UniqueRecGroups::insert(std::vector<HeapType> types) {
-  auto& group = *groups.emplace(groups.end(), std::move(types));
-  if (shapes.emplace(RecGroupShape(group, features)).second) {
+  auto groupIt = groups.emplace(groups.end(), std::move(types));
+  auto& group = *groupIt;
+  if (shapes.emplace(RecGroupShape(group, features), groupIt).second) {
     // The types are already unique.
     return group;
   }
   // There is a conflict. Find a brand that makes the group unique.
   BrandTypeIterator brand;
   group.push_back(*brand);
-  while (!shapes.emplace(RecGroupShape(group, features)).second) {
+  while (!shapes.emplace(RecGroupShape(group, features), groupIt).second) {
     group.back() = *++brand;
   }
   // Rebuild the rec group to include the brand. Map the old types (excluding
@@ -405,6 +409,17 @@ UniqueRecGroups::insert(std::vector<HeapType> types) {
   return group;
 }
 
+const std::vector<HeapType>&
+UniqueRecGroups::insertOrGet(std::vector<HeapType> types) {
+  auto groupIt = groups.emplace(groups.end(), std::move(types));
+  auto [it, inserted] =
+    shapes.emplace(RecGroupShape(*groupIt, features), groupIt);
+  if (!inserted) {
+    groups.erase(groupIt);
+  }
+  return *it->second;
+}
+
 } // namespace wasm
 
 namespace std {

src/wasm/wasm-type.cpp

@@ -15,20 +15,15 @@
  */
 
 #include <algorithm>
-#include <array>
 #include <cassert>
-#include <map>
-#include <shared_mutex>
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
-#include <variant>
 
-#include "compiler-support.h"
 #include "support/hash.h"
-#include "support/insert_ordered.h"
 #include "wasm-features.h"
 #include "wasm-type-printing.h"
+#include "wasm-type-shape.h"
 #include "wasm-type.h"
 
 #define TRACE_CANONICALIZATION 0
@@ -623,7 +618,7 @@ bool Type::isDefaultable() const {
   return isConcrete() && !isNonNullable();
 }
 
-bool Type::isCastable() { return isRef() && getHeapType().isCastable(); }
+bool Type::isCastable() const { return isRef() && getHeapType().isCastable(); }
 
 unsigned Type::getByteSize() const {
   // TODO: alignment?
@@ -1468,6 +1463,9 @@ std::ostream& operator<<(std::ostream& os, TypeBuilder::ErrorReason reason) {
       return os << "Heap type describes an invalid unshared type";
     case TypeBuilder::ErrorReason::RequiresCustomDescriptors:
       return os << "custom descriptors required but not enabled";
+    case TypeBuilder::ErrorReason::RecGroupCollision:
+      return os
+             << "distinct rec groups would be identical after binary writing";
   }
   WASM_UNREACHABLE("Unexpected error reason");
 }
@@ -2260,7 +2258,19 @@ struct TypeBuilder::Impl {
   // We will validate features as we go.
   FeatureSet features;
 
-  Impl(size_t n, FeatureSet features) : entries(n), features(features) {}
+  // We allow some types to be used even if their corresponding features are not
+  // enabled. For example, we allow exact references without custom descriptors
+  // and typed function references without GC. Allowing these more-refined types
+  // in the IR helps the optimizer be more powerful. However, these disallowed
+  // refinements will be erased when a module is written out as a binary, which
+  // could cause distinct rec groups to become identical and potentially change
+  // the results of casts, etc. To avoid this, we must disallow building rec
+  // groups that vary only in some refinement that will be removed in binary
+  // writing. Track this with a UniqueRecGroups set, which is feature-aware.
+  UniqueRecGroups unique;
+
+  Impl(size_t n, FeatureSet features)
+    : entries(n), features(features), unique(features) {}
 };
 
 TypeBuilder::TypeBuilder(size_t n, FeatureSet features) {
@@ -2692,6 +2702,17 @@ TypeBuilder::BuildResult TypeBuilder::build() {
     assert(built->size() == groupSize);
     results.insert(results.end(), built->begin(), built->end());
 
+    // If we are building multiple groups, make sure there will be no conflicts
+    // after disallowed features are taken into account.
+    if (groupSize > 0 && groupSize != entryCount) {
+      auto expectedFirst = (*built)[0];
+      auto& types = impl->unique.insertOrGet(*built);
+      if (types[0] != expectedFirst) {
+        return {TypeBuilder::Error{
+          groupStart, TypeBuilder::ErrorReason::RecGroupCollision}};
+      }
+    }
+
     groupStart += groupSize;
   }
 

test/lit/binary/erase-trivial-cast-exactness.wast

@@ -0,0 +1,39 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s -all --disable-custom-descriptors --roundtrip -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $foo (struct))
+  (type $foo (struct))
+  ;; CHECK:      (func $trivial-exact-cast (type $1) (param $x i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (ref.cast (ref $foo)
+  ;; CHECK-NEXT:    (if (result (ref $foo))
+  ;; CHECK-NEXT:     (local.get $x)
+  ;; CHECK-NEXT:     (then
+  ;; CHECK-NEXT:      (struct.new_default $foo)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:     (else
+  ;; CHECK-NEXT:      (struct.new_default $foo)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $trivial-exact-cast (param $x i32)
+    (drop
+      ;; We allow trivial exact casts even without custom descriptors enabled,
+      ;; so this is valid. However, when we round trip, the exactness in the if
+      ;; result will be erased. If we fail to erase the exactness in the cast,
+      ;; we will be emitting a binary that uses a disabled feature, and also we
+      ;; will fail validation when we read the module back in because the cast
+      ;; will no longer be trivial.
+      (ref.cast (ref (exact $foo))
+        (if (result (ref (exact $foo)))
+          (local.get $x)
+          (then (struct.new $foo))
+          (else (struct.new $foo))
+        )
+      )
+    )
+  )
+)

test/lit/validation/no-type-collision-stringref.wast

@@ -0,0 +1,24 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
+;; RUN: wasm-opt %s -all --disable-gc -S -o - | filecheck %s
+
+;; No collision - we should not write a stringref as an externref.
+(module
+  ;; CHECK:      (type $A (func (param externref)))
+  (type $A (func (param externref)))
+  ;; CHECK:      (type $B (func (param stringref)))
+  (type $B (func (param stringref)))
+
+  ;; CHECK:      (func $a (param $0 externref)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $a (type $A)
+    (nop)
+  )
+
+  ;; CHECK:      (func $b (param $0 stringref)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $b (type $B)
+    (nop)
+  )
+)