Generalize binary writing for tuples (#2670)

Updates `BinaryInstWriter::mapLocalsAndEmitHeader` so it no longer hardcodes
each possible local type. Also adds a new inner loop over the elements of any
local tuple type in the IR. Updates the map from IR local indices to binary
indices to be additionally keyed on the index within a tuple type. Since we do
not generate tuple types yet, this additional index is hardcoded to zero
everywhere it is used for now. A later PR adding tuple creation operations will
extend this functionality and add tests.

2 files changed, 26 insertions, 91 deletions

diff --git a/src/wasm-stack.h b/src/wasm-stack.h
index 863bbbf51..b42ae1253 100644
--- a/src/wasm-stack.h
+++ b/src/wasm-stack.h
@@ -166,8 +166,8 @@ private:
 
   // type => number of locals of that type in the compact form
   std::map<Type, size_t> numLocalsByType;
-  // local index => index in compact form of [all int32s][all int64s]etc
-  std::map<Index, size_t> mappedLocals;
+  // (local index, tuple index) => binary local index
+  std::map<std::pair<Index, Index>, size_t> mappedLocals;
 };
 
 // Takes binaryen IR and converts it to something else (binary or stack IR)
diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index dee0caa82..2e4cd4344 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -76,12 +76,13 @@ void BinaryInstWriter::visitCallIndirect(CallIndirect* curr) {
 }
 
 void BinaryInstWriter::visitLocalGet(LocalGet* curr) {
-  o << int8_t(BinaryConsts::LocalGet) << U32LEB(mappedLocals[curr->index]);
+  o << int8_t(BinaryConsts::LocalGet)
+    << U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
 }
 
 void BinaryInstWriter::visitLocalSet(LocalSet* curr) {
   o << int8_t(curr->isTee() ? BinaryConsts::LocalTee : BinaryConsts::LocalSet)
-    << U32LEB(mappedLocals[curr->index]);
+    << U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
 }
 
 void BinaryInstWriter::visitGlobalGet(GlobalGet* curr) {
@@ -1656,102 +1657,36 @@ void BinaryInstWriter::emitUnreachable() {
 
 void BinaryInstWriter::mapLocalsAndEmitHeader() {
   assert(func && "BinaryInstWriter: function is not set");
-  // Map them
+  // Map params
   for (Index i = 0; i < func->getNumParams(); i++) {
     size_t curr = mappedLocals.size();
-    mappedLocals[i] = curr;
+    mappedLocals[std::make_pair(i, 0)] = curr;
   }
   for (auto type : func->vars) {
-    numLocalsByType[type]++;
+    for (auto t : type.expand()) {
+      numLocalsByType[t]++;
+    }
   }
   std::map<Type, size_t> currLocalsByType;
   for (Index i = func->getVarIndexBase(); i < func->getNumLocals(); i++) {
-    size_t index = func->getVarIndexBase();
-    Type type = func->getLocalType(i);
-    // increment now for simplicity, must decrement it in returns
-    currLocalsByType[type]++;
-    if (type == Type::i32) {
-      mappedLocals[i] = index + currLocalsByType[Type::i32] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::i32];
-    if (type == Type::i64) {
-      mappedLocals[i] = index + currLocalsByType[Type::i64] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::i64];
-    if (type == Type::f32) {
-      mappedLocals[i] = index + currLocalsByType[Type::f32] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::f32];
-    if (type == Type::f64) {
-      mappedLocals[i] = index + currLocalsByType[Type::f64] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::f64];
-    if (type == Type::v128) {
-      mappedLocals[i] = index + currLocalsByType[Type::v128] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::v128];
-    if (type == Type::funcref) {
-      mappedLocals[i] = index + currLocalsByType[Type::funcref] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::funcref];
-    if (type == Type::anyref) {
-      mappedLocals[i] = index + currLocalsByType[Type::anyref] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::anyref];
-    if (type == Type::nullref) {
-      mappedLocals[i] = index + currLocalsByType[Type::nullref] - 1;
-      continue;
-    }
-    index += numLocalsByType[Type::nullref];
-    if (type == Type::exnref) {
-      mappedLocals[i] = index + currLocalsByType[Type::exnref] - 1;
-      continue;
+    const std::vector<Type> types = func->getLocalType(i).expand();
+    for (Index j = 0; j < types.size(); j++) {
+      Type type = types[j];
+      auto fullIndex = std::make_pair(i, j);
+      size_t index = func->getVarIndexBase();
+      for (auto& typeCount : numLocalsByType) {
+        if (type == typeCount.first) {
+          mappedLocals[fullIndex] = index + currLocalsByType[typeCount.first];
+          currLocalsByType[type]++;
+          break;
+        }
+        index += typeCount.second;
+      }
     }
-    WASM_UNREACHABLE("unexpected type");
-  }
-  // Emit them.
-  o << U32LEB((numLocalsByType[Type::i32] ? 1 : 0) +
-              (numLocalsByType[Type::i64] ? 1 : 0) +
-              (numLocalsByType[Type::f32] ? 1 : 0) +
-              (numLocalsByType[Type::f64] ? 1 : 0) +
-              (numLocalsByType[Type::v128] ? 1 : 0) +
-              (numLocalsByType[Type::funcref] ? 1 : 0) +
-              (numLocalsByType[Type::anyref] ? 1 : 0) +
-              (numLocalsByType[Type::nullref] ? 1 : 0) +
-              (numLocalsByType[Type::exnref] ? 1 : 0));
-  if (numLocalsByType[Type::i32]) {
-    o << U32LEB(numLocalsByType[Type::i32]) << binaryType(Type::i32);
-  }
-  if (numLocalsByType[Type::i64]) {
-    o << U32LEB(numLocalsByType[Type::i64]) << binaryType(Type::i64);
-  }
-  if (numLocalsByType[Type::f32]) {
-    o << U32LEB(numLocalsByType[Type::f32]) << binaryType(Type::f32);
-  }
-  if (numLocalsByType[Type::f64]) {
-    o << U32LEB(numLocalsByType[Type::f64]) << binaryType(Type::f64);
-  }
-  if (numLocalsByType[Type::v128]) {
-    o << U32LEB(numLocalsByType[Type::v128]) << binaryType(Type::v128);
-  }
-  if (numLocalsByType[Type::funcref]) {
-    o << U32LEB(numLocalsByType[Type::funcref]) << binaryType(Type::funcref);
-  }
-  if (numLocalsByType[Type::anyref]) {
-    o << U32LEB(numLocalsByType[Type::anyref]) << binaryType(Type::anyref);
-  }
-  if (numLocalsByType[Type::nullref]) {
-    o << U32LEB(numLocalsByType[Type::nullref]) << binaryType(Type::nullref);
   }
-  if (numLocalsByType[Type::exnref]) {
-    o << U32LEB(numLocalsByType[Type::exnref]) << binaryType(Type::exnref);
+  o << U32LEB(numLocalsByType.size());
+  for (auto& typeCount : numLocalsByType) {
+    o << U32LEB(typeCount.second) << binaryType(typeCount.first);
   }
 }