[Strings] Keep public and private types separate in StringLowering (#6642)

We need StringLowering to modify even public types, as it must replace every
single stringref with externref, even if that modifies the ABI. To achieve that
we told it that all string-using types were private, which let TypeUpdater update
them, but the problem is that it moves all private types to a new single
rec group, which meant public and private types ended up in the same group.
As a result, a single public type would make it all public, preventing optimizations
and breaking things as in #6630 #6640.

Ideally TypeUpdater would modify public types while keeping them in the same
rec groups, but this may be a very specific issue for StringLowering, and that
might be a lot of work. Instead, just make StringLowering handle public types of
functions in a manual way, which is simple and should handle all cases that
matter in practice, at least in J2Wasm.

3 files changed, 141 insertions, 41 deletions

diff --git a/src/passes/StringLowering.cpp b/src/passes/StringLowering.cpp
index 9903f090e..ca0ba773c 100644
--- a/src/passes/StringLowering.cpp
+++ b/src/passes/StringLowering.cpp
@@ -267,6 +267,44 @@ struct StringLowering : public StringGathering {
   Type nnExt = Type(HeapType::ext, NonNullable);
 
   void updateTypes(Module* module) {
+    // TypeMapper will not handle public types, but we do want to modify them as
+    // well: we are modifying the public ABI here. We can't simply tell
+    // TypeMapper to consider them private, as then they'd end up in the new big
+    // rec group with the private types (and as they are public, that would make
+    // the entire rec group public, and all types in the module with it).
+    // Instead, manually handle singleton-rec groups of function types. This
+    // keeps them at size 1, as expected, and handles the cases of function
+    // imports and exports. If we need more (non-function types, non-singleton
+    // rec groups, etc.) then more work will be necessary TODO
+    //
+    // Note that we do this before TypeMapper, which allows it to then fix up
+    // things like the types of parameters (which depend on the type of the
+    // function, which must be modified either in TypeMapper - but as just
+    // explained we cannot do that - or before it, which is what we do here).
+    for (auto& func : module->functions) {
+      if (func->type.getRecGroup().size() != 1 ||
+          !Type(func->type, Nullable).getFeatures().hasStrings()) {
+        continue;
+      }
+
+      // Fix up the stringrefs in this type that uses strings and is in a
+      // singleton rec group.
+      std::vector<Type> params, results;
+      auto fix = [](Type t) {
+        if (t.isRef() && t.getHeapType() == HeapType::string) {
+          t = Type(HeapType::ext, t.getNullability());
+        }
+        return t;
+      };
+      for (auto param : func->type.getSignature().params) {
+        params.push_back(fix(param));
+      }
+      for (auto result : func->type.getSignature().results) {
+        results.push_back(fix(result));
+      }
+      func->type = Signature(params, results);
+    }
+
     TypeMapper::TypeUpdates updates;
 
     // Strings turn into externref.
@@ -288,19 +326,7 @@ struct StringLowering : public StringGathering {
       }
     }
 
-    // We consider all types that use strings as modifiable, which means we
-    // mark them as non-public. That is, we are doing something TypeMapper
-    // normally does not, as we are changing the external interface/ABI of the
-    // module: we are changing that ABI from using strings to externs.
-    auto publicTypes = ModuleUtils::getPublicHeapTypes(*module);
-    std::vector<HeapType> stringUsers;
-    for (auto t : publicTypes) {
-      if (Type(t, Nullable).getFeatures().hasStrings()) {
-        stringUsers.push_back(t);
-      }
-    }
-
-    TypeMapper(*module, updates).map(stringUsers);
+    TypeMapper(*module, updates).map();
   }
 
   // Imported string functions.
diff --git a/test/lit/passes/string-lowering-instructions.wast b/test/lit/passes/string-lowering-instructions.wast
index b98b9af23..459f18170 100644
--- a/test/lit/passes/string-lowering-instructions.wast
+++ b/test/lit/passes/string-lowering-instructions.wast
@@ -8,14 +8,12 @@
 
     ;; CHECK:      (type $1 (func))
 
-    ;; CHECK:      (type $2 (func (param externref externref) (result i32)))
+    ;; CHECK:      (type $2 (func (result externref)))
 
-    ;; CHECK:      (rec
-    ;; CHECK-NEXT:  (type $3 (func (param externref i32 externref)))
-
-    ;; CHECK:       (type $4 (func (result externref)))
+    ;; CHECK:      (type $3 (func (param externref externref) (result i32)))
 
-    ;; CHECK:       (type $5 (func (param externref)))
+    ;; CHECK:      (rec
+    ;; CHECK-NEXT:  (type $4 (func (param externref)))
 
     ;; CHECK:       (type $struct-of-string (struct (field externref) (field i32) (field anyref)))
     (type $struct-of-string (struct (field stringref) (field i32) (field anyref)))
@@ -39,17 +37,19 @@
     (type $array16 (array (mut i16)))
   )
 
-  ;; CHECK:       (type $13 (func (param externref) (result externref)))
+  ;; CHECK:       (type $12 (func (param externref) (result externref)))
+
+  ;; CHECK:       (type $13 (func (param externref) (result i32)))
 
-  ;; CHECK:       (type $14 (func (param externref) (result i32)))
+  ;; CHECK:       (type $14 (func (param externref externref) (result i32)))
 
-  ;; CHECK:       (type $15 (func (param externref externref) (result i32)))
+  ;; CHECK:       (type $15 (func (param externref (ref $0)) (result i32)))
 
-  ;; CHECK:       (type $16 (func (param externref (ref $0)) (result i32)))
+  ;; CHECK:       (type $16 (func (param externref externref) (result (ref extern))))
 
-  ;; CHECK:       (type $17 (func (param externref externref) (result (ref extern))))
+  ;; CHECK:       (type $17 (func (param (ref $0))))
 
-  ;; CHECK:       (type $18 (func (param (ref $0))))
+  ;; CHECK:      (type $18 (func (param externref i32 externref)))
 
   ;; CHECK:      (type $19 (func (param (ref null $0) i32 i32) (result (ref extern))))
 
@@ -81,9 +81,9 @@
 
   ;; CHECK:      (import "wasm:js-string" "intoCharCodeArray" (func $intoCharCodeArray (type $22) (param externref (ref null $0) i32) (result i32)))
 
-  ;; CHECK:      (import "wasm:js-string" "equals" (func $equals (type $2) (param externref externref) (result i32)))
+  ;; CHECK:      (import "wasm:js-string" "equals" (func $equals (type $3) (param externref externref) (result i32)))
 
-  ;; CHECK:      (import "wasm:js-string" "compare" (func $compare (type $2) (param externref externref) (result i32)))
+  ;; CHECK:      (import "wasm:js-string" "compare" (func $compare (type $3) (param externref externref) (result i32)))
 
   ;; CHECK:      (import "wasm:js-string" "length" (func $length (type $23) (param externref) (result i32)))
 
@@ -98,7 +98,7 @@
 
   ;; CHECK:      (export "export.2" (func $exported-string-receiver))
 
-  ;; CHECK:      (func $string.new.gc (type $18) (param $array16 (ref $0))
+  ;; CHECK:      (func $string.new.gc (type $17) (param $array16 (ref $0))
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (call $fromCharCodeArray
   ;; CHECK-NEXT:    (local.get $array16)
@@ -117,7 +117,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.from_code_point (type $4) (result externref)
+  ;; CHECK:      (func $string.from_code_point (type $2) (result externref)
   ;; CHECK-NEXT:  (call $fromCodePoint_18
   ;; CHECK-NEXT:   (i32.const 1)
   ;; CHECK-NEXT:  )
@@ -128,7 +128,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.concat (type $17) (param $0 externref) (param $1 externref) (result (ref extern))
+  ;; CHECK:      (func $string.concat (type $16) (param $0 externref) (param $1 externref) (result (ref extern))
   ;; CHECK-NEXT:  (call $concat
   ;; CHECK-NEXT:   (local.get $0)
   ;; CHECK-NEXT:   (local.get $1)
@@ -141,7 +141,7 @@
    )
   )
 
-  ;; CHECK:      (func $string.encode (type $16) (param $ref externref) (param $array16 (ref $0)) (result i32)
+  ;; CHECK:      (func $string.encode (type $15) (param $ref externref) (param $array16 (ref $0)) (result i32)
   ;; CHECK-NEXT:  (call $intoCharCodeArray
   ;; CHECK-NEXT:   (local.get $ref)
   ;; CHECK-NEXT:   (local.get $array16)
@@ -156,7 +156,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.eq (type $15) (param $a externref) (param $b externref) (result i32)
+  ;; CHECK:      (func $string.eq (type $14) (param $a externref) (param $b externref) (result i32)
   ;; CHECK-NEXT:  (call $equals
   ;; CHECK-NEXT:   (local.get $a)
   ;; CHECK-NEXT:   (local.get $b)
@@ -169,7 +169,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.compare (type $15) (param $a externref) (param $b externref) (result i32)
+  ;; CHECK:      (func $string.compare (type $14) (param $a externref) (param $b externref) (result i32)
   ;; CHECK-NEXT:  (call $compare
   ;; CHECK-NEXT:   (local.get $a)
   ;; CHECK-NEXT:   (local.get $b)
@@ -182,7 +182,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.length (type $14) (param $ref externref) (result i32)
+  ;; CHECK:      (func $string.length (type $13) (param $ref externref) (result i32)
   ;; CHECK-NEXT:  (call $length
   ;; CHECK-NEXT:   (local.get $ref)
   ;; CHECK-NEXT:  )
@@ -193,7 +193,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.get_codeunit (type $14) (param $ref externref) (result i32)
+  ;; CHECK:      (func $string.get_codeunit (type $13) (param $ref externref) (result i32)
   ;; CHECK-NEXT:  (call $charCodeAt
   ;; CHECK-NEXT:   (local.get $ref)
   ;; CHECK-NEXT:   (i32.const 2)
@@ -206,7 +206,7 @@
     )
   )
 
-  ;; CHECK:      (func $string.slice (type $13) (param $ref externref) (result externref)
+  ;; CHECK:      (func $string.slice (type $12) (param $ref externref) (result externref)
   ;; CHECK-NEXT:  (call $substring
   ;; CHECK-NEXT:   (local.get $ref)
   ;; CHECK-NEXT:   (i32.const 2)
@@ -221,7 +221,7 @@
     )
   )
 
-  ;; CHECK:      (func $if.string (type $13) (param $ref externref) (result externref)
+  ;; CHECK:      (func $if.string (type $12) (param $ref externref) (result externref)
   ;; CHECK-NEXT:  (if (result externref)
   ;; CHECK-NEXT:   (i32.const 0)
   ;; CHECK-NEXT:   (then
@@ -244,7 +244,7 @@
     )
   )
 
-  ;; CHECK:      (func $if.string.flip (type $13) (param $ref externref) (result externref)
+  ;; CHECK:      (func $if.string.flip (type $12) (param $ref externref) (result externref)
   ;; CHECK-NEXT:  (if (result externref)
   ;; CHECK-NEXT:   (i32.const 0)
   ;; CHECK-NEXT:   (then
@@ -268,7 +268,7 @@
     )
   )
 
-  ;; CHECK:      (func $exported-string-returner (type $4) (result externref)
+  ;; CHECK:      (func $exported-string-returner (type $2) (result externref)
   ;; CHECK-NEXT:  (global.get $string.const_exported)
   ;; CHECK-NEXT: )
   (func $exported-string-returner (export "export.1") (result stringref)
@@ -277,7 +277,7 @@
     (string.const "exported")
   )
 
-  ;; CHECK:      (func $exported-string-receiver (type $3) (param $x externref) (param $y i32) (param $z externref)
+  ;; CHECK:      (func $exported-string-receiver (type $18) (param $x externref) (param $y i32) (param $z externref)
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (local.get $x)
   ;; CHECK-NEXT:  )
@@ -398,7 +398,7 @@
     )
   )
 
-  ;; CHECK:      (func $call-param-null (type $5) (param $str externref)
+  ;; CHECK:      (func $call-param-null (type $4) (param $str externref)
   ;; CHECK-NEXT:  (call $call-param-null
   ;; CHECK-NEXT:   (ref.null noextern)
   ;; CHECK-NEXT:  )
diff --git a/test/lit/passes/string-lowering_types.wast b/test/lit/passes/string-lowering_types.wast
new file mode 100644
index 000000000..54367e85a
--- /dev/null
+++ b/test/lit/passes/string-lowering_types.wast
@@ -0,0 +1,74 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --string-lowering  -all -S -o - | filecheck %s
+
+;; A private type exists, and a public type is used by imports (one explicitly,
+;; one implicitly). When we lower stringref into externref the import's types
+;; should not be part of a rec group with the private type: public and private
+;; types must remain separate.
+(module
+  (rec
+    ;; CHECK:      (type $0 (func (param externref externref) (result i32)))
+
+    ;; CHECK:      (type $1 (array (mut i16)))
+
+    ;; CHECK:      (type $2 (func (param externref)))
+
+    ;; CHECK:      (type $3 (func (param (ref extern))))
+
+    ;; CHECK:      (type $4 (func))
+
+    ;; CHECK:      (type $private (struct (field externref)))
+    (type $private (struct (field stringref)))
+  )
+  (type $public (func (param stringref)))
+
+  ;; CHECK:      (type $6 (func (param (ref null $1) i32 i32) (result (ref extern))))
+
+  ;; CHECK:      (type $7 (func (param i32) (result (ref extern))))
+
+  ;; CHECK:      (type $8 (func (param externref externref) (result (ref extern))))
+
+  ;; CHECK:      (type $9 (func (param externref (ref null $1) i32) (result i32)))
+
+  ;; CHECK:      (type $10 (func (param externref) (result i32)))
+
+  ;; CHECK:      (type $11 (func (param externref i32) (result i32)))
+
+  ;; CHECK:      (type $12 (func (param externref i32 i32) (result (ref extern))))
+
+  ;; CHECK:      (import "a" "b" (func $import (type $2) (param externref)))
+  (import "a" "b" (func $import (type $public) (param stringref)))
+
+  ;; CHECK:      (import "a" "b" (func $import-implicit (type $3) (param (ref extern))))
+  (import "a" "b" (func $import-implicit (param (ref string))))
+
+  ;; CHECK:      (import "wasm:js-string" "fromCharCodeArray" (func $fromCharCodeArray (type $6) (param (ref null $1) i32 i32) (result (ref extern))))
+
+  ;; CHECK:      (import "wasm:js-string" "fromCodePoint" (func $fromCodePoint (type $7) (param i32) (result (ref extern))))
+
+  ;; CHECK:      (import "wasm:js-string" "concat" (func $concat (type $8) (param externref externref) (result (ref extern))))
+
+  ;; CHECK:      (import "wasm:js-string" "intoCharCodeArray" (func $intoCharCodeArray (type $9) (param externref (ref null $1) i32) (result i32)))
+
+  ;; CHECK:      (import "wasm:js-string" "equals" (func $equals (type $0) (param externref externref) (result i32)))
+
+  ;; CHECK:      (import "wasm:js-string" "compare" (func $compare (type $0) (param externref externref) (result i32)))
+
+  ;; CHECK:      (import "wasm:js-string" "length" (func $length (type $10) (param externref) (result i32)))
+
+  ;; CHECK:      (import "wasm:js-string" "charCodeAt" (func $charCodeAt (type $11) (param externref i32) (result i32)))
+
+  ;; CHECK:      (import "wasm:js-string" "substring" (func $substring (type $12) (param externref i32 i32) (result (ref extern))))
+
+  ;; CHECK:      (export "export" (func $export))
+
+  ;; CHECK:      (func $export (type $4)
+  ;; CHECK-NEXT:  (local $0 (ref $private))
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $export (export "export")
+    ;; Keep the private type alive.
+    (local (ref $private))
+  )
+)