Fix RTTs for RTT-less instructions (#4294)

Allocation and cast instructions without explicit RTTs should use the canonical
RTTs for the given types. Furthermore, the RTTs for nominal types should reflect
the static type hierarchy. Previously, however, we implemented allocations and
casts without RTTs using an alternative system that only used static types
rather than RTT values. This alternative system would work fine in a world
without first-class RTTs, but it did not properly allow mixing instructions that
use RTTs and instructions that do not use RTTs as intended by the M4 GC spec.

This PR fixes the issue by using canonical RTTs where appropriate and cleans up
the relevant casting code using std::variant.

7 files changed, 160 insertions, 155 deletions

diff --git a/src/literal.h b/src/literal.h
index ce3f80902..ae0d4253f 100644
--- a/src/literal.h
+++ b/src/literal.h
@@ -62,9 +62,6 @@ class Literal {
     // as the Literal class itself.
     // To support the experimental RttFreshSub instruction, we not only store
     // the type, but also a reference to an allocation.
-    // The above describes dynamic data, that is with an actual RTT. The static
-    // case just has a static type in its GCData.
-    // See struct RttSuper below for more details.
     std::unique_ptr<RttSupers> rttSupers;
     // TODO: Literals of type `externref` can only be `null` currently but we
     // will need to represent extern values eventually, to
@@ -264,6 +261,10 @@ public:
     return lit;
   }
 
+  // Get the canonical RTT value for a given HeapType. For nominal types, the
+  // canonical RTT reflects the static supertype chain.
+  static Literal makeCanonicalRtt(HeapType type);
+
   Literal castToF32();
   Literal castToF64();
   Literal castToI32();
@@ -700,23 +701,18 @@ std::ostream& operator<<(std::ostream& o, wasm::Literals literals);
 // is. In the case of static (rtt-free) typing, the rtt is not present and
 // instead we have a static type.
 struct GCData {
-  // Either the RTT or the type must be present, but not both.
-  std::variant<Literal, HeapType> typeInfo;
+  // The runtime type info for this struct or array.
+  Literal rtt;
 
+  // The element or field values.
   Literals values;
 
-  GCData(HeapType type, Literals values) : typeInfo(type), values(values) {}
-  GCData(Literal rtt, Literals values) : typeInfo(rtt), values(values) {}
-
-  bool hasRtt() { return std::get_if<Literal>(&typeInfo); }
-  bool hasHeapType() { return std::get_if<HeapType>(&typeInfo); }
-  Literal getRtt() { return std::get<Literal>(typeInfo); }
-  HeapType getHeapType() { return std::get<HeapType>(typeInfo); }
+  GCData(Literal rtt, Literals values) : rtt(rtt), values(values) {}
 };
 
 struct RttSuper {
   // The type of the super.
-  Type type;
+  HeapType type;
   // A shared allocation, used to implement rtt.fresh_sub. This is null for a
   // normal sub, and for a fresh one we allocate a value here, which can then be
   // used to differentiate rtts. (The allocation is shared so that when copying
@@ -724,7 +720,7 @@ struct RttSuper {
   // TODO: Remove or optimize this when the spec stabilizes.
   std::shared_ptr<size_t> freshPtr;
 
-  RttSuper(Type type) : type(type) {}
+  RttSuper(HeapType type) : type(type) {}
 
   void makeFresh() { freshPtr = std::make_shared<size_t>(); }
 
diff --git a/src/wasm-builder.h b/src/wasm-builder.h
index b192afbac..5914a8afd 100644
--- a/src/wasm-builder.h
+++ b/src/wasm-builder.h
@@ -837,7 +837,7 @@ public:
   }
   RttCanon* makeRttCanon(HeapType heapType) {
     auto* ret = wasm.allocator.alloc<RttCanon>();
-    ret->type = Type(Rtt(0, heapType));
+    ret->type = Type(Rtt(heapType.getDepth(), heapType));
     ret->finalize();
     return ret;
   }
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 61a08c14f..c67ac6768 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -26,6 +26,7 @@
 #include <cmath>
 #include <limits.h>
 #include <sstream>
+#include <variant>
 
 #include "ir/module-utils.h"
 #include "support/bits.h"
@@ -1404,154 +1405,139 @@ public:
   // Helper for ref.test, ref.cast, and br_on_cast, which share almost all their
   // logic except for what they return.
   struct Cast {
-    enum Outcome {
-      // We took a break before doing anything.
-      Break,
-      // The input was null.
-      Null,
-      // The cast succeeded.
-      Success,
-      // The cast failed.
-      Failure
-    } outcome;
-
-    Flow breaking;
-    Literal originalRef;
-    Literal castRef;
+    // The control flow that preempts the cast.
+    struct Breaking : Flow {
+      Breaking(Flow breaking) : Flow(breaking) {}
+    };
+    // The null input to the cast.
+    struct Null : Literal {
+      Null(Literal original) : Literal(original) {}
+    };
+    // The result of the successful cast.
+    struct Success : Literal {
+      Success(Literal result) : Literal(result) {}
+    };
+    // The input to a failed cast.
+    struct Failure : Literal {
+      Failure(Literal original) : Literal(original) {}
+    };
+
+    std::variant<Breaking, Null, Success, Failure> state;
+
+    template<class T> Cast(T state) : state(state) {}
+    Flow* getBreaking() { return std::get_if<Breaking>(&state); }
+    Literal* getNull() { return std::get_if<Null>(&state); }
+    Literal* getSuccess() { return std::get_if<Success>(&state); }
+    Literal* getFailure() { return std::get_if<Failure>(&state); }
+    Literal* getNullOrFailure() {
+      if (auto* original = getNull()) {
+        return original;
+      } else {
+        return getFailure();
+      }
+    }
   };
 
   template<typename T> Cast doCast(T* curr) {
-    Cast cast;
     Flow ref = this->visit(curr->ref);
     if (ref.breaking()) {
-      cast.outcome = cast.Break;
-      cast.breaking = ref;
-      return cast;
+      return typename Cast::Breaking{ref};
     }
+    // The RTT value for the type we are trying to cast to.
     Literal intendedRtt;
     if (curr->rtt) {
-      // This is a dynamic check with an rtt.
+      // This is a dynamic check with an RTT.
       Flow rtt = this->visit(curr->rtt);
       if (rtt.breaking()) {
-        cast.outcome = cast.Break;
-        cast.breaking = rtt;
-        return cast;
+        return typename Cast::Breaking{ref};
       }
       intendedRtt = rtt.getSingleValue();
+    } else {
+      // If there is no explicit RTT, use the canonical RTT for the static type.
+      intendedRtt = Literal::makeCanonicalRtt(curr->intendedType);
     }
-    cast.originalRef = ref.getSingleValue();
-    if (cast.originalRef.isNull()) {
-      cast.outcome = cast.Null;
-      return cast;
+    Literal original = ref.getSingleValue();
+    if (original.isNull()) {
+      return typename Cast::Null{original};
     }
     // The input may not be GC data or a function; for example it could be an
-    // anyref of null (already handled above) or anything else (handled here,
-    // but this is for future use as atm the binaryen interpreter cannot
-    // represent external references).
-    if (!cast.originalRef.isData() && !cast.originalRef.isFunction()) {
-      cast.outcome = cast.Failure;
-      return cast;
-    }
-    if (cast.originalRef.isFunction()) {
-      // Function casts are simple in that they have no RTT hierarchies; instead
-      // each reference has the canonical RTT for the signature.
-      // We must have a module in order to perform the cast, to get the type. If
-      // we do not have one, or if the function is not present (which may happen
-      // if we are optimizing a function before the entire module is built),
-      // then this is something we cannot precompute.
-      auto* func = module
-                     ? module->getFunctionOrNull(cast.originalRef.getFunc())
-                     : nullptr;
+    // externref or an i31. The cast definitely fails in these cases.
+    if (!original.isData() && !original.isFunction()) {
+      return typename Cast::Failure{original};
+    }
+    Literal actualRtt;
+    if (original.isFunction()) {
+      // Function references always have the canonical RTTs of the functions
+      // they reference. We must have a module to look up the function's type to
+      // get that canonical RTT.
+      auto* func =
+        module ? module->getFunctionOrNull(original.getFunc()) : nullptr;
       if (!func) {
-        cast.outcome = cast.Break;
-        cast.breaking = NONCONSTANT_FLOW;
-        return cast;
-      }
-      if (curr->rtt) {
-        Literal seenRtt = Literal(Type(Rtt(0, func->type)));
-        if (!seenRtt.isSubRtt(intendedRtt)) {
-          cast.outcome = cast.Failure;
-          return cast;
-        }
-        cast.castRef = Literal(
-          func->name, Type(intendedRtt.type.getHeapType(), NonNullable));
-      } else {
-        if (!HeapType::isSubType(func->type, curr->intendedType)) {
-          cast.outcome = cast.Failure;
-          return cast;
-        }
-        cast.castRef =
-          Literal(func->name, Type(curr->intendedType, NonNullable));
+        return typename Cast::Breaking{NONCONSTANT_FLOW};
       }
+      actualRtt = Literal::makeCanonicalRtt(func->type);
     } else {
-      // GC data store an RTT in each instance.
-      assert(cast.originalRef.isData());
-      auto gcData = cast.originalRef.getGCData();
-      assert(bool(curr->rtt) == gcData->hasRtt());
-      if (curr->rtt) {
-        auto seenRtt = gcData->getRtt();
-        if (!seenRtt.isSubRtt(intendedRtt)) {
-          cast.outcome = cast.Failure;
-          return cast;
-        }
-        cast.castRef =
-          Literal(gcData, Type(intendedRtt.type.getHeapType(), NonNullable));
+      assert(original.isData());
+      actualRtt = original.getGCData()->rtt;
+    };
+    // We have the actual and intended RTTs, so perform the cast.
+    if (actualRtt.isSubRtt(intendedRtt)) {
+      Type resultType(intendedRtt.type.getHeapType(), NonNullable);
+      if (original.isFunction()) {
+        return typename Cast::Success{Literal{original.getFunc(), resultType}};
       } else {
-        auto seenType = gcData->getHeapType();
-        if (!HeapType::isSubType(seenType, curr->intendedType)) {
-          cast.outcome = cast.Failure;
-          return cast;
-        }
-        cast.castRef = Literal(gcData, Type(seenType, NonNullable));
+        return
+          typename Cast::Success{Literal(original.getGCData(), resultType)};
       }
+    } else {
+      return typename Cast::Failure{original};
     }
-    cast.outcome = cast.Success;
-    return cast;
   }
 
   Flow visitRefTest(RefTest* curr) {
     NOTE_ENTER("RefTest");
     auto cast = doCast(curr);
-    if (cast.outcome == cast.Break) {
-      return cast.breaking;
+    if (auto* breaking = cast.getBreaking()) {
+      return *breaking;
+    } else {
+      return Literal(int32_t(bool(cast.getSuccess())));
     }
-    return Literal(int32_t(cast.outcome == cast.Success));
   }
   Flow visitRefCast(RefCast* curr) {
     NOTE_ENTER("RefCast");
     auto cast = doCast(curr);
-    if (cast.outcome == cast.Break) {
-      return cast.breaking;
-    }
-    if (cast.outcome == cast.Null) {
+    if (auto* breaking = cast.getBreaking()) {
+      return *breaking;
+    } else if (cast.getNull()) {
       return Literal::makeNull(Type(curr->type.getHeapType(), Nullable));
+    } else if (auto* result = cast.getSuccess()) {
+      return *result;
     }
-    if (cast.outcome == cast.Failure) {
-      trap("cast error");
-    }
-    assert(cast.outcome == cast.Success);
-    return cast.castRef;
+    assert(cast.getFailure());
+    trap("cast error");
+    WASM_UNREACHABLE("unreachable");
   }
   Flow visitBrOn(BrOn* curr) {
     NOTE_ENTER("BrOn");
-    // BrOnCast* uses the casting infrastructure, so handle it first.
+    // BrOnCast* uses the casting infrastructure, so handle them first.
     if (curr->op == BrOnCast || curr->op == BrOnCastFail) {
       auto cast = doCast(curr);
-      if (cast.outcome == cast.Break) {
-        return cast.breaking;
-      }
-      if (cast.outcome == cast.Null || cast.outcome == cast.Failure) {
+      if (auto* breaking = cast.getBreaking()) {
+        return *breaking;
+      } else if (auto* original = cast.getNullOrFailure()) {
         if (curr->op == BrOnCast) {
-          return cast.originalRef;
+          return *original;
         } else {
-          return Flow(curr->name, cast.originalRef);
+          return Flow(curr->name, *original);
         }
-      }
-      assert(cast.outcome == cast.Success);
-      if (curr->op == BrOnCast) {
-        return Flow(curr->name, cast.castRef);
       } else {
-        return cast.castRef;
+        auto* result = cast.getSuccess();
+        assert(result);
+        if (curr->op == BrOnCast) {
+          return Flow(curr->name, *result);
+        } else {
+          return *result;
+        }
       }
     }
     // The others do a simpler check for the type.
@@ -1619,7 +1605,9 @@ public:
     }
     return {value};
   }
-  Flow visitRttCanon(RttCanon* curr) { return Literal(curr->type); }
+  Flow visitRttCanon(RttCanon* curr) {
+    return Literal::makeCanonicalRtt(curr->type.getHeapType());
+  }
   Flow visitRttSub(RttSub* curr) {
     Flow parent = this->visit(curr->parent);
     if (parent.breaking()) {
@@ -1627,34 +1615,22 @@ public:
     }
     auto parentValue = parent.getSingleValue();
     auto newSupers = std::make_unique<RttSupers>(parentValue.getRttSupers());
-    newSupers->push_back(parentValue.type);
+    newSupers->push_back(parentValue.type.getHeapType());
     if (curr->fresh) {
       newSupers->back().makeFresh();
     }
     return Literal(std::move(newSupers), curr->type);
   }
 
-  // Generates GC data for either dynamic (with an RTT) or static (with a type)
-  // typing. Dynamic typing will provide an rtt expression and an rtt flow with
-  // the value, while static typing only provides a heap type directly.
-  template<typename T>
-  std::shared_ptr<GCData>
-  makeGCData(Expression* rttExpr, Flow& rttFlow, HeapType type, T& data) {
-    if (rttExpr) {
-      return std::make_shared<GCData>(rttFlow.getSingleValue(), data);
-    } else {
-      return std::make_shared<GCData>(type, data);
-    }
-  }
-
   Flow visitStructNew(StructNew* curr) {
     NOTE_ENTER("StructNew");
-    Flow rtt;
+    Literal rttVal;
     if (curr->rtt) {
-      rtt = this->visit(curr->rtt);
+      Flow rtt = this->visit(curr->rtt);
       if (rtt.breaking()) {
         return rtt;
       }
+      rttVal = rtt.getSingleValue();
     }
     if (curr->type == Type::unreachable) {
       // We cannot proceed to compute the heap type, as there isn't one. Just
@@ -1681,8 +1657,10 @@ public:
         data[i] = value.getSingleValue();
       }
     }
-    return Flow(
-      Literal(makeGCData(curr->rtt, rtt, heapType, data), curr->type));
+    if (!curr->rtt) {
+      rttVal = Literal::makeCanonicalRtt(heapType);
+    }
+    return Literal(std::make_shared<GCData>(rttVal, data), curr->type);
   }
   Flow visitStructGet(StructGet* curr) {
     NOTE_ENTER("StructGet");
@@ -1725,12 +1703,13 @@ public:
 
   Flow visitArrayNew(ArrayNew* curr) {
     NOTE_ENTER("ArrayNew");
-    Flow rtt;
+    Literal rttVal;
     if (curr->rtt) {
-      rtt = this->visit(curr->rtt);
+      Flow rtt = this->visit(curr->rtt);
       if (rtt.breaking()) {
         return rtt;
       }
+      rttVal = rtt.getSingleValue();
     }
     auto size = this->visit(curr->size);
     if (size.breaking()) {
@@ -1765,17 +1744,20 @@ public:
         data[i] = value;
       }
     }
-    return Flow(
-      Literal(makeGCData(curr->rtt, rtt, heapType, data), curr->type));
+    if (!curr->rtt) {
+      rttVal = Literal::makeCanonicalRtt(heapType);
+    }
+    return Literal(std::make_shared<GCData>(rttVal, data), curr->type);
   }
   Flow visitArrayInit(ArrayInit* curr) {
     NOTE_ENTER("ArrayInit");
-    Flow rtt;
+    Literal rttVal;
     if (curr->rtt) {
-      rtt = this->visit(curr->rtt);
+      Flow rtt = this->visit(curr->rtt);
       if (rtt.breaking()) {
         return rtt;
       }
+      rttVal = rtt.getSingleValue();
     }
     Index num = curr->values.size();
     if (num >= ArrayLimit) {
@@ -1802,8 +1784,10 @@ public:
       }
       data[i] = truncateForPacking(value.getSingleValue(), field);
     }
-    return Flow(
-      Literal(makeGCData(curr->rtt, rtt, heapType, data), curr->type));
+    if (!curr->rtt) {
+      rttVal = Literal::makeCanonicalRtt(heapType);
+    }
+    return Literal(std::make_shared<GCData>(rttVal, data), curr->type);
   }
   Flow visitArrayGet(ArrayGet* curr) {
     NOTE_ENTER("ArrayGet");
diff --git a/src/wasm-type.h b/src/wasm-type.h
index 4c401be6b..1f6c5a7de 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -390,6 +390,10 @@ public:
   // `TypeSystem::Equirecursive` mode.
   std::optional<HeapType> getSuperType() const;
 
+  // Return the depth of this heap type in the nominal type hierarchy, i.e. the
+  // number of supertypes in its supertype chain.
+  size_t getDepth() const;
+
   // Whether this is a nominal type in the sense of being a GC Milestone 4
   // nominal type. Although all non-basic HeapTypes are nominal in
   // `TypeSystem::Nominal` mode, this will still return false unless the type is
diff --git a/src/wasm/literal.cpp b/src/wasm/literal.cpp
index 131a6d443..4753e5ca1 100644
--- a/src/wasm/literal.cpp
+++ b/src/wasm/literal.cpp
@@ -38,8 +38,7 @@ Literal::Literal(Type type) : type(type) {
     if (isData()) {
       new (&gcData) std::shared_ptr<GCData>();
     } else if (type.isRtt()) {
-      // Allocate a new RttSupers (with no data).
-      new (&rttSupers) auto(std::make_unique<RttSupers>());
+      new (this) Literal(Literal::makeCanonicalRtt(type.getHeapType()));
     } else {
       memset(&v128, 0, 16);
     }
@@ -148,6 +147,18 @@ Literal& Literal::operator=(const Literal& other) {
   return *this;
 }
 
+Literal Literal::makeCanonicalRtt(HeapType type) {
+  auto supers = std::make_unique<RttSupers>();
+  std::optional<HeapType> supertype;
+  for (auto curr = type; (supertype = curr.getSuperType()); curr = *supertype) {
+    supers->emplace_back(*supertype);
+  }
+  // We want the highest types to be first.
+  std::reverse(supers->begin(), supers->end());
+  size_t depth = supers->size();
+  return Literal(std::move(supers), Type(Rtt(depth, type)));
+}
+
 template<typename LaneT, int Lanes>
 static void extractBytes(uint8_t (&dest)[16], const LaneArray<Lanes>& lanes) {
   std::array<uint8_t, 16> bytes;
@@ -495,9 +506,7 @@ std::ostream& operator<<(std::ostream& o, Literal literal) {
     if (literal.isData()) {
       auto data = literal.getGCData();
       if (data) {
-        o << "[ref ";
-        std::visit([&](auto& info) { o << info; }, data->typeInfo);
-        o << ' ' << data->values << ']';
+        o << "[ref " << data->rtt << ' ' << data->values << ']';
       } else {
         o << "[ref null " << literal.type << ']';
       }
@@ -2560,7 +2569,7 @@ bool Literal::isSubRtt(const Literal& other) const {
   // we have the same amount of supers, and must be completely identical to
   // other.
   if (otherSupers.size() < supers.size()) {
-    return other.type == supers[otherSupers.size()].type;
+    return other.type.getHeapType() == supers[otherSupers.size()].type;
   } else {
     return other.type == type;
   }
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 94c6e96c6..002b401b5 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -1217,6 +1217,15 @@ std::optional<HeapType> HeapType::getSuperType() const {
   return {};
 }
 
+size_t HeapType::getDepth() const {
+  size_t depth = 0;
+  std::optional<HeapType> super;
+  for (auto curr = *this; (super = curr.getSuperType()); curr = *super) {
+    ++depth;
+  }
+  return depth;
+}
+
 bool HeapType::isNominal() const {
   if (isBasic()) {
     return false;
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index f8669961c..bb0e99cc2 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -2361,7 +2361,10 @@ void FunctionValidator::visitRttCanon(RttCanon* curr) {
     getModule()->features.hasGC(), curr, "rtt.canon requires gc to be enabled");
   shouldBeTrue(curr->type.isRtt(), curr, "rtt.canon must have RTT type");
   auto rtt = curr->type.getRtt();
-  shouldBeEqual(rtt.depth, Index(0), curr, "rtt.canon has a depth of 0");
+  shouldBeEqual(rtt.depth,
+                Index(curr->type.getHeapType().getDepth()),
+                curr,
+                "rtt.canon must have the depth of its heap type");
 }
 
 void FunctionValidator::visitRttSub(RttSub* curr) {