From d7955a34e332a78cc71bf77800114cba008185bb Mon Sep 17 00:00:00 2001
From: Thomas Lively <tlively@google.com>
Date: Tue, 13 Aug 2024 00:32:49 -0400
Subject: Add a TypeBuilder API for copying a heap type (#6828)

Given a function that maps the old child heap types to new child heap
types, the new API takes care of copying the rest of the structure of a
given heap type into a TypeBuilder slot.

Use the new API in GlobalTypeRewriter::rebuildTypes. It will also be
used in an upcoming type optimization. This refactoring also required
adding the ability to clear the supertype of a TypeBuilder slot, which
was previously not possible.
---
 src/wasm-type.h | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 70 insertions(+), 2 deletions(-)

(limited to 'src/wasm-type.h')

diff --git a/src/wasm-type.h b/src/wasm-type.h
index e2eec5f35..ff1358a12 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -660,6 +660,67 @@ struct TypeBuilder {
   void setHeapType(size_t i, Struct&& struct_);
   void setHeapType(size_t i, Array array);
 
+  // Sets the heap type at index `i` to be a copy of the given heap type with
+  // its referenced HeapTypes to be replaced according to the provided mapping
+  // function.
+  template<typename F> void copyHeapType(size_t i, HeapType type, F map) {
+    assert(!type.isBasic());
+    if (auto super = type.getDeclaredSuperType()) {
+      setSubType(i, map(*super));
+    }
+    setOpen(i, type.isOpen());
+    setShared(i, type.getShared());
+
+    auto copySingleType = [&](Type t) -> Type {
+      if (t.isBasic()) {
+        return t;
+      }
+      assert(t.isRef());
+      return getTempRefType(map(t.getHeapType()), t.getNullability());
+    };
+    auto copyType = [&](Type t) -> Type {
+      if (t.isTuple()) {
+        std::vector<Type> elems;
+        elems.reserve(t.size());
+        for (auto elem : t) {
+          elems.push_back(copySingleType(elem));
+        }
+        return getTempTupleType(elems);
+      }
+      return copySingleType(t);
+    };
+    switch (type.getKind()) {
+      case HeapTypeKind::Func: {
+        auto sig = type.getSignature();
+        setHeapType(i, Signature(copyType(sig.params), copyType(sig.results)));
+        return;
+      }
+      case HeapTypeKind::Struct: {
+        const auto& struct_ = type.getStruct();
+        std::vector<Field> fields;
+        fields.reserve(struct_.fields.size());
+        for (auto field : struct_.fields) {
+          field.type = copyType(field.type);
+          fields.push_back(field);
+        }
+        setHeapType(i, Struct(fields));
+        return;
+      }
+      case HeapTypeKind::Array: {
+        auto elem = type.getArray().element;
+        elem.type = copyType(elem.type);
+        // MSVC gets confused without this disambiguation.
+        setHeapType(i, wasm::Array(elem));
+        return;
+      }
+      case HeapTypeKind::Cont:
+        setHeapType(i, Continuation(map(type.getContinuation().type)));
+        return;
+      case HeapTypeKind::Basic:
+        WASM_UNREACHABLE("unexpected kind");
+    }
+  }
+
   // Gets the temporary HeapType at index `i`. This HeapType should only be used
   // to construct temporary Types using the methods below.
   HeapType getTempHeapType(size_t i);
@@ -672,7 +733,7 @@ struct TypeBuilder {
 
   // Declare the HeapType being built at index `i` to be an immediate subtype of
   // the given HeapType.
-  void setSubType(size_t i, HeapType super);
+  void setSubType(size_t i, std::optional<HeapType> super);
 
   // Create a new recursion group covering slots [i, i + length). Groups must
   // not overlap or go out of bounds.
@@ -746,7 +807,7 @@ struct TypeBuilder {
       builder.setHeapType(index, array);
       return *this;
     }
-    Entry& subTypeOf(HeapType other) {
+    Entry& subTypeOf(std::optional<HeapType> other) {
       builder.setSubType(index, other);
       return *this;
     }
@@ -758,6 +819,13 @@ struct TypeBuilder {
       builder.setShared(index, share);
       return *this;
     }
+    template<typename F> Entry& copy(HeapType type, F map) {
+      builder.copyHeapType(index, type, map);
+      return *this;
+    }
+    Entry& copy(HeapType type) {
+      return copy(type, [](HeapType t) { return t; });
+    }
   };
 
   Entry operator[](size_t i) { return Entry{*this, i}; }
-- 
cgit v1.2.3