[WasmGC] Heap2Local: Optimize RefIs and RefTest (#6705)

src/passes/Heap2Local.cpp

@@ -349,6 +349,12 @@ struct EscapeAnalyzer {
  void visitLocalSet(LocalSet* curr) { escapes = false; }
 
  // Reference operations. TODO add more
+ void visitRefIsNull(RefIsNull* curr) {
+ // The reference is compared to null, but nothing more.
+ escapes = false;
+ fullyConsumes = true;
+ }
+
  void visitRefEq(RefEq* curr) {
  // The reference is compared for identity, but nothing more.
  escapes = false;
@@ -367,6 +373,11 @@ struct EscapeAnalyzer {
  }
  }
 
+ void visitRefTest(RefTest* curr) {
+ escapes = false;
+ fullyConsumes = true;
+ }
+
  void visitRefCast(RefCast* curr) {
  // As it is our allocation that flows through here, we need to
  // check that the cast will not trap, so that we can continue
@@ -707,6 +718,17 @@ struct Struct2Local : PostWalker<Struct2Local> {
  replaceCurrent(builder.makeBlock(contents));
  }
 
+ void visitRefIsNull(RefIsNull* curr) {
+ if (!analyzer.reached.count(curr)) {
+ return;
+ }
+
+ // The result must be 0, since the allocation is not null. Drop the RefIs
+ // and append that.
+ replaceCurrent(builder.makeSequence(
+ builder.makeDrop(curr), builder.makeConst(Literal(int32_t(0)))));
+ }
+
  void visitRefEq(RefEq* curr) {
  if (!analyzer.reached.count(curr)) {
  return;
@@ -739,6 +761,23 @@ struct Struct2Local : PostWalker<Struct2Local> {
  replaceCurrent(curr->value);
  }
 
+ void visitRefTest(RefTest* curr) {
+ if (!analyzer.reached.count(curr)) {
+ return;
+ }
+
+ // This test operates on the allocation, which means we can compute whether
+ // it will succeed statically. We do not even need
+ // GCTypeUtils::evaluateCastCheck because we know the allocation's type
+ // precisely (it cannot be a strict subtype of the type - it is the type).
+ int32_t result = Type::isSubType(allocation->type, curr->castType);
+ // Remove the RefTest and leave only its reference child. If we kept it,
+ // we'd need to refinalize (as the input to the test changes, since the
+ // reference becomes a null, which has a different type).
+ replaceCurrent(builder.makeSequence(builder.makeDrop(curr->ref),
+ builder.makeConst(Literal(result))));
+ }
+
  void visitRefCast(RefCast* curr) {
  if (!analyzer.reached.count(curr)) {
  return;
@@ -820,12 +859,15 @@ struct Array2Struct : PostWalker<Array2Struct> {
  EscapeAnalyzer& analyzer;
  Function* func;
  Builder builder;
+ // The original type of the allocation, before we turn it into a struct.
+ Type originalType;
 
  Array2Struct(Expression* allocation,
  EscapeAnalyzer& analyzer,
  Function* func,
  Module& wasm)
- : allocation(allocation), analyzer(analyzer), func(func), builder(wasm) {
+ : allocation(allocation), analyzer(analyzer), func(func), builder(wasm),
+ originalType(allocation->type) {
 
  // Build the struct type we need: as many fields as the size of the array,
  // all of the same type as the array's element.
@@ -989,6 +1031,25 @@ struct Array2Struct : PostWalker<Array2Struct> {
  noteCurrentIsReached();
  }
 
+ // Some additional operations need special handling
+ void visitRefTest(RefTest* curr) {
+ if (!analyzer.reached.count(curr)) {
+ return;
+ }
+
+ // When we ref.test an array allocation, we cannot simply turn the array
+ // into a struct, as then the test will behave different. (Note that this is
+ // not a problem for ref.*cast*, as the cast simply goes away when the value
+ // flows through, and we verify it will do so in the escape analysis.) To
+ // handle this, check if the test succeeds or not, and write out the outcome
+ // here (similar to Struct2Local::visitRefTest). Note that we test on
+ // |originalType| here and not |allocation->type|, as the allocation has
+ // been turned into a struct.
+ int32_t result = Type::isSubType(originalType, curr->castType);
+ replaceCurrent(builder.makeSequence(builder.makeDrop(curr),
+ builder.makeConst(Literal(result))));
+ }
+
  // Get the value in an expression we know must contain a constant index.
  Index getIndex(Expression* curr) {
  return curr->cast<Const>()->value.getUnsigned();