Support tuple locals in Asyncify (#2696)

Iterate over tuple locals and separately load or store each component.

4 files changed, 154 insertions, 110 deletions

diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index f8fd6085b..05c1c8954 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -1115,9 +1115,7 @@ private:
     auto* func = getFunction();
     Index total = 0;
     for (Index i = 0; i < numPreservableLocals; i++) {
-      auto type = func->getLocalType(i);
-      auto size = type.getByteSize();
-      total += size;
+      total += func->getLocalType(i).getByteSize();
     }
     auto* block = builder->makeBlock();
     block->list.push_back(builder->makeIncStackPos(-total));
@@ -1126,19 +1124,31 @@ private:
       builder->makeLocalSet(tempIndex, builder->makeGetStackPos()));
     Index offset = 0;
     for (Index i = 0; i < numPreservableLocals; i++) {
-      auto type = func->getLocalType(i);
-      auto size = type.getByteSize();
-      assert(size % STACK_ALIGN == 0);
-      // TODO: higher alignment?
-      block->list.push_back(builder->makeLocalSet(
-        i,
-        builder->makeLoad(size,
-                          true,
-                          offset,
-                          STACK_ALIGN,
-                          builder->makeLocalGet(tempIndex, Type::i32),
-                          type)));
-      offset += size;
+      const auto& types = func->getLocalType(i).expand();
+      SmallVector<Expression*, 1> loads;
+      for (Index j = 0; j < types.size(); j++) {
+        auto type = types[j];
+        auto size = type.getByteSize();
+        assert(size % STACK_ALIGN == 0);
+        // TODO: higher alignment?
+        loads.push_back(
+          builder->makeLoad(size,
+                            true,
+                            offset,
+                            STACK_ALIGN,
+                            builder->makeLocalGet(tempIndex, Type::i32),
+                            type));
+        offset += size;
+      }
+      Expression* load;
+      if (loads.size() == 1) {
+        load = loads[0];
+      } else if (types.size() > 1) {
+        load = builder->makeTupleMake(std::move(loads));
+      } else {
+        WASM_UNREACHABLE("Unexpected empty type");
+      }
+      block->list.push_back(builder->makeLocalSet(i, load));
     }
     block->finalize();
     return block;
@@ -1155,18 +1165,26 @@ private:
       builder->makeLocalSet(tempIndex, builder->makeGetStackPos()));
     Index offset = 0;
     for (Index i = 0; i < numPreservableLocals; i++) {
-      auto type = func->getLocalType(i);
-      auto size = type.getByteSize();
-      assert(size % STACK_ALIGN == 0);
-      // TODO: higher alignment?
-      block->list.push_back(
-        builder->makeStore(size,
-                           offset,
-                           STACK_ALIGN,
-                           builder->makeLocalGet(tempIndex, Type::i32),
-                           builder->makeLocalGet(i, type),
-                           type));
-      offset += size;
+      auto localType = func->getLocalType(i);
+      const auto& types = localType.expand();
+      for (Index j = 0; j < types.size(); j++) {
+        auto type = types[j];
+        auto size = type.getByteSize();
+        Expression* localGet = builder->makeLocalGet(i, localType);
+        if (types.size() > 1) {
+          localGet = builder->makeTupleExtract(localGet, j);
+        }
+        assert(size % STACK_ALIGN == 0);
+        // TODO: higher alignment?
+        block->list.push_back(
+          builder->makeStore(size,
+                             offset,
+                             STACK_ALIGN,
+                             builder->makeLocalGet(tempIndex, Type::i32),
+                             localGet,
+                             type));
+        offset += size;
+      }
     }
     block->list.push_back(builder->makeIncStackPos(offset));
     block->finalize();
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 024894dfb..a98d10be4 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -145,28 +145,31 @@ bool Type::operator<(const Type& other) const {
 }
 
 unsigned Type::getByteSize() const {
-  assert(isSingle() && "getByteSize does not works with single types");
-  Type singleType = *expand().begin();
-  switch (singleType.getSingle()) {
-    case Type::i32:
-      return 4;
-    case Type::i64:
-      return 8;
-    case Type::f32:
-      return 4;
-    case Type::f64:
-      return 8;
-    case Type::v128:
-      return 16;
-    case Type::funcref:
-    case Type::anyref:
-    case Type::nullref:
-    case Type::exnref:
-    case Type::none:
-    case Type::unreachable:
-      WASM_UNREACHABLE("invalid type");
+  // TODO: alignment?
+  unsigned size = 0;
+  for (auto t : expand()) {
+    switch (t.getSingle()) {
+      case Type::i32:
+      case Type::f32:
+        size += 4;
+        break;
+      case Type::i64:
+      case Type::f64:
+        size += 8;
+        break;
+      case Type::v128:
+        size += 16;
+        break;
+      case Type::funcref:
+      case Type::anyref:
+      case Type::nullref:
+      case Type::exnref:
+      case Type::none:
+      case Type::unreachable:
+        WASM_UNREACHABLE("invalid type");
+    }
   }
-  WASM_UNREACHABLE("invalid type");
+  return size;
 }
 
 Type Type::reinterpret() const {
diff --git a/test/passes/asyncify.txt b/test/passes/asyncify.txt
index 8e349a637..07fecce10 100644
--- a/test/passes/asyncify.txt
+++ b/test/passes/asyncify.txt
@@ -779,7 +779,7 @@
  )
  (func $many-locals (; 7 ;) (param $x i32) (result i32)
   (local $y i32)
-  (local $2 i32)
+  (local $z (f32 i64))
   (local $3 i32)
   (local $4 i32)
   (local $5 i32)
@@ -791,6 +791,7 @@
   (local $11 i32)
   (local $12 i32)
   (local $13 i32)
+  (local $14 i32)
   (if
    (i32.eq
     (global.get $__asyncify_state)
@@ -803,67 +804,77 @@
       (i32.load
        (global.get $__asyncify_data)
       )
-      (i32.const -40)
+      (i32.const -52)
      )
     )
-    (local.set $12
+    (local.set $13
      (i32.load
       (global.get $__asyncify_data)
      )
     )
     (local.set $x
      (i32.load
-      (local.get $12)
+      (local.get $13)
      )
     )
     (local.set $y
      (i32.load offset=4
-      (local.get $12)
+      (local.get $13)
      )
     )
-    (local.set $2
-     (i32.load offset=8
-      (local.get $12)
+    (local.set $z
+     (tuple.make
+      (f32.load offset=8
+       (local.get $13)
+      )
+      (i64.load offset=12 align=4
+       (local.get $13)
+      )
      )
     )
     (local.set $3
-     (i32.load offset=12
-      (local.get $12)
+     (i32.load offset=20
+      (local.get $13)
      )
     )
     (local.set $4
-     (i32.load offset=16
-      (local.get $12)
+     (i32.load offset=24
+      (local.get $13)
      )
     )
     (local.set $5
-     (i32.load offset=20
-      (local.get $12)
+     (i32.load offset=28
+      (local.get $13)
      )
     )
     (local.set $6
-     (i32.load offset=24
-      (local.get $12)
+     (i32.load offset=32
+      (local.get $13)
      )
     )
     (local.set $7
-     (i32.load offset=28
-      (local.get $12)
+     (i32.load offset=36
+      (local.get $13)
      )
     )
     (local.set $8
-     (i32.load offset=32
-      (local.get $12)
+     (i32.load offset=40
+      (local.get $13)
      )
     )
     (local.set $9
-     (i32.load offset=36
-      (local.get $12)
+     (i32.load offset=44
+      (local.get $13)
+     )
+    )
+    (local.set $10
+     (i32.load offset=48
+      (local.get $13)
      )
     )
    )
   )
-  (local.set $10
+  (local.set $11
    (block $__asyncify_unwind (result i32)
     (block
      (block
@@ -882,7 +893,7 @@
           (i32.const -4)
          )
         )
-        (local.set $11
+        (local.set $12
          (i32.load
           (i32.load
            (global.get $__asyncify_data)
@@ -898,35 +909,35 @@
          (i32.const 0)
         )
         (loop $l
-         (local.set $2
+         (local.set $3
           (local.get $y)
          )
-         (local.set $3
+         (local.set $4
           (i32.add
-           (local.get $2)
+           (local.get $3)
            (i32.const 1)
           )
          )
          (local.set $x
-          (local.get $3)
+          (local.get $4)
          )
-         (local.set $4
+         (local.set $5
           (local.get $x)
          )
-         (local.set $5
+         (local.set $6
           (i32.div_s
-           (local.get $4)
+           (local.get $5)
            (i32.const 3)
           )
          )
          (local.set $y
-          (local.get $5)
+          (local.get $6)
          )
-         (local.set $6
+         (local.set $7
           (local.get $y)
          )
          (br_if $l
-          (local.get $6)
+          (local.get $7)
          )
         )
        )
@@ -938,7 +949,7 @@
          )
          (i32.const 1)
          (i32.eq
-          (local.get $11)
+          (local.get $12)
           (i32.const 0)
          )
         )
@@ -961,11 +972,11 @@
          (i32.const 0)
         )
         (block
-         (local.set $7
+         (local.set $8
           (local.get $y)
          )
          (return
-          (local.get $7)
+          (local.get $8)
          )
         )
        )
@@ -982,7 +993,7 @@
     (i32.load
      (global.get $__asyncify_data)
     )
-    (local.get $10)
+    (local.get $11)
    )
    (i32.store
     (global.get $__asyncify_data)
@@ -995,58 +1006,70 @@
    )
   )
   (block
-   (local.set $13
+   (local.set $14
     (i32.load
      (global.get $__asyncify_data)
     )
    )
    (i32.store
-    (local.get $13)
+    (local.get $14)
     (local.get $x)
    )
    (i32.store offset=4
-    (local.get $13)
+    (local.get $14)
     (local.get $y)
    )
-   (i32.store offset=8
-    (local.get $13)
-    (local.get $2)
+   (f32.store offset=8
+    (local.get $14)
+    (tuple.extract 0
+     (local.get $z)
+    )
    )
-   (i32.store offset=12
-    (local.get $13)
+   (i64.store offset=12 align=4
+    (local.get $14)
+    (tuple.extract 1
+     (local.get $z)
+    )
+   )
+   (i32.store offset=20
+    (local.get $14)
     (local.get $3)
    )
-   (i32.store offset=16
-    (local.get $13)
+   (i32.store offset=24
+    (local.get $14)
     (local.get $4)
    )
-   (i32.store offset=20
-    (local.get $13)
+   (i32.store offset=28
+    (local.get $14)
     (local.get $5)
    )
-   (i32.store offset=24
-    (local.get $13)
+   (i32.store offset=32
+    (local.get $14)
     (local.get $6)
    )
-   (i32.store offset=28
-    (local.get $13)
+   (i32.store offset=36
+    (local.get $14)
     (local.get $7)
    )
-   (i32.store offset=32
-    (local.get $13)
+   (i32.store offset=40
+    (local.get $14)
     (local.get $8)
    )
-   (i32.store offset=36
-    (local.get $13)
+   (i32.store offset=44
+    (local.get $14)
     (local.get $9)
    )
+   (i32.store offset=48
+    (local.get $14)
+    (local.get $10)
+   )
    (i32.store
     (global.get $__asyncify_data)
     (i32.add
      (i32.load
       (global.get $__asyncify_data)
      )
-     (i32.const 40)
+     (i32.const 52)
     )
    )
   )
diff --git a/test/passes/asyncify.wast b/test/passes/asyncify.wast
index 6f4754b74..b34338bd0 100644
--- a/test/passes/asyncify.wast
+++ b/test/passes/asyncify.wast
@@ -35,7 +35,7 @@
     ;; work will sleep, so we exit through here while it is paused
   )
   ;; the second event called from the main event loop: to resume $work,
-  ;; stop the unwind, then prepare a rewind, and initiate it by doing 
+  ;; stop the unwind, then prepare a rewind, and initiate it by doing
   ;; the call to rewind the call stack back up to where it was
   (func $second_event
     (call $asyncify_stop_unwind)
@@ -71,6 +71,7 @@
   )
   (func $many-locals (param $x i32) (result i32)
     (local $y i32)
+    (local $z (f32 i64))
     (loop $l
       (local.set $x
         (i32.add (local.get $y) (i32.const 1))
@@ -149,4 +150,3 @@
 ;; empty module, in particular with no memory
 (module
 )
-