SignaturePruning: Properly handle public types

src/ir/type-updating.cpp

@@ -28,7 +28,10 @@ namespace wasm {
 
 GlobalTypeRewriter::GlobalTypeRewriter(Module& wasm) : wasm(wasm) {}
 
-void GlobalTypeRewriter::update() { mapTypes(rebuildTypes()); }
+void GlobalTypeRewriter::update(
+  const std::vector<HeapType>& additionalPrivateTypes) {
+  mapTypes(rebuildTypes(additionalPrivateTypes));
+}
 
 GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
   const std::vector<HeapType>& additionalPrivateTypes) {
@@ -40,8 +43,17 @@ GlobalTypeRewriter::TypeMap GlobalTypeRewriter::rebuildTypes(
   Index i = 0;
   auto privateTypes = ModuleUtils::getPrivateHeapTypes(wasm);
 
-  for (auto t : additionalPrivateTypes) {
-    privateTypes.push_back(t);
+  if (!additionalPrivateTypes.empty()) {
+    // Only add additional private types that are not already in the list.
+    std::unordered_set<HeapType> privateTypesSet(privateTypes.begin(),
+                                                 privateTypes.end());
+
+    for (auto t : additionalPrivateTypes) {
+      if (!privateTypesSet.count(t)) {
+        privateTypes.push_back(t);
+        privateTypesSet.insert(t);
+      }
+    }
   }
 
   // Topological sort to have supertypes first, but we have to account for the

src/ir/type-updating.h

@@ -352,7 +352,12 @@ class GlobalTypeRewriter {
   // Main entry point. This performs the entire process of creating new heap
   // types and calling the hooks below, then applies the new types throughout
   // the module.
-  void update();
+  //
+  // This only operates on private types (so as not to modify the module's
+  // external ABI). It takes as a parameter a list of public types to consider
+  // private, which allows more flexibility (e.g. in closed world if a pass
+  // knows a type is safe to modify despite being public, it can add it).
+  void update(const std::vector<HeapType>& additionalPrivateTypes = {});
 
   using TypeMap = std::unordered_map<HeapType, HeapType>;
 
@@ -398,7 +403,10 @@ class GlobalTypeRewriter {
 
   // Helper for the repeating pattern of just updating Signature types using a
   // map of old heap type => new Signature.
-  static void updateSignatures(const SignatureUpdates& updates, Module& wasm) {
+  static void
+  updateSignatures(const SignatureUpdates& updates,
+                   Module& wasm,
+                   const std::vector<HeapType>& additionalPrivateTypes = {}) {
     if (updates.empty()) {
       return;
     }
@@ -407,9 +415,11 @@ class GlobalTypeRewriter {
       const SignatureUpdates& updates;
 
     public:
-      SignatureRewriter(Module& wasm, const SignatureUpdates& updates)
+      SignatureRewriter(Module& wasm,
+                        const SignatureUpdates& updates,
+                        const std::vector<HeapType>& additionalPrivateTypes)
         : GlobalTypeRewriter(wasm), updates(updates) {
-        update();
+        update(additionalPrivateTypes);
       }
 
       void modifySignature(HeapType oldSignatureType, Signature& sig) override {
@@ -419,17 +429,15 @@ class GlobalTypeRewriter {
           sig.results = getTempType(iter->second.results);
         }
       }
-    } rewriter(wasm, updates);
+    } rewriter(wasm, updates, additionalPrivateTypes);
   }
 
 protected:
   // Builds new types after updating their contents using the hooks below and
   // returns a map from the old types to the modified types. Used internally in
   // update().
   //
-  // This only operates on private types (so as not to modify the module's
-  // external ABI). It takes as a parameter a list of public types to consider
-  // private, which allows more flexibility.
+  // See above regarding private types.
   TypeMap
   rebuildTypes(const std::vector<HeapType>& additionalPrivateTypes = {});
 

src/passes/SignaturePruning.cpp

@@ -45,11 +45,6 @@ namespace wasm {
 namespace {
 
 struct SignaturePruning : public Pass {
-  // Maps each heap type to the possible pruned heap type. We will fill this
-  // during analysis and then use it while doing an update of the types. If a
-  // type has no improvement that we can find, it will not appear in this map.
-  std::unordered_map<HeapType, Signature> newSignatures;
-
   void run(Module* module) override {
     if (!module->features.hasGC()) {
       return;
@@ -182,6 +177,11 @@ struct SignaturePruning : public Pass {
     //      types with subtyping relations at once.
     SubTypes subTypes(*module);
 
+    // Maps each heap type to the possible pruned heap type. We will fill this
+    // during analysis and then use it while doing an update of the types. If a
+    // type has no improvement that we can find, it will not appear in this map.
+    std::unordered_map<HeapType, Signature> newSignatures;
+
     // Find parameters to prune.
     //
     // TODO: The order matters here, and more than one cycle can find more work
@@ -291,8 +291,16 @@ struct SignaturePruning : public Pass {
       }
     }
 
-    // Rewrite the types.
-    GlobalTypeRewriter::updateSignatures(newSignatures, *module);
+    // Rewrite the types. We pass in all the types we intend to modify as being
+    // "additional private types" because we have proven above that they are
+    // safe to modify, even if they are technically public (e.g. they may be in
+    // a singleton big rec group that is public because one member is public).
+    std::vector<HeapType> additionalPrivateTypes;
+    for (auto& [type, sig] : newSignatures) {
+      additionalPrivateTypes.push_back(type);
+    }
+    GlobalTypeRewriter::updateSignatures(
+      newSignatures, *module, additionalPrivateTypes);
 
     if (callTargetsToLocalize.empty()) {
       return false;

test/lit/passes/signature-pruning.wast

@@ -1151,3 +1151,52 @@
     )
   )
 )
+
+;; $exported is exported. The entire rec group becomes exported as well, which
+;; causes $unused-param's type to be public, which means we cannot normally
+;; modify it. However, in closed world we allow such changes, and we can remove
+;; the unused param there. What happens is that we keep the original public rec
+;; group as-is, and add a new rec group for private types, put the pruned type
+;; there, and use that pruned type on $unused-param.
+(module
+  (rec
+   ;; CHECK:      (rec
+   ;; CHECK-NEXT:  (type $much (func))
+
+   ;; CHECK:       (type $1 (func))
+
+   ;; CHECK:      (rec
+   ;; CHECK-NEXT:  (type $none (func))
+   (type $none (func))
+   ;; CHECK:       (type $much (func (param i32)))
+   (type $much (func (param i32)))
+  )
+
+  ;; CHECK:      (export "exported" (func $exported))
+
+  ;; CHECK:      (func $exported (type $none)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $exported (export "exported") (type $none)
+  )
+
+  ;; CHECK:      (func $unused-param (type $much)
+  ;; CHECK-NEXT:  (local $0 i32)
+  ;; CHECK-NEXT:  (local.set $0
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $unused-param (type $much) (param $param i32)
+  )
+
+  ;; CHECK:      (func $caller (type $1)
+  ;; CHECK-NEXT:  (call $unused-param)
+  ;; CHECK-NEXT: )
+  (func $caller
+    (call $unused-param
+      (i32.const 0)
+    )
+  )
+)
+