[Wasm GC] Add array support to TypeSSA (#5327)

Previously it only handled structs.

2 files changed, 309 insertions, 37 deletions

diff --git a/src/passes/TypeSSA.cpp b/src/passes/TypeSSA.cpp
index 75c2d4651..e25795cdc 100644
--- a/src/passes/TypeSSA.cpp
+++ b/src/passes/TypeSSA.cpp
@@ -57,17 +57,16 @@ namespace wasm {
 
 namespace {
 
-struct News {
-  std::vector<StructNew*> structNews;
-  // TODO: arrayNews of all kinds
-};
+// A vector of struct.new or one of the variations on array.new.
+using News = std::vector<Expression*>;
 
 struct NewFinder : public PostWalker<NewFinder> {
   News news;
 
-  void visitStructNew(StructNew* curr) { news.structNews.push_back(curr); }
-
-  // TODO arrays
+  void visitStructNew(StructNew* curr) { news.push_back(curr); }
+  void visitArrayNew(ArrayNew* curr) { news.push_back(curr); }
+  void visitArrayNewSeg(ArrayNewSeg* curr) { news.push_back(curr); }
+  void visitArrayInit(ArrayInit* curr) { news.push_back(curr); }
 };
 
 struct TypeSSA : public Pass {
@@ -118,16 +117,15 @@ struct TypeSSA : public Pass {
   Index nameCounter = 0;
 
   void processNews(const News& news) {
-    for (auto* curr : news.structNews) {
+    for (auto* curr : news) {
       if (isInteresting(curr)) {
-        newsToModify.structNews.push_back(curr);
+        newsToModify.push_back(curr);
       }
     }
   }
 
   void modifyNews() {
-    auto& structNews = newsToModify.structNews;
-    auto num = structNews.size();
+    auto num = newsToModify.size();
     if (num == 0) {
       return;
     }
@@ -148,9 +146,13 @@ struct TypeSSA : public Pass {
 
     TypeBuilder builder(num);
     for (Index i = 0; i < num; i++) {
-      auto* curr = structNews[i];
+      auto* curr = newsToModify[i];
       auto oldType = curr->type.getHeapType();
-      builder[i] = oldType.getStruct();
+      if (oldType.isStruct()) {
+        builder[i] = oldType.getStruct();
+      } else {
+        builder[i] = oldType.getArray();
+      }
       builder[i].subTypeOf(oldType);
     }
     builder.createRecGroup(0, num);
@@ -192,7 +194,7 @@ struct TypeSSA : public Pass {
     }
 
     for (Index i = 0; i < num; i++) {
-      auto* curr = structNews[i];
+      auto* curr = newsToModify[i];
       auto oldType = curr->type.getHeapType();
       auto newType = newTypes[i];
       curr->type = Type(newType, NonNullable);
@@ -210,41 +212,28 @@ struct TypeSSA : public Pass {
         existingTypeNames.insert(newName);
       }
     }
-
-    // TODO: arrays
   }
 
-  // An interesting StructNew, which we think is worth creating a new type for,
-  // is one that can be optimized better with a new type. That means it must
-  // have something interesting for optimizations to work with.
+  // An interesting *.new, which we think is worth creating a new type for, is
+  // one that can be optimized better with a new type. That means it must have
+  // something interesting for optimizations to work with.
   //
   // TODO: We may add new optimizations in the future that can benefit from more
   //       things, so it may be interesting to experiment with considering all
   //       news as "interesting" when we add major new type-based optimization
   //       passes.
-  bool isInteresting(StructNew* curr) {
+  bool isInteresting(Expression* curr) {
     if (curr->type == Type::unreachable) {
       // This is dead code anyhow.
       return false;
     }
 
-    if (curr->isWithDefault()) {
-      // This starts with all default values - zeros and nulls - and that might
-      // be useful.
-      //
-      // (A struct whose fields are all bottom types only has a single possible
-      // value in each field anyhow, so that is not interesting, but also
-      // unreasonable to occur in practice as other optimizations should handle
-      // it.)
-      return true;
-    }
-
-    // Look for at least one interesting operand.
-    auto& fields = curr->type.getHeapType().getStruct().fields;
-    for (Index i = 0; i < fields.size(); i++) {
-      assert(i <= curr->operands.size());
-      auto* operand = curr->operands[i];
-      if (operand->type != fields[i].type) {
+    // Look for at least one interesting operand. We will consider each operand
+    // against the declared type, that is, the type declared for where it is
+    // stored. If it has more information than the declared type then it is
+    // interesting.
+    auto isInterestingRelevantTo = [&](Expression* operand, Type declaredType) {
+      if (operand->type != declaredType) {
         // Excellent, this field has an interesting type - more refined than the
         // declared type, and which optimizations might benefit from.
         //
@@ -261,6 +250,59 @@ struct TypeSSA : public Pass {
         // This is a constant that passes may benefit from.
         return true;
       }
+
+      return false;
+    };
+
+    auto type = curr->type.getHeapType();
+
+    if (auto* structNew = curr->dynCast<StructNew>()) {
+      if (structNew->isWithDefault()) {
+        // This starts with all default values - zeros and nulls - and that
+        // might be useful.
+        //
+        // (An item whose fields are all bottom types only has a single possible
+        // value in each field anyhow, so that is not interesting, but also
+        // unreasonable to occur in practice as other optimizations should
+        // handle it.)
+        return true;
+      }
+
+      auto& fields = type.getStruct().fields;
+      for (Index i = 0; i < fields.size(); i++) {
+        assert(i <= structNew->operands.size());
+        if (isInterestingRelevantTo(structNew->operands[i], fields[i].type)) {
+          return true;
+        }
+      }
+    } else if (auto* arrayNew = curr->dynCast<ArrayNew>()) {
+      if (arrayNew->isWithDefault()) {
+        return true;
+      }
+
+      auto element = type.getArray().element;
+      if (isInterestingRelevantTo(arrayNew->init, element.type)) {
+        return true;
+      }
+    } else if (curr->is<ArrayNewSeg>()) {
+      // TODO: If the element segment is immutable perhaps we could inspect it.
+      return true;
+    } else if (auto* arrayInit = curr->dynCast<ArrayInit>()) {
+      // All the items must be interesting for us to consider this interesting,
+      // as we only track a single value for all indexes in the array, so one
+      // boring value means it is all boring.
+      //
+      // Note that we consider the empty array to be interesting (though atm no
+      // pass tracks the length - we might add one later though).
+      auto element = type.getArray().element;
+      for (auto* value : arrayInit->values) {
+        if (!isInterestingRelevantTo(value, element.type)) {
+          return false;
+        }
+      }
+      return true;
+    } else {
+      WASM_UNREACHABLE("unknown new");
     }
 
     // Nothing interesting.
diff --git a/test/lit/passes/type-ssa.wast b/test/lit/passes/type-ssa.wast
index f799cd871..790c6488a 100644
--- a/test/lit/passes/type-ssa.wast
+++ b/test/lit/passes/type-ssa.wast
@@ -219,3 +219,233 @@
     )
   )
 )
+
+(module
+  ;; CHECK:      (type $array (array (mut anyref)))
+  ;; NOMNL:      (type $array (array (mut anyref)))
+  (type $array (array (mut (ref null any))))
+
+  ;; CHECK:      (type $ref|i31|_anyref_=>_none (func (param (ref i31) anyref)))
+
+  ;; CHECK:      (type $array-func (array (mut funcref)))
+  ;; NOMNL:      (type $ref|i31|_anyref_=>_none (func (param (ref i31) anyref)))
+
+  ;; NOMNL:      (type $array$1 (array_subtype (mut anyref) $array))
+
+  ;; NOMNL:      (type $array$2 (array_subtype (mut anyref) $array))
+
+  ;; NOMNL:      (type $array$3 (array_subtype (mut anyref) $array))
+
+  ;; NOMNL:      (type $none_=>_none (func))
+
+  ;; NOMNL:      (type $array-func (array (mut funcref)))
+  (type $array-func (array (mut funcref)))
+
+
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $array$1 (array_subtype (mut anyref) $array))
+
+  ;; CHECK:       (type $array$2 (array_subtype (mut anyref) $array))
+
+  ;; CHECK:       (type $array$3 (array_subtype (mut anyref) $array))
+
+  ;; CHECK:       (type $array-func$4 (array_subtype (mut funcref) $array-func))
+
+  ;; CHECK:       (type $array$5 (array_subtype (mut anyref) $array))
+
+  ;; CHECK:       (type $array$6 (array_subtype (mut anyref) $array))
+
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (elem func $array.new)
+  ;; NOMNL:      (type $array-func$4 (array_subtype (mut funcref) $array-func))
+
+  ;; NOMNL:      (type $array$5 (array_subtype (mut anyref) $array))
+
+  ;; NOMNL:      (type $array$6 (array_subtype (mut anyref) $array))
+
+  ;; NOMNL:      (elem func $array.new)
+  (elem func $array.new)
+
+  ;; CHECK:      (func $array.new (type $ref|i31|_anyref_=>_none) (param $refined (ref i31)) (param $null-any anyref)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.new_default $array$1
+  ;; CHECK-NEXT:    (i32.const 5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.new $array$2
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:    (i32.const 5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.new $array$3
+  ;; CHECK-NEXT:    (local.get $refined)
+  ;; CHECK-NEXT:    (i32.const 5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.new $array
+  ;; CHECK-NEXT:    (local.get $null-any)
+  ;; CHECK-NEXT:    (i32.const 5)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $array.new (type $ref|i31|_anyref_=>_none) (param $refined (ref i31)) (param $null-any anyref)
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.new_default $array$1
+  ;; NOMNL-NEXT:    (i32.const 5)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.new $array$2
+  ;; NOMNL-NEXT:    (ref.null none)
+  ;; NOMNL-NEXT:    (i32.const 5)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.new $array$3
+  ;; NOMNL-NEXT:    (local.get $refined)
+  ;; NOMNL-NEXT:    (i32.const 5)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.new $array
+  ;; NOMNL-NEXT:    (local.get $null-any)
+  ;; NOMNL-NEXT:    (i32.const 5)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $array.new (param $refined (ref i31)) (param $null-any (ref null any))
+    ;; Default null, an interesting value, so we get a new type.
+    (drop
+      (array.new_default $array
+        (i32.const 5)
+      )
+    )
+    ;; Given null, also interesting.
+    (drop
+      (array.new $array
+        (ref.null none)
+        (i32.const 5)
+      )
+    )
+    ;; More refined type, interesting.
+    (drop
+      (array.new $array
+        (local.get $refined)
+        (i32.const 5)
+      )
+    )
+    ;; Same type as declared - boring, no new type.
+    (drop
+      (array.new $array
+        (local.get $null-any)
+        (i32.const 5)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.new_seg (type $none_=>_none)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.new_elem $array-func$4 0
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:    (i32.const 3)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $array.new_seg (type $none_=>_none)
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.new_elem $array-func$4 0
+  ;; NOMNL-NEXT:    (i32.const 0)
+  ;; NOMNL-NEXT:    (i32.const 3)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $array.new_seg
+    ;; We consider all new_elem to be interesting as we don't look at the elem
+    ;; data yet.
+    (drop
+      (array.new_elem $array-func 0
+        (i32.const 0)
+        (i32.const 3)
+      )
+    )
+  )
+
+  ;; CHECK:      (func $array.init (type $ref|i31|_anyref_=>_none) (param $refined (ref i31)) (param $null-any anyref)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.init_static $array$5
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.init_static $array$6
+  ;; CHECK-NEXT:    (local.get $refined)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.init_static $array
+  ;; CHECK-NEXT:    (local.get $null-any)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (array.init_static $array
+  ;; CHECK-NEXT:    (local.get $refined)
+  ;; CHECK-NEXT:    (local.get $null-any)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; NOMNL:      (func $array.init (type $ref|i31|_anyref_=>_none) (param $refined (ref i31)) (param $null-any anyref)
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.init_static $array$5
+  ;; NOMNL-NEXT:    (ref.null none)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.init_static $array$6
+  ;; NOMNL-NEXT:    (local.get $refined)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.init_static $array
+  ;; NOMNL-NEXT:    (local.get $null-any)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT:  (drop
+  ;; NOMNL-NEXT:   (array.init_static $array
+  ;; NOMNL-NEXT:    (local.get $refined)
+  ;; NOMNL-NEXT:    (local.get $null-any)
+  ;; NOMNL-NEXT:   )
+  ;; NOMNL-NEXT:  )
+  ;; NOMNL-NEXT: )
+  (func $array.init (param $refined (ref i31)) (param $null-any (ref null any))
+    ;; Null, interesting, so we get a new type.
+    (drop
+      (array.init_static $array
+        (ref.null none)
+      )
+    )
+    ;; More refined type, interesting.
+    (drop
+      (array.init_static $array
+        (local.get $refined)
+      )
+    )
+    ;; Same type as declared - boring, no new type.
+    (drop
+      (array.init_static $array
+        (local.get $null-any)
+      )
+    )
+    ;; Mixture of boring and interesting => boring (since we infer a single type
+    ;; for the entire array).
+    (drop
+      (array.init_static $array
+        (local.get $refined)
+        (local.get $null-any)
+      )
+    )
+  )
+)