From e9b012ff6e0e55d73d57b9a9bc4f64c15521bde1 Mon Sep 17 00:00:00 2001
From: Thomas Lively <tlively@google.com>
Date: Wed, 13 Dec 2023 14:05:19 -0800
Subject: Preserve multivalue drops in IRBuilder (#6150)

In Binaryen IR, we allow single `Drop` expressions to drop multiple values
packaged up as a tuple. When using IRBuilder to rebuild IR containing such a
drop, it previously treated the drop as a normal WebAssembly drop that dropped
only a single value, producing invalid IR that had extra, undropped values. Fix
the problem by preserving the arity of `Drop` inputs in IRBuilder. To avoid
bloating the IR, thread the size of the desired value through IRBuilder's pop
implementation so that tuple values do not need to be split up and recombined.
---
 src/wasm-ir-builder.h        | 12 +++++++++---
 src/wasm/wasm-ir-builder.cpp | 44 +++++++++++++++++++++++++++++++++++++-------
 2 files changed, 46 insertions(+), 10 deletions(-)

(limited to 'src')

diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index 8b01977be..e1d2ce1fd 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -200,6 +200,8 @@ public:
 
   // Private functions that must be public for technical reasons.
   [[nodiscard]] Result<> visitExpression(Expression*);
+  [[nodiscard]] Result<>
+  visitDrop(Drop*, std::optional<uint32_t> arity = std::nullopt);
   [[nodiscard]] Result<> visitIf(If*);
   [[nodiscard]] Result<> visitReturn(Return*);
   [[nodiscard]] Result<> visitStructNew(StructNew*);
@@ -463,7 +465,7 @@ private:
   [[nodiscard]] Result<Name> getLabelName(Index label);
   [[nodiscard]] Result<Name> getDelegateLabelName(Index label);
   [[nodiscard]] Result<Index> addScratchLocal(Type);
-  [[nodiscard]] Result<Expression*> pop();
+  [[nodiscard]] Result<Expression*> pop(size_t size = 1);
 
   struct HoistedVal {
     // The index in the stack of the original value-producing expression.
@@ -478,8 +480,12 @@ private:
   // Transform the stack as necessary such that the original producer of the
   // hoisted value will be popped along with the final expression that produces
   // the value, if they are different. May only be called directly after
-  // hoistLastValue().
-  [[nodiscard]] Result<> packageHoistedValue(const HoistedVal&);
+  // hoistLastValue(). `sizeHint` is the size of the type we ultimately want to
+  // consume, so if the hoisted value has `sizeHint` elements, it is left intact
+  // even if it is a tuple. Otherwise, hoisted tuple values will be broken into
+  // pieces.
+  [[nodiscard]] Result<> packageHoistedValue(const HoistedVal&,
+                                             size_t sizeHint = 1);
 
   [[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
                                                    std::optional<Index> label);
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index d3f48b1c0..fbb116aa9 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -90,7 +90,8 @@ MaybeResult<IRBuilder::HoistedVal> IRBuilder::hoistLastValue() {
   return HoistedVal{Index(index), get};
 }
 
-Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted) {
+Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted,
+                                        size_t sizeHint) {
   auto& scope = getScope();
   assert(!scope.exprStack.empty());
 
@@ -106,7 +107,7 @@ Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted) {
 
   auto type = scope.exprStack.back()->type;
 
-  if (!type.isTuple()) {
+  if (type.size() == sizeHint) {
     if (hoisted.get) {
       packageAsBlock(type);
     }
@@ -158,7 +159,8 @@ void IRBuilder::push(Expression* expr) {
   DBG(dump());
 }
 
-Result<Expression*> IRBuilder::pop() {
+Result<Expression*> IRBuilder::pop(size_t size) {
+  assert(size >= 1);
   auto& scope = getScope();
 
   // Find the suffix of expressions that do not produce values.
@@ -173,11 +175,25 @@ Result<Expression*> IRBuilder::pop() {
     return Err{"popping from empty stack"};
   }
 
-  CHECK_ERR(packageHoistedValue(*hoisted));
+  CHECK_ERR(packageHoistedValue(*hoisted, size));
 
   auto* ret = scope.exprStack.back();
-  scope.exprStack.pop_back();
-  return ret;
+  if (ret->type.size() == size) {
+    scope.exprStack.pop_back();
+    return ret;
+  }
+
+  // The last value-producing expression did not produce exactly the right
+  // number of values, so we need to construct a tuple piecewise instead.
+  assert(size > 1);
+  std::vector<Expression*> elems;
+  elems.resize(size);
+  for (int i = size - 1; i >= 0; --i) {
+    auto elem = pop();
+    CHECK_ERR(elem);
+    elems[i] = *elem;
+  }
+  return builder.makeTupleMake(elems);
 }
 
 Result<Expression*> IRBuilder::build() {
@@ -319,6 +335,20 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitDrop(Drop* curr, std::optional<uint32_t> arity) {
+  // Multivalue drops must remain multivalue drops.
+  if (!arity) {
+    arity = curr->value->type.size();
+  }
+  if (*arity >= 2) {
+    auto val = pop(*arity);
+    CHECK_ERR(val);
+    curr->value = *val;
+    return Ok{};
+  }
+  return visitExpression(curr);
+}
+
 Result<> IRBuilder::visitIf(If* curr) {
   // Only the condition is popped from the stack. The ifTrue and ifFalse are
   // self-contained so we do not modify them.
@@ -1183,7 +1213,7 @@ Result<> IRBuilder::makeSelect(std::optional<Type> type) {
 
 Result<> IRBuilder::makeDrop() {
   Drop curr;
-  CHECK_ERR(visitDrop(&curr));
+  CHECK_ERR(visitDrop(&curr, 1));
   push(builder.makeDrop(curr.value));
   return Ok{};
 }
-- 
cgit v1.2.3