Skip to content

Commit

Permalink
[Strings] Keep public and private types separate in StringLowering (#…
Browse files Browse the repository at this point in the history
…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.
  • Loading branch information
kripken committed Jun 10, 2024
1 parent 0a1a59a commit 3d0e687
Show file tree
Hide file tree
Showing 3 changed files with 141 additions and 41 deletions.
52 changes: 39 additions & 13 deletions src/passes/StringLowering.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down
56 changes: 28 additions & 28 deletions test/lit/passes/string-lowering-instructions.wast
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
Expand All @@ -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))))

Expand Down Expand Up @@ -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)))

Expand All @@ -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)
Expand All @@ -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: )
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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: )
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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: )
Expand Down Expand Up @@ -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: )
Expand Down
74 changes: 74 additions & 0 deletions test/lit/passes/string-lowering_types.wast
Original file line number Diff line number Diff line change
@@ -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))
)
)

0 comments on commit 3d0e687

Please sign in to comment.