[Parser] Parse tuple operations (#6174)

Parse `tuple.make`, `tuple.extract`, and `tuple.drop`. Also slightly improve the way we break up tuples into individual elements in IRBuilder by using a `local.tee` instead of a block containing a `local.set` and `local.get`.
WebAssembly · Dec 13, 2023 · 7adc82b · 7adc82b
1 parent e9b012f
commit 7adc82b
Show file tree

Hide file tree

Showing 7 changed files with 194 additions and 137 deletions.
diff --git a/src/parser/contexts.h b/src/parser/contexts.h
@@ -430,6 +430,9 @@ struct NullInstrParserCtx {
   Result<> makeTableCopy(Index, TableIdxT*, TableIdxT*) { return Ok{}; }
   Result<> makeThrow(Index, TagIdxT) { return Ok{}; }
   Result<> makeRethrow(Index, LabelIdxT) { return Ok{}; }
+  Result<> makeTupleMake(Index, uint32_t) { return Ok{}; }
+  Result<> makeTupleExtract(Index, uint32_t, uint32_t) { return Ok{}; }
+  Result<> makeTupleDrop(Index, uint32_t) { return Ok{}; }
   template<typename HeapTypeT> Result<> makeCallRef(Index, HeapTypeT, bool) {
     return Ok{};
   }
@@ -1624,6 +1627,18 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return withLoc(pos, irBuilder.makeRethrow(label));
   }
 
+  Result<> makeTupleMake(Index pos, uint32_t arity) {
+    return withLoc(pos, irBuilder.makeTupleMake(arity));
+  }
+
+  Result<> makeTupleExtract(Index pos, uint32_t arity, uint32_t index) {
+    return withLoc(pos, irBuilder.makeTupleExtract(arity, index));
+  }
+
+  Result<> makeTupleDrop(Index pos, uint32_t arity) {
+    return withLoc(pos, irBuilder.makeTupleDrop(arity));
+  }
+
   Result<> makeCallRef(Index pos, HeapType type, bool isReturn) {
     return withLoc(pos, irBuilder.makeCallRef(type, isReturn));
   }

diff --git a/src/parser/parsers.h b/src/parser/parsers.h
@@ -41,6 +41,7 @@ template<typename Ctx> Result<typename Ctx::LimitsT> limits64(Ctx&);
 template<typename Ctx> Result<typename Ctx::MemTypeT> memtype(Ctx&);
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx&);
 template<typename Ctx> Result<typename Ctx::GlobalTypeT> globaltype(Ctx&);
+template<typename Ctx> Result<uint32_t> tupleArity(Ctx&);
 
 // Instructions
 template<typename Ctx> MaybeResult<> foldedBlockinstr(Ctx&);
@@ -605,6 +606,18 @@ template<typename Ctx> Result<typename Ctx::GlobalTypeT> globaltype(Ctx& ctx) {
   return ctx.makeGlobalType(mutability, *type);
 }
 
+// arity ::= x:u32    (if x >=2 )
+template<typename Ctx> Result<uint32_t> tupleArity(Ctx& ctx) {
+  auto arity = ctx.in.takeU32();
+  if (!arity) {
+    return ctx.in.err("expected tuple arity");
+  }
+  if (*arity < 2) {
+    return ctx.in.err("tuple arity must be at least 2");
+  }
+  return *arity;
+}
+
 // ============
 // Instructions
 // ============
@@ -1512,15 +1525,25 @@ template<typename Ctx> Result<> makeRethrow(Ctx& ctx, Index pos) {
 }
 
 template<typename Ctx> Result<> makeTupleMake(Ctx& ctx, Index pos) {
-  return ctx.in.err("unimplemented instruction");
+  auto arity = tupleArity(ctx);
+  CHECK_ERR(arity);
+  return ctx.makeTupleMake(pos, *arity);
 }
 
 template<typename Ctx> Result<> makeTupleExtract(Ctx& ctx, Index pos) {
-  return ctx.in.err("unimplemented instruction");
+  auto arity = tupleArity(ctx);
+  CHECK_ERR(arity);
+  auto index = ctx.in.takeU32();
+  if (!index) {
+    return ctx.in.err("expected tuple index");
+  }
+  return ctx.makeTupleExtract(pos, *arity, *index);
 }
 
 template<typename Ctx> Result<> makeTupleDrop(Ctx& ctx, Index pos) {
-  return ctx.in.err("unimplemented instruction");
+  auto arity = tupleArity(ctx);
+  CHECK_ERR(arity);
+  return ctx.makeTupleDrop(pos, *arity);
 }
 
 template<typename Ctx>

diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
@@ -157,8 +157,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   [[nodiscard]] Result<> makeTry(Name label, Type type);
   [[nodiscard]] Result<> makeThrow(Name tag);
   [[nodiscard]] Result<> makeRethrow(Index label);
-  // [[nodiscard]] Result<> makeTupleMake();
-  // [[nodiscard]] Result<> makeTupleExtract();
+  [[nodiscard]] Result<> makeTupleMake(uint32_t arity);
+  [[nodiscard]] Result<> makeTupleExtract(uint32_t arity, uint32_t index);
+  [[nodiscard]] Result<> makeTupleDrop(uint32_t arity);
   [[nodiscard]] Result<> makeRefI31();
   [[nodiscard]] Result<> makeI31Get(bool signed_);
   [[nodiscard]] Result<> makeCallRef(HeapType type, bool isReturn);
@@ -217,6 +218,10 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   [[nodiscard]] Result<> visitThrow(Throw*);
   [[nodiscard]] Result<> visitStringNew(StringNew*);
   [[nodiscard]] Result<> visitStringEncode(StringEncode*);
+  [[nodiscard]] Result<> visitTupleMake(TupleMake*);
+  [[nodiscard]] Result<>
+  visitTupleExtract(TupleExtract*,
+                    std::optional<uint32_t> arity = std::nullopt);
 
 private:
   Module& wasm;

diff --git a/src/wasm.h b/src/wasm.h
@@ -1501,6 +1501,7 @@ class TupleMake : public SpecificExpression<Expression::TupleMakeId> {
 
 class TupleExtract : public SpecificExpression<Expression::TupleExtractId> {
 public:
+  TupleExtract() = default;
   TupleExtract(MixedArena& allocator) {}
 
   Expression* tuple;

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
@@ -107,17 +107,17 @@ Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted,
 
   auto type = scope.exprStack.back()->type;
 
-  if (type.size() == sizeHint) {
+  if (type.size() == sizeHint || type.size() <= 1) {
     if (hoisted.get) {
       packageAsBlock(type);
     }
     return Ok{};
   }
 
-  // We need to break up the hoisted tuple. Create and push a block setting the
-  // tuple to a local and returning its first element, then push additional gets
-  // of each of its subsequent elements. Reuse the scratch local we used for
-  // hoisting, if it exists.
+  // We need to break up the hoisted tuple. Create and push an expression
+  // setting the tuple to a local and returning its first element, then push
+  // additional gets of each of its subsequent elements. Reuse the scratch local
+  // we used for hoisting, if it exists.
   Index scratchIdx;
   if (hoisted.get) {
     // Update the get on top of the stack to just return the first element.
@@ -127,12 +127,8 @@ Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted,
   } else {
     auto scratch = addScratchLocal(type);
     CHECK_ERR(scratch);
-    auto* block = builder.makeSequence(
-      builder.makeLocalSet(*scratch, scope.exprStack.back()),
-      builder.makeTupleExtract(builder.makeLocalGet(*scratch, type), 0),
-      type[0]);
-    scope.exprStack.pop_back();
-    push(block);
+    scope.exprStack.back() = builder.makeTupleExtract(
+      builder.makeLocalTee(*scratch, scope.exprStack.back(), type), 0);
     scratchIdx = *scratch;
   }
   for (Index i = 1, size = type.size(); i < size; ++i) {
@@ -560,6 +556,33 @@ Result<> IRBuilder::visitStringEncode(StringEncode* curr) {
   WASM_UNREACHABLE("unexpected op");
 }
 
+Result<> IRBuilder::visitTupleMake(TupleMake* curr) {
+  assert(curr->operands.size() >= 2);
+  for (size_t i = 0, size = curr->operands.size(); i < size; ++i) {
+    auto elem = pop();
+    CHECK_ERR(elem);
+    curr->operands[size - 1 - i] = *elem;
+  }
+  return Ok{};
+}
+
+Result<> IRBuilder::visitTupleExtract(TupleExtract* curr,
+                                      std::optional<uint32_t> arity) {
+  if (!arity) {
+    if (curr->tuple->type == Type::unreachable) {
+      // Fallback to an arbitrary valid arity.
+      arity = 2;
+    } else {
+      arity = curr->tuple->type.size();
+    }
+  }
+  assert(*arity >= 2);
+  auto tuple = pop(*arity);
+  CHECK_ERR(tuple);
+  curr->tuple = *tuple;
+  return Ok{};
+}
+
 Result<> IRBuilder::visitFunctionStart(Function* func) {
   if (!scopeStack.empty()) {
     return Err{"unexpected start of function"};
@@ -1332,9 +1355,27 @@ Result<> IRBuilder::makeRethrow(Index label) {
   return Ok{};
 }
 
-// Result<> IRBuilder::makeTupleMake() {}
+Result<> IRBuilder::makeTupleMake(uint32_t arity) {
+  TupleMake curr(wasm.allocator);
+  curr.operands.resize(arity);
+  CHECK_ERR(visitTupleMake(&curr));
+  push(builder.makeTupleMake(curr.operands));
+  return Ok{};
+}
+
+Result<> IRBuilder::makeTupleExtract(uint32_t arity, uint32_t index) {
+  TupleExtract curr;
+  CHECK_ERR(visitTupleExtract(&curr, arity));
+  push(builder.makeTupleExtract(curr.tuple, index));
+  return Ok{};
+}
 
-// Result<> IRBuilder::makeTupleExtract() {}
+Result<> IRBuilder::makeTupleDrop(uint32_t arity) {
+  Drop curr;
+  CHECK_ERR(visitDrop(&curr, arity));
+  push(builder.makeDrop(curr.value));
+  return Ok{};
+}
 
 Result<> IRBuilder::makeRefI31() {
   RefI31 curr;

diff --git a/test/lit/passes/outlining.wast b/test/lit/passes/outlining.wast
@@ -645,13 +645,10 @@
   ;; CHECK-NEXT:  (local $scratch_1 (i32 i32))
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (i32.add
-  ;; CHECK-NEXT:    (block (result i32)
-  ;; CHECK-NEXT:     (local.set $scratch
+  ;; CHECK-NEXT:    (tuple.extract 2 0
+  ;; CHECK-NEXT:     (local.tee $scratch
   ;; CHECK-NEXT:      (call $outline$)
   ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (tuple.extract 2 0
-  ;; CHECK-NEXT:      (local.get $scratch)
-  ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (tuple.extract 2 1
   ;; CHECK-NEXT:     (local.get $scratch)
@@ -660,13 +657,10 @@
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
   ;; CHECK-NEXT:   (i32.mul
-  ;; CHECK-NEXT:    (block (result i32)
-  ;; CHECK-NEXT:     (local.set $scratch_1
+  ;; CHECK-NEXT:    (tuple.extract 2 0
+  ;; CHECK-NEXT:     (local.tee $scratch_1
   ;; CHECK-NEXT:      (call $outline$)
   ;; CHECK-NEXT:     )
-  ;; CHECK-NEXT:     (tuple.extract 2 0
-  ;; CHECK-NEXT:      (local.get $scratch_1)
-  ;; CHECK-NEXT:     )
   ;; CHECK-NEXT:    )
   ;; CHECK-NEXT:    (tuple.extract 2 1
   ;; CHECK-NEXT:     (local.get $scratch_1)