Fix conflicting generalized types in TypeSSA

src/passes/TypeSSA.cpp

@@ -54,93 +54,37 @@
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/hash.h"
+#include "wasm-type-shape.h"
 #include "wasm.h"
 
 namespace wasm {
 
 namespace {
 
-// Given some TypeBuilder items that we want to build new types with, this
-// function builds the types in a new rec group.
-//
-// This is almost the same as just calling build(), but there is a risk of a
-// collision with an existing rec group. This function handles that by finding a
-// way to ensure that the new types are in fact in a new rec group.
-//
-// TODO: Move this outside if we find more uses.
-std::vector<HeapType> ensureTypesAreInNewRecGroup(RecGroup recGroup,
+// Ensure there are no conflicts between the newly built types and any existing
+// types.
+std::vector<HeapType> ensureTypesAreInNewRecGroup(std::vector<HeapType>&& types,
                                                   Module& wasm) {
-  auto num = recGroup.size();
-
-  std::vector<HeapType> types;
-  types.reserve(num);
-  for (auto type : recGroup) {
-    types.push_back(type);
+  std::unordered_set<RecGroup> existing;
+  for (auto type : ModuleUtils::collectHeapTypes(wasm)) {
+    existing.insert(type.getRecGroup());
   }
 
-  // Find all the heap types present before we create the new ones. The new
-  // types must not appear in |existingSet|.
-  std::vector<HeapType> existing = ModuleUtils::collectHeapTypes(wasm);
-  std::unordered_set<HeapType> existingSet(existing.begin(), existing.end());
-
-  // Check for a collision with an existing rec group. Note that it is enough to
-  // check one of the types: either the entire rec group gets merged, so they
-  // are all merged, or not.
-  if (existingSet.count(types[0])) {
-    // Unfortunately there is a conflict. Handle it by adding a "hash" - a
-    // "random" extra item in the rec group that is so outlandish it will
-    // surely (?) never collide with anything. We must loop while doing so,
-    // until we find a hash that does not collide.
-    //
-    // Note that we use uint64_t here, and deterministic_hash_combine below, to
-    // ensure our output is fully deterministic - the types we add here are
-    // observable in the output.
-    uint64_t hashSize = num + 10;
-    uint64_t random = num;
-    while (1) {
-      // Make a builder and add a slot for the hash.
-      TypeBuilder builder(num + 1);
-      for (Index i = 0; i < num; i++) {
-        builder[i].copy(types[i]);
-      }
-
-      // Implement the hash as a struct with "random" fields, and add it.
-      Struct hashStruct;
-      for (Index i = 0; i < hashSize; i++) {
-        // TODO: a denser encoding?
-        auto type = (random & 1) ? Type::i32 : Type::f64;
-        deterministic_hash_combine(random, hashSize + i);
-        hashStruct.fields.push_back(Field(type, Mutable));
-      }
-      builder[num] = hashStruct;
-
-      // Build and hope for the best.
-      builder.createRecGroup(0, num + 1);
-      auto result = builder.build();
-      assert(!result.getError());
-      types = *result;
-      assert(types.size() == num + 1);
-
-      if (existingSet.count(types[0])) {
-        // There is still a collision. Exponentially use larger hashes to
-        // quickly find one that works. Note that we also use different
-        // pseudorandom values while doing so in the for-loop above.
-        hashSize *= 2;
-      } else {
-        // Success! Leave the loop.
-        break;
-      }
-    }
+  UniqueRecGroups unique(wasm.features);
+  for (auto group : existing) {
+    std::vector<HeapType> types(group.begin(), group.end());
+    [[maybe_unused]] auto uniqueTypes = unique.insert(std::move(types));
+    assert(uniqueTypes.size() == group.size() && "unexpected collision");
   }
 
-#ifndef NDEBUG
-  // Verify the lack of a collision, just to be safe.
-  for (auto newType : types) {
-    assert(!existingSet.count(newType));
+  auto num = types.size();
+  std::vector<HeapType> uniqueTypes = unique.insert(std::move(types));
+  if (uniqueTypes.size() != num) {
+    // Remove the brand type, which we do not need to consider further.
+    uniqueTypes.pop_back();
   }
-#endif
-
-  return types;
+  assert(uniqueTypes.size() == num);
+  return uniqueTypes;
 }
 
 // A vector of struct.new or one of the variations on array.new.
@@ -407,8 +351,7 @@ struct TypeSSA : public Pass {
     assert(newTypes.size() == num);
 
     // Make sure this is actually a new rec group.
-    auto recGroup = newTypes[0].getRecGroup();
-    newTypes = ensureTypesAreInNewRecGroup(recGroup, *module);
+    newTypes = ensureTypesAreInNewRecGroup(std::move(newTypes), *module);
 
     // Success: we can apply the new types.
 

test/lit/passes/type-ssa-shared.wast

@@ -86,7 +86,7 @@
   ;; CHECK:      (rec
   ;; CHECK-NEXT:  (type $A_1 (sub $A (shared (array (mut i32)))))
 
-  ;; CHECK:       (type $4 (struct (field (mut i32)) (field (mut i32)) (field (mut f64)) (field (mut f64)) (field (mut i32)) (field (mut f64)) (field (mut f64)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32))))
+  ;; CHECK:       (type $4 (struct))
 
   ;; CHECK:      (func $func (type $1)
   ;; CHECK-NEXT:  (local $local (ref $B))

test/lit/passes/type-ssa.wast

@@ -460,7 +460,7 @@
   ;; CHECK:      (rec
   ;; CHECK-NEXT:  (type $array_1 (sub $array (array (mut f32))))
 
-  ;; CHECK:       (type $4 (struct (field (mut i32)) (field (mut i32)) (field (mut f64)) (field (mut f64)) (field (mut i32)) (field (mut f64)) (field (mut f64)) (field (mut i32)) (field (mut i32)) (field (mut i32)) (field (mut i32))))
+  ;; CHECK:       (type $4 (struct))
 
   ;; CHECK:      (func $1 (type $2) (param $ref (ref $subarray))
   ;; CHECK-NEXT:  (drop

test/lit/passes/types-ssa-generalized-exact.wast

@@ -0,0 +1,54 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: wasm-opt %s -all --disable-custom-descriptors \
+;; RUN:     --type-ssa --roundtrip --fuzz-exec -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $foo (struct))
+  (type $foo (struct))
+  ;; CHECK:      (type $super (sub (struct (field (ref $foo)))))
+  (type $super (sub (struct (field (ref $foo)))))
+  ;; CHECK:      (type $sub (sub $super (struct (field (ref $foo)))))
+  (type $sub (sub $super (struct (field (ref (exact $foo))))))
+
+  ;; CHECK:      (type $3 (func (result i32)))
+
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $super_1 (sub $super (struct (field (ref $foo)))))
+
+  ;; CHECK:       (type $5 (struct))
+
+  ;; CHECK:      (export "test" (func $test))
+
+  ;; CHECK:      (func $test (type $3) (result i32)
+  ;; CHECK-NEXT:  (local $sub (ref $sub))
+  ;; CHECK-NEXT:  (local $any anyref)
+  ;; CHECK-NEXT:  (ref.test (ref $sub)
+  ;; CHECK-NEXT:   (select (result anyref)
+  ;; CHECK-NEXT:    (struct.new $super_1
+  ;; CHECK-NEXT:     (struct.new_default $foo)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $any)
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $test (export "test") (result i32)
+    (local $sub (ref $sub))
+    (local $any anyref)
+    ;; TypeSSA will create another subtype of $super, which will differ from
+    ;; $sub because its field will not be exact. However, since exactness will
+    ;; be erased by the round trip, we still need a brand type to distinguish
+    ;; $sub and the new subtype. If we did not insert a brand type, this
+    ;; ref.test would incorrectly return 1 after optimization.
+    (ref.test (ref $sub)
+      ;; The select stops the ref.test from being optimized by finalization.
+      (select (result anyref)
+        (struct.new $super
+          (struct.new $foo)
+        )
+        (local.get $any)
+        (i32.const 1)
+      )
+    )
+  )
+)