TypeUpdating: Fix heap type updates on heap type children (#7848)

`GlobalTypeRewriter::getTempType` had the right logic to handle a
Type, which included the right logic to handle the HeapType of a
Type: Look in `typeIndices` and use that mapping if it is present. We
need the same thing for HeapTypes as well, which this adds in a
new `getTempHeapType`.

This was not noticed before because Continuations are the first
thing to have a direct HeapType child which is not through a Type.
So Types worked before, but not direct HeapType children, which
ended up mapped to the old type before the rewriting.

Changes to `wasm-type.h` are NFC but I think make it easier to read,
as it makes that logic entirely parallel to the lines above it.

One existing test has an update, which fixes it by removing the
ref to the old type.

Author: kripken

scripts/test/fuzzing.py

@@ -121,6 +121,7 @@
     'cont_export_throw.wast',
     'type-merging-cont.wast',
     'remove-unused-module-elements-cont.wast',
+    'abstract-type-refining-cont.wast',
     # TODO: fix split_wast() on tricky escaping situations like a string ending
     #       in \\" (the " is not escaped - there is an escaped \ before it)
     'string-lifting-section.wast',

src/ir/type-updating.cpp

@@ -359,9 +359,10 @@ Type GlobalTypeRewriter::getTempType(Type type) {
   }
   if (type.isRef()) {
     auto heapType = type.getHeapType();
-    if (auto it = typeIndices.find(heapType); it != typeIndices.end()) {
+    auto tempHeapType = getTempHeapType(heapType);
+    if (tempHeapType != heapType) {
       return typeBuilder.getTempRefType(
-        typeBuilder[it->second], type.getNullability(), type.getExactness());
+        tempHeapType, type.getNullability(), type.getExactness());
     }
     // This type is not one that is eligible for optimizing. That is fine; just
     // use it unmodified.
@@ -377,6 +378,13 @@ Type GlobalTypeRewriter::getTempType(Type type) {
   WASM_UNREACHABLE("bad type");
 }
 
+HeapType GlobalTypeRewriter::getTempHeapType(HeapType type) {
+  if (auto it = typeIndices.find(type); it != typeIndices.end()) {
+    return typeBuilder[it->second];
+  }
+  return type;
+}
+
 Type GlobalTypeRewriter::getTempTupleType(Tuple tuple) {
   return typeBuilder.getTempTupleType(tuple);
 }

src/ir/type-updating.h

@@ -404,6 +404,7 @@ class GlobalTypeRewriter {
   // so that they can use a proper temp type of the TypeBuilder while modifying
   // things.
   Type getTempType(Type type);
+  HeapType getTempHeapType(HeapType type);
   Type getTempTupleType(Tuple tuple);
 
   using SignatureUpdates = std::unordered_map<HeapType, Signature>;
@@ -496,7 +497,7 @@ class TypeMapper : public GlobalTypeRewriter {
     if (iter != mapping.end()) {
       return iter->second;
     }
-    return type;
+    return getTempHeapType(type);
   }
 
   Type getNewType(Type type) {

src/wasm-type.h

@@ -805,9 +805,12 @@ struct TypeBuilder {
         setHeapType(i, wasm::Array(elem));
         return;
       }
-      case HeapTypeKind::Cont:
-        setHeapType(i, Continuation(map(type.getContinuation().type)));
+      case HeapTypeKind::Cont: {
+        auto cont = type.getContinuation();
+        cont.type = map(cont.type);
+        setHeapType(i, cont);
         return;
+      }
       case HeapTypeKind::Basic:
         WASM_UNREACHABLE("unexpected kind");
     }

src/wasm/wasm-type.cpp

@@ -1324,6 +1324,9 @@ FeatureSet HeapType::getFeatures() const {
 
       // In addition, scan their non-ref children, to add dependencies on
       // things like SIMD.
+      // XXX This will not scan HeapType children that are not also children of
+      //     Type children, which happens with Continuation (has a HeapType
+      //     child that is not a Type).
       for (auto child : heapType.getTypeChildren()) {
         if (!child.isRef()) {
           feats |= child.getFeatures();

test/lit/passes/abstract-type-refining-cont.wast

@@ -0,0 +1,43 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --abstract-type-refining -all --closed-world -S -o - | filecheck %s
+
+;; $uncreated is never created, so we optimize and rebuild types. While doing
+;; so we should not get confused as we copy the continuation type, which
+;; should end up referring properly to the corresponding func type.
+
+(module
+ (rec
+  ;; CHECK:      (rec
+  ;; CHECK-NEXT:  (type $cont (cont $func))
+
+  ;; CHECK:       (type $func (func))
+  (type $func (func))
+  (type $cont (cont $func))
+
+  (type $uncreated (struct))
+ )
+
+ ;; CHECK:      (elem declare func $func)
+
+ ;; CHECK:      (func $func (type $func)
+ ;; CHECK-NEXT: )
+ (func $func (type $func)
+ )
+
+ ;; CHECK:      (func $test (type $func)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (cont.new $cont
+ ;; CHECK-NEXT:    (ref.func $func)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $test (type $func)
+  (drop
+   (cont.new $cont
+    (ref.func $func)
+   )
+  )
+ )
+)
+

test/lit/passes/type-merging-cont.wast

@@ -10,15 +10,12 @@
   ;; CHECK:      (rec
   ;; CHECK-NEXT:  (type $f (func))
   (type $f (func))
-  ;; CHECK:       (type $f_0 (func))
-
-  ;; CHECK:       (type $k (cont $f_0))
+  ;; CHECK:       (type $k (cont $f))
   (type $k (cont $f))
 
+  ;; CHECK:       (type $f-i32 (func (result i32)))
   (type $f-i32 (func (result i32)))
-  ;; CHECK:       (type $f-i32_0 (func (result i32)))
-
-  ;; CHECK:       (type $k-i32 (cont $f-i32_0))
+  ;; CHECK:       (type $k-i32 (cont $f-i32))
   (type $k-i32 (cont $f-i32))
 
   (rec