1 files changed, 62 insertions, 1 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index 39c77bec4..7aec23ccf 100644
--- a/src/passes/Heap2Local.cpp
+++ b/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();