Add IR, parsing, printing, and binary for atomic cmpxchg (#1083)

12 files changed, 287 insertions, 14 deletions

diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp
index d0c9603f5..33e06d7a5 100644
--- a/src/passes/Print.cpp
+++ b/src/passes/Print.cpp
@@ -350,15 +350,14 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     printFullLine(curr->value);
     decIndent();
   }
-  void visitAtomicRMW(AtomicRMW* curr) {
-    o << '(';
-    prepareColor(o) << printWasmType(curr->type) << ".atomic.rmw";
-    if (curr->bytes != getWasmTypeSize(curr->type)) {
-      if (curr->bytes == 1) {
+  static void printRMWSize(std::ostream& o, WasmType type, uint8_t bytes) {
+    prepareColor(o) << printWasmType(type) << ".atomic.rmw";
+    if (bytes != getWasmTypeSize(type)) {
+      if (bytes == 1) {
         o << '8';
-      } else if (curr->bytes == 2) {
+      } else if (bytes == 2) {
         o << "16";
-      } else if (curr->bytes == 4) {
+      } else if (bytes == 4) {
         o << "32";
       } else {
         WASM_UNREACHABLE();
@@ -366,6 +365,10 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
       o << "_u";
     }
     o << '.';
+  }
+  void visitAtomicRMW(AtomicRMW* curr) {
+    o << '(';
+    printRMWSize(o, curr->type, curr->bytes);
     switch (curr->op) {
       case Add:  o << "add";  break;
       case Sub:  o << "sub";  break;
@@ -383,6 +386,20 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     printFullLine(curr->value);
     decIndent();
   }
+  void visitAtomicCmpxchg(AtomicCmpxchg* curr) {
+    o << '(';
+    printRMWSize(o, curr->type, curr->bytes);
+    o << "cmpxchg";
+    restoreNormalColor(o);
+    if (curr->offset) {
+      o << " offset=" << curr->offset;
+    }
+    incIndent();
+    printFullLine(curr->ptr);
+    printFullLine(curr->expected);
+    printFullLine(curr->replacement);
+    decIndent();
+  }
   void visitConst(Const *curr) {
     o << curr->value;
   }
diff --git a/src/wasm-binary.h b/src/wasm-binary.h
index 332b4a7b6..6b984642c 100644
--- a/src/wasm-binary.h
+++ b/src/wasm-binary.h
@@ -571,6 +571,16 @@ enum AtomicOpcodes {
   I64AtomicRMWXchg16U = 0x46,
   I64AtomicRMWXchg32U = 0x47,
   AtomicRMWOps_End = 0x47,
+
+  AtomicCmpxchgOps_Begin = 0x48,
+  I32AtomicCmpxchg = 0x48,
+  I64AtomicCmpxchg = 0x49,
+  I32AtomicCmpxchg8U = 0x4a,
+  I32AtomicCmpxchg16U = 0x4b,
+  I64AtomicCmpxchg8U = 0x4c,
+  I64AtomicCmpxchg16U = 0x4d,
+  I64AtomicCmpxchg32U = 0x4e,
+  AtomicCmpxchgOps_End = 0x4e
 };
 
 
@@ -723,6 +733,7 @@ public:
   void visitLoad(Load *curr);
   void visitStore(Store *curr);
   void visitAtomicRMW(AtomicRMW *curr);
+  void visitAtomicCmpxchg(AtomicCmpxchg *curr);
   void visitConst(Const *curr);
   void visitUnary(Unary *curr);
   void visitBinary(Binary *curr);
@@ -881,6 +892,7 @@ public:
   bool maybeVisitLoad(Expression*& out, uint8_t code, bool isAtomic);
   bool maybeVisitStore(Expression*& out, uint8_t code, bool isAtomic);
   bool maybeVisitAtomicRMW(Expression*& out, uint8_t code);
+  bool maybeVisitAtomicCmpxchg(Expression*& out, uint8_t code);
   bool maybeVisitConst(Expression*& out, uint8_t code);
   bool maybeVisitUnary(Expression*& out, uint8_t code);
   bool maybeVisitBinary(Expression*& out, uint8_t code);
diff --git a/src/wasm-s-parser.h b/src/wasm-s-parser.h
index 161256c2f..d621128de 100644
--- a/src/wasm-s-parser.h
+++ b/src/wasm-s-parser.h
@@ -177,7 +177,9 @@ private:
   Expression* makeConst(Element& s, WasmType type);
   Expression* makeLoad(Element& s, WasmType type, bool isAtomic);
   Expression* makeStore(Element& s, WasmType type, bool isAtomic);
-  Expression* makeAtomicRMW(Element& s, WasmType type);
+  Expression* makeAtomicRMWOrCmpxchg(Element& s, WasmType type);
+  Expression* makeAtomicRMW(Element& s, WasmType type, uint8_t bytes, const char* extra);
+  Expression* makeAtomicCmpxchg(Element& s, WasmType type, uint8_t bytes, const char* extra);
   Expression* makeIf(Element& s);
   Expression* makeMaybeBlock(Element& s, size_t i, WasmType type);
   Expression* makeLoop(Element& s);
diff --git a/src/wasm-traversal.h b/src/wasm-traversal.h
index 11332b5c3..f56bfca96 100644
--- a/src/wasm-traversal.h
+++ b/src/wasm-traversal.h
@@ -50,6 +50,7 @@ struct Visitor {
   ReturnType visitLoad(Load* curr) {}
   ReturnType visitStore(Store* curr) {}
   ReturnType visitAtomicRMW(AtomicRMW* curr) {return ReturnType();} //Stub impl so not every pass has to implement this yet.
+  ReturnType visitAtomicCmpxchg(AtomicCmpxchg* curr) {return ReturnType();} //Stub impl so not every pass has to implement this yet.
   ReturnType visitConst(Const* curr) {}
   ReturnType visitUnary(Unary* curr) {}
   ReturnType visitBinary(Binary* curr) {}
@@ -92,6 +93,7 @@ struct Visitor {
       case Expression::Id::LoadId: DELEGATE(Load);
       case Expression::Id::StoreId: DELEGATE(Store);
       case Expression::Id::AtomicRMWId: DELEGATE(AtomicRMW);
+      case Expression::Id::AtomicCmpxchgId: DELEGATE(AtomicCmpxchg);
       case Expression::Id::ConstId: DELEGATE(Const);
       case Expression::Id::UnaryId: DELEGATE(Unary);
       case Expression::Id::BinaryId: DELEGATE(Binary);
@@ -133,6 +135,7 @@ struct UnifiedExpressionVisitor : public Visitor<SubType> {
   ReturnType visitLoad(Load* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
   ReturnType visitStore(Store* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
   ReturnType visitAtomicRMW(AtomicRMW* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
+  ReturnType visitAtomicCmpxchg(AtomicCmpxchg* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
   ReturnType visitConst(Const* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
   ReturnType visitUnary(Unary* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
   ReturnType visitBinary(Binary* curr) { return static_cast<SubType*>(this)->visitExpression(curr); }
@@ -310,6 +313,7 @@ struct Walker : public VisitorType {
   static void doVisitLoad(SubType* self, Expression** currp)         { self->visitLoad((*currp)->cast<Load>()); }
   static void doVisitStore(SubType* self, Expression** currp)        { self->visitStore((*currp)->cast<Store>()); }
   static void doVisitAtomicRMW(SubType* self, Expression** currp)    { self->visitAtomicRMW((*currp)->cast<AtomicRMW>()); }
+  static void doVisitAtomicCmpxchg(SubType* self, Expression** currp){ self->visitAtomicCmpxchg((*currp)->cast<AtomicCmpxchg>()); }
   static void doVisitConst(SubType* self, Expression** currp)        { self->visitConst((*currp)->cast<Const>()); }
   static void doVisitUnary(SubType* self, Expression** currp)        { self->visitUnary((*currp)->cast<Unary>()); }
   static void doVisitBinary(SubType* self, Expression** currp)       { self->visitBinary((*currp)->cast<Binary>()); }
@@ -438,6 +442,13 @@ struct PostWalker : public Walker<SubType, VisitorType> {
         self->pushTask(SubType::scan, &curr->cast<AtomicRMW>()->ptr);
         break;
       }
+      case Expression::Id::AtomicCmpxchgId: {
+        self->pushTask(SubType::doVisitAtomicCmpxchg, currp);
+        self->pushTask(SubType::scan, &curr->cast<AtomicCmpxchg>()->replacement);
+        self->pushTask(SubType::scan, &curr->cast<AtomicCmpxchg>()->expected);
+        self->pushTask(SubType::scan, &curr->cast<AtomicCmpxchg>()->ptr);
+        break;
+      }
       case Expression::Id::ConstId: {
         self->pushTask(SubType::doVisitConst, currp);
         break;
diff --git a/src/wasm.h b/src/wasm.h
index e98c4db63..1f9e1c7af 100644
--- a/src/wasm.h
+++ b/src/wasm.h
@@ -131,7 +131,7 @@ enum HostOp {
 };
 
 enum AtomicRMWOp {
-  Add, Sub, And, Or, Xor, Xchg,
+  Add, Sub, And, Or, Xor, Xchg
 };
 
 //
@@ -445,6 +445,15 @@ class AtomicRMW : public SpecificExpression<Expression::AtomicRMWId> {
 class AtomicCmpxchg : public SpecificExpression<Expression::AtomicCmpxchgId> {
  public:
   AtomicCmpxchg() = default;
+  AtomicCmpxchg(MixedArena& allocator) : AtomicCmpxchg() {}
+
+  uint8_t bytes;
+  Address offset;
+  Expression* ptr;
+  Expression* expected;
+  Expression* replacement;
+
+  void finalize();
 };
 
 class Const : public SpecificExpression<Expression::ConstId> {
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 69d3ecdbb..efbd82e8b 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -877,6 +877,37 @@ void WasmBinaryWriter::visitAtomicRMW(AtomicRMW *curr) {
   emitMemoryAccess(curr->bytes, curr->bytes, curr->offset);
 }
 
+void WasmBinaryWriter::visitAtomicCmpxchg(AtomicCmpxchg *curr) {
+  if (debug) std::cerr << "zz node: AtomicCmpxchg" << std::endl;
+  recurse(curr->ptr);
+  recurse(curr->expected);
+  recurse(curr->replacement);
+
+  o << int8_t(BinaryConsts::AtomicPrefix);
+  switch (curr->type) {
+    case i32:
+      switch (curr->bytes) {
+        case 1: o << int8_t(BinaryConsts::I32AtomicCmpxchg8U); break;
+        case 2: o << int8_t(BinaryConsts::I32AtomicCmpxchg16U); break;
+        case 4: o << int8_t(BinaryConsts::I32AtomicCmpxchg); break;
+        default: WASM_UNREACHABLE();
+      }
+      break;
+    case i64:
+      switch (curr->bytes) {
+        case 1: o << int8_t(BinaryConsts::I64AtomicCmpxchg8U); break;
+        case 2: o << int8_t(BinaryConsts::I64AtomicCmpxchg16U); break;
+        case 4: o << int8_t(BinaryConsts::I64AtomicCmpxchg32U); break;
+        case 8: o << int8_t(BinaryConsts::I64AtomicCmpxchg); break;
+        default: WASM_UNREACHABLE();
+      }
+      break;
+    default: WASM_UNREACHABLE();
+  }
+  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) {
@@ -1980,6 +2011,7 @@ BinaryConsts::ASTNodes WasmBinaryBuilder::readExpression(Expression*& curr) {
       if (maybeVisitLoad(curr, code, /*isAtomic=*/true)) break;
       if (maybeVisitStore(curr, code, /*isAtomic=*/true)) break;
       if (maybeVisitAtomicRMW(curr, code)) break;
+      if (maybeVisitAtomicCmpxchg(curr, code)) break;
       throw ParseException("invalid code after atomic prefix: " + std::to_string(code));
     }
     default: {
@@ -2372,6 +2404,38 @@ bool WasmBinaryBuilder::maybeVisitAtomicRMW(Expression*& out, uint8_t code) {
   return true;
 }
 
+bool WasmBinaryBuilder::maybeVisitAtomicCmpxchg(Expression*& out, uint8_t code) {
+  if (code < BinaryConsts::AtomicCmpxchgOps_Begin || code > BinaryConsts::AtomicCmpxchgOps_End) return false;
+  auto* curr = allocator.alloc<AtomicCmpxchg>();
+
+  // Set curr to the given type and size.
+#define SET(optype, size)         \
+  curr->type = optype;            \
+  curr->bytes = size
+
+  switch (code) {
+    case BinaryConsts::I32AtomicCmpxchg: SET(i32, 4); break;
+    case BinaryConsts::I64AtomicCmpxchg: SET(i64, 8); break;
+    case BinaryConsts::I32AtomicCmpxchg8U: SET(i32, 1); break;
+    case BinaryConsts::I32AtomicCmpxchg16U: SET(i32, 2); break;
+    case BinaryConsts::I64AtomicCmpxchg8U: SET(i64, 1); break;
+    case BinaryConsts::I64AtomicCmpxchg16U: SET(i64, 2); break;
+    case BinaryConsts::I64AtomicCmpxchg32U: SET(i64, 4); break;
+    default: WASM_UNREACHABLE();
+  }
+
+  if (debug) std::cerr << "zz node: AtomicCmpxchg" << std::endl;
+  Address readAlign;
+  readMemoryAccess(readAlign, curr->bytes, curr->offset);
+  if (readAlign != curr->bytes) throw ParseException("Align of AtomicCpxchg must match size");
+  curr->replacement = popNonVoidExpression();
+  curr->expected = 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;
diff --git a/src/wasm/wasm-s-parser.cpp b/src/wasm/wasm-s-parser.cpp
index c6331afc0..2fede28fc 100644
--- a/src/wasm/wasm-s-parser.cpp
+++ b/src/wasm/wasm-s-parser.cpp
@@ -666,7 +666,7 @@ Expression* SExpressionWasmBuilder::makeExpression(Element& s) {
         if (op[1] == 't' && !strncmp(op, "atomic.", strlen("atomic."))) {
           if (op[7] == 'l') return makeLoad(s, type, /*isAtomic=*/true);
           if (op[7] == 's') return makeStore(s, type, /*isAtomic=*/true);
-          if (op[7] == 'r') return makeAtomicRMW(s, type);
+          if (op[7] == 'r') return makeAtomicRMWOrCmpxchg(s, type);
         }
         abort_on(op);
       }
@@ -1197,14 +1197,20 @@ Expression* SExpressionWasmBuilder::makeStore(Element& s, WasmType type, bool is
   return ret;
 }
 
-Expression* SExpressionWasmBuilder::makeAtomicRMW(Element& s, WasmType type) {
+Expression* SExpressionWasmBuilder::makeAtomicRMWOrCmpxchg(Element& s, WasmType type) {
   const char* extra = strchr(s[0]->c_str(), '.') + 11; // afer "type.atomic.rmw"
-  auto ret = allocator.alloc<AtomicRMW>();
-  ret->type = type;
-  ret->bytes = parseMemBytes(&extra, getWasmTypeSize(type));
+  auto bytes = parseMemBytes(&extra, getWasmTypeSize(type));
   extra = strchr(extra, '.'); // after the optional '_u' and before the opcode
   if (!extra) throw ParseException("malformed atomic rmw instruction");
   extra++; // after the '.'
+  if (!strncmp(extra, "cmpxchg", 7)) return makeAtomicCmpxchg(s, type, bytes, extra);
+  return makeAtomicRMW(s, type, bytes, extra);
+}
+
+Expression* SExpressionWasmBuilder::makeAtomicRMW(Element& s, WasmType type, uint8_t bytes, const char* extra) {
+  auto ret = allocator.alloc<AtomicRMW>();
+  ret->type = type;
+  ret->bytes = bytes;
   if (!strncmp(extra, "add", 3)) ret->op = Add;
   else if (!strncmp(extra, "and", 3)) ret->op = And;
   else if (!strncmp(extra, "or", 2)) ret->op = Or;
@@ -1221,6 +1227,20 @@ Expression* SExpressionWasmBuilder::makeAtomicRMW(Element& s, WasmType type) {
   return ret;
 }
 
+Expression* SExpressionWasmBuilder::makeAtomicCmpxchg(Element& s, WasmType type, uint8_t bytes, const char* extra) {
+  auto ret = allocator.alloc<AtomicCmpxchg>();
+  ret->type = type;
+  ret->bytes = bytes;
+  Address align;
+  size_t i = parseMemAttributes(s, &ret->offset, &align, ret->bytes);
+  if (align != ret->bytes) throw ParseException("Align of Atomic Cmpxchg must match size");
+  ret->ptr = parseExpression(s[i]);
+  ret->expected = parseExpression(s[i+1]);
+  ret->replacement = parseExpression(s[i+2]);
+  ret->finalize();
+  return ret;
+}
+
 Expression* SExpressionWasmBuilder::makeIf(Element& s) {
   auto ret = allocator.alloc<If>();
   Index i = 1;
diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index cbfcbccad..d331ba3a2 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -359,6 +359,12 @@ void AtomicRMW::finalize() {
   }
 }
 
+void AtomicCmpxchg::finalize() {
+  if (ptr->type == unreachable || expected->type == unreachable || replacement->type == unreachable) {
+    type = unreachable;
+  }
+}
+
 Const* Const::set(Literal value_) {
   value = value_;
   type = value.type;
diff --git a/test/atomics.wast b/test/atomics.wast
index 26aebdb0d..6b85c80ef 100644
--- a/test/atomics.wast
+++ b/test/atomics.wast
@@ -102,4 +102,36 @@
    )
   )
  )
+ (func $atomic-cmpxchg (type $0)
+  (local $0 i32)
+  (local $1 i32)
+  (drop
+   (i32.atomic.rmw.cmpxchg offset=4
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i32.atomic.rmw8_u.cmpxchg
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i64.atomic.rmw.cmpxchg offset=4
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i64.atomic.rmw32_u.cmpxchg align=4
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+ )
 )
diff --git a/test/atomics.wast.from-wast b/test/atomics.wast.from-wast
index ef15de7e2..01735c1fc 100644
--- a/test/atomics.wast.from-wast
+++ b/test/atomics.wast.from-wast
@@ -102,4 +102,36 @@
    )
   )
  )
+ (func $atomic-cmpxchg (type $0)
+  (local $0 i32)
+  (local $1 i32)
+  (drop
+   (i32.atomic.rmw.cmpxchg offset=4
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i32.atomic.rmw8_u.cmpxchg
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i64.atomic.rmw.cmpxchg offset=4
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+  (drop
+   (i64.atomic.rmw32_u.cmpxchg
+    (get_local $0)
+    (get_local $0)
+    (get_local $0)
+   )
+  )
+ )
 )
diff --git a/test/atomics.wast.fromBinary b/test/atomics.wast.fromBinary
index b3bce034a..63809f1ed 100644
--- a/test/atomics.wast.fromBinary
+++ b/test/atomics.wast.fromBinary
@@ -106,5 +106,39 @@
    )
   )
  )
+ (func $atomic-cmpxchg (type $0)
+  (local $var$0 i32)
+  (local $var$1 i32)
+  (block $label$0
+   (drop
+    (i32.atomic.rmw.cmpxchg offset=4
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i32.atomic.rmw8_u.cmpxchg
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i64.atomic.rmw.cmpxchg offset=4
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i64.atomic.rmw32_u.cmpxchg
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+  )
+ )
 )
 
diff --git a/test/atomics.wast.fromBinary.noDebugInfo b/test/atomics.wast.fromBinary.noDebugInfo
index 3777417bf..8dcca9d15 100644
--- a/test/atomics.wast.fromBinary.noDebugInfo
+++ b/test/atomics.wast.fromBinary.noDebugInfo
@@ -106,5 +106,39 @@
    )
   )
  )
+ (func $2 (type $0)
+  (local $var$0 i32)
+  (local $var$1 i32)
+  (block $label$0
+   (drop
+    (i32.atomic.rmw.cmpxchg offset=4
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i32.atomic.rmw8_u.cmpxchg
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i64.atomic.rmw.cmpxchg offset=4
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+   (drop
+    (i64.atomic.rmw32_u.cmpxchg
+     (get_local $var$0)
+     (get_local $var$0)
+     (get_local $var$0)
+    )
+   )
+  )
+ )
 )