1 files changed, 90 insertions, 0 deletions

diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 4844dd5b4..69d3ecdbb 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -832,6 +832,51 @@ void WasmBinaryWriter::visitStore(Store *curr) {
   emitMemoryAccess(curr->align, curr->bytes, curr->offset);
 }
 
+void WasmBinaryWriter::visitAtomicRMW(AtomicRMW *curr) {
+  if (debug) std::cerr << "zz node: AtomicRMW" << std::endl;
+  recurse(curr->ptr);
+  recurse(curr->value);
+
+  o << int8_t(BinaryConsts::AtomicPrefix);
+
+#define CASE_FOR_OP(Op) \
+  case Op: \
+    switch (curr->type) {                                               \
+      case i32:                                                         \
+        switch (curr->bytes) {                                          \
+          case 1: o << int8_t(BinaryConsts::I32AtomicRMW##Op##8U); break; \
+          case 2: o << int8_t(BinaryConsts::I32AtomicRMW##Op##16U); break; \
+          case 4: o << int8_t(BinaryConsts::I32AtomicRMW##Op); break;   \
+          default: WASM_UNREACHABLE();                                  \
+        }                                                               \
+        break;                                                          \
+      case i64:                                                         \
+        switch (curr->bytes) {                                          \
+          case 1: o << int8_t(BinaryConsts::I64AtomicRMW##Op##8U); break; \
+          case 2: o << int8_t(BinaryConsts::I64AtomicRMW##Op##16U); break; \
+          case 4: o << int8_t(BinaryConsts::I64AtomicRMW##Op##32U); break; \
+          case 8: o << int8_t(BinaryConsts::I64AtomicRMW##Op); break;   \
+          default: WASM_UNREACHABLE();                                  \
+        }                                                               \
+        break;                                                          \
+      default: WASM_UNREACHABLE();                                      \
+    }                                                                   \
+    break
+
+  switch(curr->op) {
+    CASE_FOR_OP(Add);
+    CASE_FOR_OP(Sub);
+    CASE_FOR_OP(And);
+    CASE_FOR_OP(Or);
+    CASE_FOR_OP(Xor);
+    CASE_FOR_OP(Xchg);
+    default: WASM_UNREACHABLE();
+  }
+#undef CASE_FOR_OP
+
+  emitMemoryAccess(curr->bytes, curr->bytes, curr->offset);
+}
+
 void WasmBinaryWriter::visitConst(Const *curr) {
   if (debug) std::cerr << "zz node: Const" << curr << " : " << curr->type << std::endl;
   switch (curr->type) {
@@ -1934,6 +1979,7 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       code = getInt8();
       if (maybeVisitLoad(curr, code, /*isAtomic=*/true)) break;
       if (maybeVisitStore(curr, code, /*isAtomic=*/true)) break;
+      if (maybeVisitAtomicRMW(curr, code)) break;
       throw ParseException("invalid code after atomic prefix: " + std::to_string(code));
     }
     default: {
@@ -2282,6 +2328,50 @@ bool WasmBinaryBuilder::maybeVisitStore(Expression*& out, uint8_t code, bool isA
   return true;
 }
 
+
+bool WasmBinaryBuilder::maybeVisitAtomicRMW(Expression*& out, uint8_t code) {
+  if (code < BinaryConsts::AtomicRMWOps_Begin || code > BinaryConsts::AtomicRMWOps_End) return false;
+  auto* curr = allocator.alloc<AtomicRMW>();
+
+  // Set curr to the given opcode, type and size.
+#define SET(opcode, optype, size) \
+  curr->op = opcode;              \
+  curr->type = optype;            \
+  curr->bytes = size
+
+  // Handle the cases for all the valid types for a particular opcode
+#define SET_FOR_OP(Op) \
+    case BinaryConsts::I32AtomicRMW##Op: SET(Op, i32, 4); break;      \
+    case BinaryConsts::I32AtomicRMW##Op##8U: SET(Op, i32, 1); break;  \
+    case BinaryConsts::I32AtomicRMW##Op##16U: SET(Op, i32, 2); break; \
+    case BinaryConsts::I64AtomicRMW##Op: SET(Op, i64, 8); break;      \
+    case BinaryConsts::I64AtomicRMW##Op##8U: SET(Op, i64, 1); break;  \
+    case BinaryConsts::I64AtomicRMW##Op##16U: SET(Op, i64, 2); break; \
+    case BinaryConsts::I64AtomicRMW##Op##32U: SET(Op, i64, 4); break;
+
+  switch(code) {
+    SET_FOR_OP(Add);
+    SET_FOR_OP(Sub);
+    SET_FOR_OP(And);
+    SET_FOR_OP(Or);
+    SET_FOR_OP(Xor);
+    SET_FOR_OP(Xchg);
+    default: WASM_UNREACHABLE();
+  }
+#undef SET_FOR_OP
+#undef SET
+
+  if (debug) std::cerr << "zz node: AtomicRMW" << std::endl;
+  Address readAlign;
+  readMemoryAccess(readAlign, curr->bytes, curr->offset);
+  if (readAlign != curr->bytes) throw ParseException("Align of AtomicRMW must match size");
+  curr->value = popNonVoidExpression();
+  curr->ptr = popNonVoidExpression();
+  curr->finalize();
+  out = curr;
+  return true;
+}
+
 bool WasmBinaryBuilder::maybeVisitConst(Expression*& out, uint8_t code) {
   Const* curr;
   if (debug) std::cerr << "zz node: Const, code " << code << std::endl;