SignaturePruning: Properly handle public types (#6630)

The SignaturePruning pass optimizes away parameters that it proves are safe to
remove. It turns out that that does not always match the definition of private
types, which is more restrictive. Specifically, if say all the types are in one big
rec group and one of them is used on an exported function then all of them are
considered public (as the rec group is). However, in closed world, it would be ok
to leave that rec group unchanged but to create a pruned version of that type
and use it, in cases where we see it is safe to remove a parameter. (See the
testcase for a concrete example.)

To put it another way, SignaturePruning already proves that a parameter is
safe to remove in all the ways that matter. Before this PR, however, the testcase
in this PR would error - so this PR is not an optimization but a bugfix, really -
because SignaturePruning would see that a parameter is safe to remove but
then TypeUpdating would see the type is public and so it would leave it alone,
leading to a broken module.

This situation is in fact not that rare, and happens on real-world Java code.
The reason we did not notice it before is that typically there are no remaining
SignaturePruning opportunities late in the process (when other closed world
optimizations have typically led to a single big rec group).

The concrete fix here is to add additionalPrivateTypes to a few more places
in TypeUpdating. We already supported that for cases where a pass knew
better than the general logic what can be modified, and this adds that
ability to the signature-rewriting logic there. Then SignaturePruning can
send in all the types it has proven are safe to modify.

 *  Also necessary here is to only add from additionalPrivateTypes if the type
    is not already in our list (or we'd end up with duplicates in the final rec
    group).
*   Also move newSignatures in SignaturePruning out of the top level, which
    was confusing (the pass has multiple iterations, and we want each to have
    a fresh instance).

4 files changed, 95 insertions, 18 deletions

diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index 12a8c7c36..760d71b9c 100644
--- a/src/ir/type-updating.cpp
+++ b/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
diff --git a/src/ir/type-updating.h b/src/ir/type-updating.h
index 1626b40bd..60b92e585 100644
--- a/src/ir/type-updating.h
+++ b/src/ir/type-updating.h
@@ -352,7 +352,12 @@ public:
   // 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 @@ public:
 
   // 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 @@ public:
       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,7 +429,7 @@ public:
           sig.results = getTempType(iter->second.results);
         }
       }
-    } rewriter(wasm, updates);
+    } rewriter(wasm, updates, additionalPrivateTypes);
   }
 
 protected:
@@ -427,9 +437,7 @@ protected:
   // 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 = {});
 
diff --git a/src/passes/SignaturePruning.cpp b/src/passes/SignaturePruning.cpp
index 2e4be89e8..4480c1832 100644
--- a/src/passes/SignaturePruning.cpp
+++ b/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 signature. 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;
diff --git a/test/lit/passes/signature-pruning.wast b/test/lit/passes/signature-pruning.wast
index cad9af82b..c273d1588 100644
--- a/test/lit/passes/signature-pruning.wast
+++ b/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)
+    )
+  )
+)
+