[Custom Descriptors] Use placeholder describees in GTO (#7888)

tlively · kripken · web-flow · commit 4f25fee7d364 · 2025-09-10T23:53:02.000Z
GTO tries to optimize out unused descriptor types, but sometimes cannot
because the descriptor type must remain a descriptor to be a valid
supertype of a subtype that will remain a descriptor. To optimize the
original described type to not have a descriptor while simultaneously
keeping the descriptor a valid supertype, insert a new, trivial
placeholder type for the descriptor to describe. Since this new type is
not otherwise used in the module, it should be able to be optimized out
after subsequent rounds of unsubtyping and other optimizations.

---------

Co-authored-by: Alon Zakai &lt;azakai@google.com&gt;
diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
@@ -127,16 +127,15 @@ GlobalTypeRewriter::getSortedTypes(PredecessorGraph preds) {
 
 GlobalTypeRewriter::TypeMap
 GlobalTypeRewriter::rebuildTypes(std::vector<HeapType> types) {
-  Index i = 0;
-  for (auto type : types) {
-    typeIndices[type] = i++;
+  for (Index i = 0; i < types.size(); ++i) {
+    typeIndices[types[i]] = i;
   }
 
   if (typeIndices.size() == 0) {
     return {};
   }
 
-  typeBuilder.grow(typeIndices.size());
+  typeBuilder.grow(types.size());
 
   // All the input types are distinct, so we need to make sure the output
   // types are distinct as well. Further, the new types may have more
@@ -146,14 +145,14 @@ GlobalTypeRewriter::rebuildTypes(std::vector<HeapType> types) {
   typeBuilder.createRecGroup(0, typeBuilder.size());
 
   // Create the temporary heap types.
-  i = 0;
   auto map = [&](HeapType type) -> HeapType {
     if (auto it = typeIndices.find(type); it != typeIndices.end()) {
       return typeBuilder[it->second];
     }
     return type;
   };
-  for (auto [type, _] : typeIndices) {
+  for (Index i = 0; i < types.size(); ++i) {
+    auto type = types[i];
     typeBuilder[i].copy(type, map);
     switch (type.getKind()) {
       case HeapTypeKind::Func: {
@@ -191,7 +190,6 @@ GlobalTypeRewriter::rebuildTypes(std::vector<HeapType> types) {
     }
 
     modifyTypeBuilderEntry(typeBuilder, i, type);
-    ++i;
   }
 
   auto buildResults = typeBuilder.build();
diff --git a/src/passes/GlobalTypeOptimization.cpp b/src/passes/GlobalTypeOptimization.cpp
@@ -30,10 +30,9 @@
 #include "ir/struct-utils.h"
 #include "ir/subtypes.h"
 #include "ir/type-updating.h"
-#include "ir/utils.h"
 #include "pass.h"
+#include "support/insert_ordered.h"
 #include "support/permutations.h"
-#include "wasm-builder.h"
 #include "wasm-type-ordering.h"
 #include "wasm-type.h"
 #include "wasm.h"
@@ -138,6 +137,12 @@ struct GlobalTypeOptimization : public Pass {
   // The types that no longer need a descriptor.
   std::unordered_set<HeapType> haveUnneededDescriptors;
 
+  // Descriptor types that are not needed by their described types but that
+  // still need to be descriptors for their own subtypes and supertypes to be
+  // valid. We will keep them descriptors by having them describe trivial new
+  // placeholder types.
+  InsertOrderedMap<HeapType, Index> descriptorsOfPlaceholders;
+
   void run(Module* module) override {
     if (!module->features.hasGC()) {
       return;
@@ -423,45 +428,43 @@ struct GlobalTypeOptimization : public Pass {
     //        ^
     //   B -> B.desc
     //
-    // Here the descriptors subtype, but *not* the describees. We cannot
-    // remove A's descriptor without also removing $B's, so we need to propagate
-    // that "must remain a descriptor" property among descriptors.
+    // Say we want to optimize A to no longer have a descriptor. Then A.desc
+    // will no longer describe A. But A.desc still needs to be a descriptor for
+    // it to remain a valid supertype of B.desc. To allow the optimization of A
+    // to proceed, we will introduce a placeholder type for A.desc to describe,
+    // keeping it a descriptor type.
     if (!haveUnneededDescriptors.empty()) {
       StructUtils::TypeHierarchyPropagator<StructUtils::CombinableBool>
         descPropagator(subTypes);
 
       // Populate the initial data: Any descriptor we did not see was unneeded,
       // is needed.
       StructUtils::TypeHierarchyPropagator<
-        StructUtils::CombinableBool>::StructMap map;
+        StructUtils::CombinableBool>::StructMap remainingDesciptors;
       for (auto type : subTypes.types) {
         if (auto desc = type.getDescriptorType()) {
           if (!haveUnneededDescriptors.count(type)) {
             // This descriptor type is needed.
-            map[*desc].value = true;
+            remainingDesciptors[*desc].value = true;
           }
         }
       }
 
       // Propagate.
-      descPropagator.propagateToSuperAndSubTypes(map);
-
-      // Remove optimization opportunities that the propagation ruled out.
-      // TODO: We could do better here,
-      //
-      //         A -> A.desc              A    A.desc <- A2
-      //              ^           =>           ^
-      //         B -> B.desc              B -> B.desc
-      //
-      // Starting from the left, we can remove A's descriptor *but keep A.desc
-      // as being a descriptor*, by making it describe a new type A2. That would
-      // keep subtyping working for the descriptors, and later passes could
-      // remove the unused A2.
-      for (auto& [type, info] : map) {
-        if (info.value) {
-          auto described = type.getDescribedType();
-          assert(described);
-          haveUnneededDescriptors.erase(*described);
+      descPropagator.propagateToSuperAndSubTypes(remainingDesciptors);
+
+      // Determine the set of descriptor types that will need placeholder
+      // describees. Do not iterate directly on remainingDescriptors because it
+      // is not deterministically ordered.
+      for (auto type : subTypes.types) {
+        if (auto it = remainingDesciptors.find(type);
+            it != remainingDesciptors.end() && it->second.value) {
+          auto desc = type.getDescribedType();
+          assert(desc);
+          if (haveUnneededDescriptors.count(*desc)) {
+            descriptorsOfPlaceholders.insert(
+              {type, descriptorsOfPlaceholders.size()});
+          }
         }
       }
     }
@@ -484,10 +487,22 @@ struct GlobalTypeOptimization : public Pass {
   void updateTypes(Module& wasm) {
     class TypeRewriter : public GlobalTypeRewriter {
       GlobalTypeOptimization& parent;
+      InsertOrderedMap<HeapType, Index>::iterator placeholderIt;
 
     public:
       TypeRewriter(Module& wasm, GlobalTypeOptimization& parent)
-        : GlobalTypeRewriter(wasm), parent(parent) {}
+        : GlobalTypeRewriter(wasm), parent(parent),
+          placeholderIt(parent.descriptorsOfPlaceholders.begin()) {}
+
+      std::vector<HeapType> getSortedTypes(PredecessorGraph preds) override {
+        auto types = GlobalTypeRewriter::getSortedTypes(std::move(preds));
+        // Prefix the types with placeholders to be overwritten with the
+        // placeholder describees.
+        HeapType placeholder = Struct{};
+        types.insert(
+          types.begin(), parent.descriptorsOfPlaceholders.size(), placeholder);
+        return types;
+      }
 
       void modifyStruct(HeapType oldStructType, Struct& struct_) override {
         auto& newFields = struct_.fields;
@@ -549,17 +564,30 @@ struct GlobalTypeOptimization : public Pass {
           return;
         }
 
-        // Remove an unneeded descriptor.
-        if (parent.haveUnneededDescriptors.count(oldType)) {
-          typeBuilder.setDescriptor(i, std::nullopt);
+        // Until we've created all the placeholders, create a placeholder
+        // describee type for the next descriptor that needs one.
+        if (placeholderIt != parent.descriptorsOfPlaceholders.end()) {
+          typeBuilder[i].descriptor(getTempHeapType(placeholderIt->first));
+          ++placeholderIt;
+          return;
         }
 
-        // Remove an unneeded describes.
+        // Remove an unneeded describee or describe a placeholder type.
         if (auto described = oldType.getDescribedType()) {
           if (parent.haveUnneededDescriptors.count(*described)) {
-            typeBuilder.setDescribed(i, std::nullopt);
+            if (auto it = parent.descriptorsOfPlaceholders.find(oldType);
+                it != parent.descriptorsOfPlaceholders.end()) {
+              typeBuilder[i].describes(typeBuilder[it->second]);
+            } else {
+              typeBuilder[i].describes(std::nullopt);
+            }
           }
         }
+
+        // Remove an unneeded descriptor.
+        if (parent.haveUnneededDescriptors.count(oldType)) {
+          typeBuilder.setDescriptor(i, std::nullopt);
+        }
       }
     };
 
diff --git a/test/lit/passes/gto-desc-tnh.wast b/test/lit/passes/gto-desc-tnh.wast
@@ -95,18 +95,21 @@
 ;;   B -> B.desc
 ;;
 ;; $B is written a null descriptor, so we cannot optimize it when traps are
-;; possible. This also prevents optimizations on $A: we cannot remove that
-;; descriptor either, or its subtype would break.
+;; possible. This means $A.desc must remain a descriptor even as we optimize $A,
+;; so we give $A.desc a placeholder describee. With TNH, we can optimize without
+;; the placeholder.
 ;;
 ;; This tests subtyping of descriptors *without* subtyping of describees.
 (module
   (rec
     ;; CHECK:      (rec
-    ;; CHECK-NEXT:  (type $A (sub (descriptor $A.desc (struct))))
+    ;; CHECK-NEXT:  (type $0 (descriptor $A.desc (struct)))
+
+    ;; CHECK:       (type $A (sub (struct)))
     ;; T_N_H:      (rec
     ;; T_N_H-NEXT:  (type $A (sub (struct)))
     (type $A (sub (descriptor $A.desc (struct))))
-    ;; CHECK:       (type $A.desc (sub (describes $A (struct))))
+    ;; CHECK:       (type $A.desc (sub (describes $0 (struct))))
     ;; T_N_H:       (type $A.desc (sub (struct)))
     (type $A.desc (sub (describes $A (struct ))))
 
@@ -118,9 +121,9 @@
     (type $B.desc (sub $A.desc (describes $B (struct))))
   )
 
-  ;; CHECK:      (type $4 (func))
+  ;; CHECK:       (type $5 (func))
 
-  ;; CHECK:      (func $test (type $4)
+  ;; CHECK:      (func $test (type $5)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (struct.new_default $B
   ;; CHECK-NEXT:    (ref.null none)
@@ -147,30 +150,34 @@
 ;; null descriptor but a use. We cannot optimize even without traps.
 ;; Subtyping of descriptors *without* subtyping of describees.
 ;;
-;; $B's descriptor seems removable, but the subtyping of the descriptors
-;; prevents this.
+;; $A's descriptor can be removed, but $A.desc needs to be given a placeholder
+;; describee.
 (module
   (rec
     ;; CHECK:      (rec
-    ;; CHECK-NEXT:  (type $A (sub (descriptor $A.desc (struct))))
+    ;; CHECK-NEXT:  (type $0 (descriptor $B.desc (struct)))
+
+    ;; CHECK:       (type $A (sub (descriptor $A.desc (struct))))
     ;; T_N_H:      (rec
-    ;; T_N_H-NEXT:  (type $A (sub (descriptor $A.desc (struct))))
+    ;; T_N_H-NEXT:  (type $0 (descriptor $B.desc (struct)))
+
+    ;; T_N_H:       (type $A (sub (descriptor $A.desc (struct))))
     (type $A (sub (descriptor $A.desc (struct))))
     ;; CHECK:       (type $A.desc (sub (describes $A (struct))))
     ;; T_N_H:       (type $A.desc (sub (describes $A (struct))))
     (type $A.desc (sub (describes $A (struct ))))
 
-    ;; CHECK:       (type $B (sub (descriptor $B.desc (struct))))
-    ;; T_N_H:       (type $B (sub (descriptor $B.desc (struct))))
+    ;; CHECK:       (type $B (sub (struct)))
+    ;; T_N_H:       (type $B (sub (struct)))
     (type $B (sub (descriptor $B.desc (struct))))
-    ;; CHECK:       (type $B.desc (sub $A.desc (describes $B (struct))))
-    ;; T_N_H:       (type $B.desc (sub $A.desc (describes $B (struct))))
+    ;; CHECK:       (type $B.desc (sub $A.desc (describes $0 (struct))))
+    ;; T_N_H:       (type $B.desc (sub $A.desc (describes $0 (struct))))
     (type $B.desc (sub $A.desc (describes $B (struct))))
   )
 
-  ;; CHECK:      (type $4 (func))
+  ;; CHECK:       (type $5 (func))
 
-  ;; CHECK:      (func $test (type $4)
+  ;; CHECK:      (func $test (type $5)
   ;; CHECK-NEXT:  (local $B (ref $B))
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (ref.get_desc $A
@@ -180,9 +187,9 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
-  ;; T_N_H:      (type $4 (func))
+  ;; T_N_H:       (type $5 (func))
 
-  ;; T_N_H:      (func $test (type $4)
+  ;; T_N_H:      (func $test (type $5)
   ;; T_N_H-NEXT:  (local $B (ref $B))
   ;; T_N_H-NEXT:  (drop
   ;; T_N_H-NEXT:   (ref.get_desc $A
diff --git a/test/lit/passes/gto-desc.wast b/test/lit/passes/gto-desc.wast
