TypeUpdating: Fix heap type updates on heap type children

scripts/test/fuzzing.py

@@ -120,6 +120,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