Atomic wait/wake fixes (#1383)

* fix wait and wake binary format support, they have alignments and offsets

* don't emit unreachable parts of atomic operations, for simplicity and to avoid special handling

* don't emit atomic waits by default in the fuzzer, they hang in native vm support

Author: kripken

src/binaryen-c.cpp

@@ -809,7 +809,7 @@ BinaryenExpressionRef BinaryenAtomicCmpxchg(BinaryenModuleRef module, BinaryenIn
   return static_cast<Expression*>(ret);
 }
 BinaryenExpressionRef BinaryenAtomicWait(BinaryenModuleRef module, BinaryenExpressionRef ptr, BinaryenExpressionRef expected, BinaryenExpressionRef timeout, BinaryenType expectedType) {
-  auto* ret = Builder(*((Module*)module)).makeAtomicWait((Expression*)ptr, (Expression*)expected, (Expression*)timeout, WasmType(expectedType));
+  auto* ret = Builder(*((Module*)module)).makeAtomicWait((Expression*)ptr, (Expression*)expected, (Expression*)timeout, WasmType(expectedType), 0);
 
   if (tracing) {
     auto id = noteExpression(ret);
@@ -819,7 +819,7 @@ BinaryenExpressionRef BinaryenAtomicWait(BinaryenModuleRef module, BinaryenExpre
   return static_cast<Expression*>(ret);
 }
 BinaryenExpressionRef BinaryenAtomicWake(BinaryenModuleRef module, BinaryenExpressionRef ptr, BinaryenExpressionRef wakeCount) {
-  auto* ret = Builder(*((Module*)module)).makeAtomicWake((Expression*)ptr, (Expression*)wakeCount);
+  auto* ret = Builder(*((Module*)module)).makeAtomicWake((Expression*)ptr, (Expression*)wakeCount, 0);
 
   if (tracing) {
     auto id = noteExpression(ret);

src/ir/ExpressionAnalyzer.cpp

@@ -492,13 +492,15 @@ uint32_t ExpressionAnalyzer::hash(Expression* curr) {
         break;
       }
       case Expression::Id::AtomicWaitId: {
+        HASH(AtomicWait, offset);
         HASH(AtomicWait, expectedType);
         PUSH(AtomicWait, ptr);
         PUSH(AtomicWait, expected);
         PUSH(AtomicWait, timeout);
         break;
       }
       case Expression::Id::AtomicWakeId: {
+        HASH(AtomicWake, offset);
         PUSH(AtomicWake, ptr);
         PUSH(AtomicWake, wakeCount);
         break;

src/ir/ExpressionManipulator.cpp

@@ -116,10 +116,10 @@ Expression* flexibleCopy(Expression* original, Module& wasm, CustomCopier custom
                                        curr->type);
     }
     Expression* visitAtomicWait(AtomicWait* curr) {
-      return builder.makeAtomicWait(copy(curr->ptr), copy(curr->expected), copy(curr->timeout), curr->expectedType);
+      return builder.makeAtomicWait(copy(curr->ptr), copy(curr->expected), copy(curr->timeout), curr->expectedType, curr->offset);
     }
     Expression* visitAtomicWake(AtomicWake* curr) {
-      return builder.makeAtomicWake(copy(curr->ptr), copy(curr->wakeCount));
+      return builder.makeAtomicWake(copy(curr->ptr), copy(curr->wakeCount), curr->offset);
     }
     Expression* visitConst(Const *curr) {
       return builder.makeConst(curr->value);

src/passes/Print.cpp

@@ -411,6 +411,9 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
     o << '(' ;
     prepareColor(o);
     o << printWasmType(curr->expectedType) << ".wait";
+    if (curr->offset) {
+      o << " offset=" << curr->offset;
+    }
     restoreNormalColor(o);
     incIndent();
     printFullLine(curr->ptr);
@@ -420,6 +423,9 @@ struct PrintSExpression : public Visitor<PrintSExpression> {
   }
   void visitAtomicWake(AtomicWake* curr) {
     printOpening(o, "wake");
+    if (curr->offset) {
+      o << " offset=" << curr->offset;
+    }
     incIndent();
     printFullLine(curr->ptr);
     printFullLine(curr->wakeCount);

src/tools/fuzzing.h

@@ -174,6 +174,9 @@ class TranslateToFuzzReader {
   // Whether to emit atomics
   static const bool ATOMICS = true;
 
+  // Whether to emit atomic waits (which in single-threaded mode, may hang...)
+  static const bool ATOMIC_WAITS = false;
+
   // after we finish the input, we start going through it again, but xoring
   // so it's not identical
   int xorFactor = 0;
@@ -1242,16 +1245,16 @@ class TranslateToFuzzReader {
     if (!ATOMICS || (type != i32 && type != i64)) return makeTrivial(type);
     wasm.memory.shared = true;
     if (type == i32 && oneIn(2)) {
-      if (oneIn(2)) {
+      if (ATOMIC_WAITS && oneIn(2)) {
         auto* ptr = makePointer();
         auto expectedType = pick(i32, i64);
         auto* expected = make(expectedType);
         auto* timeout = make(i64);
-        return builder.makeAtomicWait(ptr, expected, timeout, expectedType);
+        return builder.makeAtomicWait(ptr, expected, timeout, expectedType, logify(get()));
       } else {
         auto* ptr = makePointer();
         auto* count = make(i32);
-        return builder.makeAtomicWake(ptr, count);
+        return builder.makeAtomicWake(ptr, count, logify(get()));
       }
     }
     Index bytes;

src/wasm-binary.h

@@ -959,7 +959,7 @@ class WasmBinaryBuilder {
   void visitSetLocal(SetLocal *curr, uint8_t code);
   void visitGetGlobal(GetGlobal *curr);
   void visitSetGlobal(SetGlobal *curr);
-  void readMemoryAccess(Address& alignment, size_t bytes, Address& offset);
+  void readMemoryAccess(Address& alignment, Address& offset);
   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);

src/wasm-builder.h

@@ -240,17 +240,19 @@ class Builder {
     load->isAtomic = true;
     return load;
   }
-  AtomicWait* makeAtomicWait(Expression* ptr, Expression* expected, Expression* timeout, WasmType expectedType) {
+  AtomicWait* makeAtomicWait(Expression* ptr, Expression* expected, Expression* timeout, WasmType expectedType, Address offset) {
     auto* wait = allocator.alloc<AtomicWait>();
+    wait->offset = offset;
     wait->ptr = ptr;
     wait->expected = expected;
     wait->timeout = timeout;
     wait->expectedType = expectedType;
     wait->finalize();
     return wait;
   }
-  AtomicWake* makeAtomicWake(Expression* ptr, Expression* wakeCount) {
+  AtomicWake* makeAtomicWake(Expression* ptr, Expression* wakeCount, Address offset) {
     auto* wake = allocator.alloc<AtomicWake>();
+    wake->offset = offset;
     wake->ptr = ptr;
     wake->wakeCount = wakeCount;
     wake->finalize();

src/wasm.h

@@ -474,6 +474,7 @@ class AtomicWait : public SpecificExpression<Expression::AtomicWaitId> {
   AtomicWait() = default;
   AtomicWait(MixedArena& allocator) : AtomicWait() {}
 
+  Address offset;
   Expression* ptr;
   Expression* expected;
   Expression* timeout;
@@ -487,6 +488,7 @@ class AtomicWake : public SpecificExpression<Expression::AtomicWakeId> {
   AtomicWake() = default;
   AtomicWake(MixedArena& allocator) : AtomicWake() {}
 
+  Address offset;
   Expression* ptr;
   Expression* wakeCount;
 

src/wasm/wasm-binary.cpp

@@ -954,7 +954,10 @@ void WasmBinaryWriter::visitStore(Store *curr) {
 void WasmBinaryWriter::visitAtomicRMW(AtomicRMW *curr) {
   if (debug) std::cerr << "zz node: AtomicRMW" << std::endl;
   recurse(curr->ptr);
+  // stop if the rest isn't reachable anyhow
+  if (curr->ptr->type == unreachable) return;
   recurse(curr->value);
+  if (curr->value->type == unreachable) return;
 
   if (curr->type == unreachable) {
     // don't even emit it; we don't know the right type
@@ -1005,8 +1008,12 @@ void WasmBinaryWriter::visitAtomicRMW(AtomicRMW *curr) {
 void WasmBinaryWriter::visitAtomicCmpxchg(AtomicCmpxchg *curr) {
   if (debug) std::cerr << "zz node: AtomicCmpxchg" << std::endl;
   recurse(curr->ptr);
+  // stop if the rest isn't reachable anyhow
+  if (curr->ptr->type == unreachable) return;
   recurse(curr->expected);
+  if (curr->expected->type == unreachable) return;
   recurse(curr->replacement);
+  if (curr->replacement->type == unreachable) return;
 
   if (curr->type == unreachable) {
     // don't even emit it; we don't know the right type
@@ -1041,23 +1048,39 @@ void WasmBinaryWriter::visitAtomicCmpxchg(AtomicCmpxchg *curr) {
 void WasmBinaryWriter::visitAtomicWait(AtomicWait *curr) {
   if (debug) std::cerr << "zz node: AtomicWait" << std::endl;
   recurse(curr->ptr);
+  // stop if the rest isn't reachable anyhow
+  if (curr->ptr->type == unreachable) return;
   recurse(curr->expected);
+  if (curr->expected->type == unreachable) return;
   recurse(curr->timeout);
+  if (curr->timeout->type == unreachable) return;
 
   o << int8_t(BinaryConsts::AtomicPrefix);
   switch (curr->expectedType) {
-    case i32: o << int8_t(BinaryConsts::I32AtomicWait); break;
-    case i64: o << int8_t(BinaryConsts::I64AtomicWait); break;
+    case i32: {
+      o << int8_t(BinaryConsts::I32AtomicWait);
+      emitMemoryAccess(4, 4, 0);
+      break;
+    }
+    case i64: {
+      o << int8_t(BinaryConsts::I64AtomicWait);
+      emitMemoryAccess(8, 8, 0);
+      break;
+    }
     default: WASM_UNREACHABLE();
   }
 }
 
 void WasmBinaryWriter::visitAtomicWake(AtomicWake *curr) {
   if (debug) std::cerr << "zz node: AtomicWake" << std::endl;
   recurse(curr->ptr);
+  // stop if the rest isn't reachable anyhow
+  if (curr->ptr->type == unreachable) return;
   recurse(curr->wakeCount);
+  if (curr->wakeCount->type == unreachable) return;
 
   o << int8_t(BinaryConsts::AtomicPrefix) << int8_t(BinaryConsts::AtomicWake);
+  emitMemoryAccess(4, 4, 0);
 }
 
 void WasmBinaryWriter::visitConst(Const *curr) {
@@ -2555,7 +2578,7 @@ void WasmBinaryBuilder::visitSetGlobal(SetGlobal *curr) {
   curr->finalize();
 }
 
-void WasmBinaryBuilder::readMemoryAccess(Address& alignment, size_t bytes, Address& offset) {
+void WasmBinaryBuilder::readMemoryAccess(Address& alignment, Address& offset) {
   auto rawAlignment = getU32LEB();
   if (rawAlignment > 4) throw ParseException("Alignment must be of a reasonable size");
   alignment = Pow2(rawAlignment);
@@ -2599,7 +2622,7 @@ bool WasmBinaryBuilder::maybeVisitLoad(Expression*& out, uint8_t code, bool isAt
   }
 
   curr->isAtomic = isAtomic;
-  readMemoryAccess(curr->align, curr->bytes, curr->offset);
+  readMemoryAccess(curr->align, curr->offset);
   curr->ptr = popNonVoidExpression();
   curr->finalize();
   out = curr;
@@ -2636,7 +2659,7 @@ bool WasmBinaryBuilder::maybeVisitStore(Expression*& out, uint8_t code, bool isA
 
   curr->isAtomic = isAtomic;
   if (debug) std::cerr << "zz node: Store" << std::endl;
-  readMemoryAccess(curr->align, curr->bytes, curr->offset);
+  readMemoryAccess(curr->align, curr->offset);
   curr->value = popNonVoidExpression();
   curr->ptr = popNonVoidExpression();
   curr->finalize();
@@ -2679,7 +2702,7 @@ bool WasmBinaryBuilder::maybeVisitAtomicRMW(Expression*& out, uint8_t code) {
 
   if (debug) std::cerr << "zz node: AtomicRMW" << std::endl;
   Address readAlign;
-  readMemoryAccess(readAlign, curr->bytes, curr->offset);
+  readMemoryAccess(readAlign, curr->offset);
   if (readAlign != curr->bytes) throw ParseException("Align of AtomicRMW must match size");
   curr->value = popNonVoidExpression();
   curr->ptr = popNonVoidExpression();
@@ -2710,7 +2733,7 @@ bool WasmBinaryBuilder::maybeVisitAtomicCmpxchg(Expression*& out, uint8_t code)
 
   if (debug) std::cerr << "zz node: AtomicCmpxchg" << std::endl;
   Address readAlign;
-  readMemoryAccess(readAlign, curr->bytes, curr->offset);
+  readMemoryAccess(readAlign, curr->offset);
   if (readAlign != curr->bytes) throw ParseException("Align of AtomicCpxchg must match size");
   curr->replacement = popNonVoidExpression();
   curr->expected = popNonVoidExpression();
@@ -2734,6 +2757,9 @@ bool WasmBinaryBuilder::maybeVisitAtomicWait(Expression*& out, uint8_t code) {
   curr->timeout = popNonVoidExpression();
   curr->expected = popNonVoidExpression();
   curr->ptr = popNonVoidExpression();
+  Address readAlign;
+  readMemoryAccess(readAlign, curr->offset);
+  if (readAlign != getWasmTypeSize(curr->expectedType)) throw ParseException("Align of AtomicWait must match size");
   curr->finalize();
   out = curr;
   return true;
@@ -2747,6 +2773,9 @@ bool WasmBinaryBuilder::maybeVisitAtomicWake(Expression*& out, uint8_t code) {
   curr->type = i32;
   curr->wakeCount = popNonVoidExpression();
   curr->ptr = popNonVoidExpression();
+  Address readAlign;
+  readMemoryAccess(readAlign, curr->offset);
+  if (readAlign != getWasmTypeSize(curr->type)) throw ParseException("Align of AtomicWake must match size");
   curr->finalize();
   out = curr;
   return true;