Fix scratch local optimizations when emitting string slice (#6649)

The binary writing of `stringview_wtf16.slice` requires scratch locals to store
the `start` and `end` operands while the string operand is converted to a
stringview. To avoid unbounded binary bloat when round-tripping, we detect the
case that `start` and `end` are already `local.get`s and avoid using scratch
locals by deferring the binary writing of the `local.get` operands until after
the stringview conversoins is emitted.

We previously optimized the scratch locals for `start` and `end` independently,
but this could produce incorrect code in the case where the `local.get` for
`start` is deferred but its value is changed by a `local.set` in the code for
`end`. Fix the problem by only optimizing to avoid scratch locals in the case
where both `start` and `end` are already `local.get`s, so they will still be
emitted in the original relative order and they cannot interfere with each other
anyway.

3 files changed, 105 insertions, 53 deletions

diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index 50c13da65..ac117e2e0 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2455,32 +2455,25 @@ void BinaryInstWriter::visitStringWTF16Get(StringWTF16Get* curr) {
 void BinaryInstWriter::visitStringSliceWTF(StringSliceWTF* curr) {
   // We need to convert the ref operand to a stringview, but it is buried under
   // the start and end operands. Put the i32s in scratch locals, emit the
-  // conversion, then get the i32s back onto the stack. If either `start` or
-  // `end` is already a local.get, then we can skip its scratch local.
-  bool startDeferred = false, endDeferred = false;
+  // conversion, then get the i32s back onto the stack. If both `start` and
+  // `end` are already local.gets, then we can skip the scratch locals.
+  bool deferred = false;
   Index startIndex, endIndex;
-  if (auto* get = curr->start->dynCast<LocalGet>()) {
-    assert(deferredGets.count(get));
-    startDeferred = true;
-    startIndex = mappedLocals[{get->index, 0}];
+  auto* startGet = curr->start->dynCast<LocalGet>();
+  auto* endGet = curr->end->dynCast<LocalGet>();
+  if (startGet && endGet) {
+    assert(deferredGets.count(startGet));
+    assert(deferredGets.count(endGet));
+    deferred = true;
+    startIndex = mappedLocals[{startGet->index, 0}];
+    endIndex = mappedLocals[{endGet->index, 0}];
   } else {
     startIndex = scratchLocals[Type::i32];
-  }
-  if (auto* get = curr->end->dynCast<LocalGet>()) {
-    assert(deferredGets.count(get));
-    endDeferred = true;
-    endIndex = mappedLocals[{get->index, 0}];
-  } else {
-    endIndex = scratchLocals[Type::i32];
-    if (!startDeferred) {
-      ++endIndex;
-    }
+    endIndex = startIndex + 1;
   }
 
-  if (!endDeferred) {
+  if (!deferred) {
     o << int8_t(BinaryConsts::LocalSet) << U32LEB(endIndex);
-  }
-  if (!startDeferred) {
     o << int8_t(BinaryConsts::LocalSet) << U32LEB(startIndex);
   }
   o << int8_t(BinaryConsts::GCPrefix) << U32LEB(BinaryConsts::StringAsWTF16);
@@ -2698,21 +2691,19 @@ InsertOrderedMap<Type, Index> BinaryInstWriter::countScratchLocals() {
       if (curr->type == Type::unreachable) {
         return;
       }
-      // If `start` or `end` are already local.gets, we can defer emitting those
-      // gets instead of using scratch locals.
-      Index numScratches = 2;
-      if (auto* get = curr->start->dynCast<LocalGet>()) {
-        parent.deferredGets.insert(get);
-        --numScratches;
-      }
-      if (auto* get = curr->end->dynCast<LocalGet>()) {
-        parent.deferredGets.insert(get);
-        --numScratches;
+      // If `start` and `end` are already local.gets, we can defer emitting
+      // those gets instead of using scratch locals.
+      auto* startGet = curr->start->dynCast<LocalGet>();
+      auto* endGet = curr->end->dynCast<LocalGet>();
+      if (startGet && endGet) {
+        parent.deferredGets.insert(startGet);
+        parent.deferredGets.insert(endGet);
+        return;
       }
       // Scratch locals to hold the `start` and `end` values while we emit a
       // stringview conversion for the `ref` value.
       auto& count = scratches[Type::i32];
-      count = std::max(count, numScratches);
+      count = std::max(count, 2u);
     }
 
     // As mentioned in BinaryInstWriter::visitBreak, the type of br_if with a
diff --git a/test/lit/exec/strings.wast b/test/lit/exec/strings.wast
index ca77f9882..f0c659d19 100644
--- a/test/lit/exec/strings.wast
+++ b/test/lit/exec/strings.wast
@@ -284,6 +284,21 @@
     )
   )
 
+  ;; CHECK:      [fuzz-exec] calling slice-ordering
+  ;; CHECK-NEXT: [fuzz-exec] note result: slice-ordering => string("h")
+  (func $slice-ordering (export "slice-ordering") (result (ref string))
+    (local $0 i32)
+    (stringview_wtf16.slice
+      (string.const "hello")
+      ;; If we were to defer emitting this get in the binary writer, it would
+      ;; end up with the wrong value.
+      (local.get $0)
+      (local.tee $0
+        (i32.const 1)
+      )
+    )
+  )
+
   ;; CHECK:      [fuzz-exec] calling new_empty
   ;; CHECK-NEXT: [fuzz-exec] note result: new_empty => string("")
   (func $new_empty (export "new_empty") (result stringref)
@@ -536,6 +551,9 @@
 ;; CHECK:      [fuzz-exec] calling slice-big
 ;; CHECK-NEXT: [fuzz-exec] note result: slice-big => string("defgh")
 
+;; CHECK:      [fuzz-exec] calling slice-ordering
+;; CHECK-NEXT: [fuzz-exec] note result: slice-ordering => string("h")
+
 ;; CHECK:      [fuzz-exec] calling new_empty
 ;; CHECK-NEXT: [fuzz-exec] note result: new_empty => string("")
 
@@ -620,6 +638,7 @@
 ;; CHECK-NEXT: [fuzz-exec] comparing new_wtf16_array
 ;; CHECK-NEXT: [fuzz-exec] comparing slice
 ;; CHECK-NEXT: [fuzz-exec] comparing slice-big
+;; CHECK-NEXT: [fuzz-exec] comparing slice-ordering
 ;; CHECK-NEXT: [fuzz-exec] comparing slice-unicode
 ;; CHECK-NEXT: [fuzz-exec] comparing string.from_code_point
 ;; CHECK-NEXT: [fuzz-exec] comparing string.measure
diff --git a/test/lit/string.as_wtf16.wast b/test/lit/string.as_wtf16.wast
index 97a5d9055..916e11d39 100644
--- a/test/lit/string.as_wtf16.wast
+++ b/test/lit/string.as_wtf16.wast
@@ -185,38 +185,59 @@
   ;; RTRIP:      (func $slice-start-get (type $0) (result stringref)
   ;; RTRIP-NEXT:  (local $start i32)
   ;; RTRIP-NEXT:  (local $1 i32)
-  ;; RTRIP-NEXT:  (local $2 (ref string))
+  ;; RTRIP-NEXT:  (local $2 i32)
+  ;; RTRIP-NEXT:  (local $3 i32)
+  ;; RTRIP-NEXT:  (local $4 (ref string))
   ;; RTRIP-NEXT:  (stringview_wtf16.slice
   ;; RTRIP-NEXT:   (block (result (ref string))
-  ;; RTRIP-NEXT:    (local.set $2
+  ;; RTRIP-NEXT:    (local.set $4
   ;; RTRIP-NEXT:     (string.const "abc")
   ;; RTRIP-NEXT:    )
   ;; RTRIP-NEXT:    (local.set $1
-  ;; RTRIP-NEXT:     (i32.const 2)
+  ;; RTRIP-NEXT:     (block (result i32)
+  ;; RTRIP-NEXT:      (local.set $3
+  ;; RTRIP-NEXT:       (local.get $start)
+  ;; RTRIP-NEXT:      )
+  ;; RTRIP-NEXT:      (local.set $2
+  ;; RTRIP-NEXT:       (i32.const 2)
+  ;; RTRIP-NEXT:      )
+  ;; RTRIP-NEXT:      (local.get $3)
+  ;; RTRIP-NEXT:     )
   ;; RTRIP-NEXT:    )
-  ;; RTRIP-NEXT:    (local.get $2)
+  ;; RTRIP-NEXT:    (local.get $4)
   ;; RTRIP-NEXT:   )
-  ;; RTRIP-NEXT:   (local.get $start)
   ;; RTRIP-NEXT:   (local.get $1)
+  ;; RTRIP-NEXT:   (local.get $2)
   ;; RTRIP-NEXT:  )
   ;; RTRIP-NEXT: )
   ;; RRTRP:      (func $slice-start-get (type $0) (result stringref)
   ;; RRTRP-NEXT:  (local $start i32)
   ;; RRTRP-NEXT:  (local $1 i32)
-  ;; RRTRP-NEXT:  (local $2 (ref string))
-  ;; RRTRP-NEXT:  (local $3 (ref string))
+  ;; RRTRP-NEXT:  (local $2 i32)
+  ;; RRTRP-NEXT:  (local $3 i32)
+  ;; RRTRP-NEXT:  (local $4 (ref string))
+  ;; RRTRP-NEXT:  (local $5 i32)
+  ;; RRTRP-NEXT:  (local $6 (ref string))
   ;; RRTRP-NEXT:  (stringview_wtf16.slice
   ;; RRTRP-NEXT:   (block (result (ref string))
-  ;; RRTRP-NEXT:    (local.set $3
+  ;; RRTRP-NEXT:    (local.set $6
   ;; RRTRP-NEXT:     (string.const "abc")
   ;; RRTRP-NEXT:    )
   ;; RRTRP-NEXT:    (local.set $1
-  ;; RRTRP-NEXT:     (i32.const 2)
+  ;; RRTRP-NEXT:     (block (result i32)
+  ;; RRTRP-NEXT:      (local.set $5
+  ;; RRTRP-NEXT:       (local.get $start)
+  ;; RRTRP-NEXT:      )
+  ;; RRTRP-NEXT:      (local.set $2
+  ;; RRTRP-NEXT:       (i32.const 2)
+  ;; RRTRP-NEXT:      )
+  ;; RRTRP-NEXT:      (local.get $5)
+  ;; RRTRP-NEXT:     )
   ;; RRTRP-NEXT:    )
-  ;; RRTRP-NEXT:    (local.get $3)
+  ;; RRTRP-NEXT:    (local.get $6)
   ;; RRTRP-NEXT:   )
-  ;; RRTRP-NEXT:   (local.get $start)
   ;; RRTRP-NEXT:   (local.get $1)
+  ;; RRTRP-NEXT:   (local.get $2)
   ;; RRTRP-NEXT:  )
   ;; RRTRP-NEXT: )
   (func $slice-start-get (result stringref)
@@ -240,38 +261,59 @@
   ;; RTRIP:      (func $slice-end-get (type $0) (result stringref)
   ;; RTRIP-NEXT:  (local $end i32)
   ;; RTRIP-NEXT:  (local $1 i32)
-  ;; RTRIP-NEXT:  (local $2 (ref string))
+  ;; RTRIP-NEXT:  (local $2 i32)
+  ;; RTRIP-NEXT:  (local $3 i32)
+  ;; RTRIP-NEXT:  (local $4 (ref string))
   ;; RTRIP-NEXT:  (stringview_wtf16.slice
   ;; RTRIP-NEXT:   (block (result (ref string))
-  ;; RTRIP-NEXT:    (local.set $2
+  ;; RTRIP-NEXT:    (local.set $4
   ;; RTRIP-NEXT:     (string.const "abc")
   ;; RTRIP-NEXT:    )
   ;; RTRIP-NEXT:    (local.set $1
-  ;; RTRIP-NEXT:     (i32.const 1)
+  ;; RTRIP-NEXT:     (block (result i32)
+  ;; RTRIP-NEXT:      (local.set $3
+  ;; RTRIP-NEXT:       (i32.const 1)
+  ;; RTRIP-NEXT:      )
+  ;; RTRIP-NEXT:      (local.set $2
+  ;; RTRIP-NEXT:       (local.get $end)
+  ;; RTRIP-NEXT:      )
+  ;; RTRIP-NEXT:      (local.get $3)
+  ;; RTRIP-NEXT:     )
   ;; RTRIP-NEXT:    )
-  ;; RTRIP-NEXT:    (local.get $2)
+  ;; RTRIP-NEXT:    (local.get $4)
   ;; RTRIP-NEXT:   )
   ;; RTRIP-NEXT:   (local.get $1)
-  ;; RTRIP-NEXT:   (local.get $end)
+  ;; RTRIP-NEXT:   (local.get $2)
   ;; RTRIP-NEXT:  )
   ;; RTRIP-NEXT: )
   ;; RRTRP:      (func $slice-end-get (type $0) (result stringref)
   ;; RRTRP-NEXT:  (local $end i32)
   ;; RRTRP-NEXT:  (local $1 i32)
-  ;; RRTRP-NEXT:  (local $2 (ref string))
-  ;; RRTRP-NEXT:  (local $3 (ref string))
+  ;; RRTRP-NEXT:  (local $2 i32)
+  ;; RRTRP-NEXT:  (local $3 i32)
+  ;; RRTRP-NEXT:  (local $4 (ref string))
+  ;; RRTRP-NEXT:  (local $5 i32)
+  ;; RRTRP-NEXT:  (local $6 (ref string))
   ;; RRTRP-NEXT:  (stringview_wtf16.slice
   ;; RRTRP-NEXT:   (block (result (ref string))
-  ;; RRTRP-NEXT:    (local.set $3
+  ;; RRTRP-NEXT:    (local.set $6
   ;; RRTRP-NEXT:     (string.const "abc")
   ;; RRTRP-NEXT:    )
   ;; RRTRP-NEXT:    (local.set $1
-  ;; RRTRP-NEXT:     (i32.const 1)
+  ;; RRTRP-NEXT:     (block (result i32)
+  ;; RRTRP-NEXT:      (local.set $5
+  ;; RRTRP-NEXT:       (i32.const 1)
+  ;; RRTRP-NEXT:      )
+  ;; RRTRP-NEXT:      (local.set $2
+  ;; RRTRP-NEXT:       (local.get $end)
+  ;; RRTRP-NEXT:      )
+  ;; RRTRP-NEXT:      (local.get $5)
+  ;; RRTRP-NEXT:     )
   ;; RRTRP-NEXT:    )
-  ;; RRTRP-NEXT:    (local.get $3)
+  ;; RRTRP-NEXT:    (local.get $6)
   ;; RRTRP-NEXT:   )
   ;; RRTRP-NEXT:   (local.get $1)
-  ;; RRTRP-NEXT:   (local.get $end)
+  ;; RRTRP-NEXT:   (local.get $2)
   ;; RRTRP-NEXT:  )
   ;; RRTRP-NEXT: )
   (func $slice-end-get (result stringref)