Allow i64 and ref eq control words in struct.wait

Author: stevenfontanella

src/ir/child-typer.h

@@ -1038,10 +1038,19 @@ template<typename Subtype> struct ChildTyper : OverriddenVisitor<Subtype> {
       }
       ht = curr->ref->type.getHeapType();
     }
+    const auto& fields = ht->getStruct().fields;
+    if (curr->index >= fields.size()) {
+      self().noteUnknown();
+      return;
+    }
 
     note(&curr->ref, Type(*ht, Nullable));
     note(&curr->waitqueue, Type(HeapTypes::sharedWaitqueue, Nullable));
-    note(&curr->expected, Type(Type::BasicType::i32));
+    auto expectedType = fields[curr->index].type;
+    if (expectedType.isRef()) {
+      expectedType = Type(HeapTypes::eq.getBasic(Shared), Nullable);
+    }
+    note(&curr->expected, expectedType);
     note(&curr->timeout, Type(Type::BasicType::i64));
   }
 

src/wasm-interpreter.h

@@ -2324,6 +2324,7 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
 
   Flow visitStructWait(StructWait* curr) {
     VISIT(ref, curr->ref)
+    VISIT(waitqueue, curr->waitqueue)
     VISIT(expected, curr->expected)
     VISIT(timeout, curr->timeout)
 
@@ -2335,8 +2336,11 @@ class ExpressionRunner : public OverriddenVisitor<SubType, Flow> {
     if (!data) {
       trap("null ref");
     }
+    if (!waitqueue.getSingleValue().getGCData()) {
+      trap("null ref");
+    }
     auto& field = data->values[curr->index];
-    if (field.geti32() != expected.getSingleValue().geti32()) {
+    if (field != expected.getSingleValue()) {
       return Literal(int32_t{1}); // not equal
     }
     // TODO: Add threads support. For now, report a host limit here, as there

src/wasm/wasm-ir-builder.cpp

@@ -2426,6 +2426,7 @@ Result<> IRBuilder::makeStructWait(HeapType type, Index index) {
   }
 
   StructWait curr(wasm.allocator);
+  curr.index = index;
   CHECK_ERR(ChildPopper{*this}.visitStructWait(&curr, type));
   CHECK_ERR(validateTypeAnnotation(type, curr.ref));
   push(builder.makeStructWait(

src/wasm/wasm-validator.cpp

@@ -600,6 +600,8 @@ struct FunctionValidator : public WalkerPass<PostWalker<FunctionValidator>> {
   bool shouldBeTrue(bool result, T curr, const char* text) {
     return info.shouldBeTrue(result, curr, text, getFunction());
   }
+
+  // Returns true if the assertion was met, i.e. returns !result.
   template<typename T>
   bool shouldBeFalse(bool result, T curr, const char* text) {
     return info.shouldBeFalse(result, curr, text, getFunction());
@@ -3644,6 +3646,10 @@ void FunctionValidator::visitStructCmpxchg(StructCmpxchg* curr) {
 }
 
 void FunctionValidator::visitStructWait(StructWait* curr) {
+  // In IRBuilder, we check that the struct ref matches the type immediate.
+  // We can't check this here because we've already discarded the type immediate
+  // at this point. All other validations are here.
+
   shouldBeTrue(
     !getModule() || getModule()->features.hasSharedEverything(),
     curr,
@@ -3653,20 +3659,44 @@ void FunctionValidator::visitStructWait(StructWait* curr) {
                   Type(HeapTypes::sharedWaitqueue, Nullable),
                   curr,
                   "struct.wait waitqueue must be a shared waitqueue reference");
-  shouldBeEqual(curr->expected->type,
-                Type(Type::BasicType::i32),
-                curr,
-                "struct.wait expected must be an i32");
   shouldBeEqual(curr->timeout->type,
                 Type(Type::BasicType::i64),
                 curr,
                 "struct.wait timeout must be an i64");
 
-  // Checks to the ref argument's type are done in IRBuilder where we have the
-  // type annotation immediate available. We check that
-  // * The reference arg is a subtype of the type immediate
-  // * The index immediate is a valid field index of the type immediate (and
-  // thus valid for the reference's type too)
+  if (curr->ref->type == Type::unreachable || curr->ref->type.isNull()) {
+    return;
+  }
+  if (!shouldBeTrue(curr->ref->type.isStruct(),
+                    curr->ref,
+                    "struct.wait ref must be a struct")) {
+    return;
+  }
+  const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
+  if (!shouldBeTrue(
+        curr->index < fields.size(), curr, "out of bounds struct.wait field")) {
+    return;
+  }
+  auto& field = fields[curr->index];
+  if (!shouldBeFalse(
+        field.isPacked(), curr, "struct.wait field must not be packed")) {
+    return;
+  }
+
+  if (
+    !shouldBeTrue(
+      field.type == Type::i32 || field.type == Type::i64 ||
+        Type::isSubType(field.type,
+                        Type(HeapTypes::eq.getBasic(Shared), Nullable)),
+      curr,
+      R"(struct.wait control word field must be i32, i64 or a subtype of (ref null (shared eq)))")) {
+    return;
+  }
+
+  shouldBeSubType(curr->expected->type,
+                  field.type,
+                  curr,
+                  "struct.wait expected value must match the field immediate");
 }
 
 void FunctionValidator::visitWaitqueueNew(WaitqueueNew* curr) {

test/lit/validation/waitqueue.wast

@@ -1,9 +1,15 @@
 ;; RUN: not wasm-opt --enable-reference-types --enable-gc %s 2>&1 | filecheck %s
+
+;; Tests feature-related validations.
+;; Other validations are in the spec test spec/waitqueue.wast.
 (module
   (type $struct (struct (field i32)))
   ;; CHECK:      waitqueue.new requires shared-everything [--enable-shared-everything]
-  (func
+  (func $new
     (drop (waitqueue.new))
   )
+  ;; CHECK:      struct.wait requires shared-everything [--enable-shared-everything]
+  (func $wait-no-feature (param $ref (ref $struct)) (param $wq (ref null (shared waitqueue)))
+    (drop (struct.wait $struct 0 (local.get $ref) (local.get $wq) (i32.const 0) (i64.const 0)))
+  )
 )
-

test/spec/waitqueue.wast

@@ -14,7 +14,7 @@
     (func (param $expected i32) (param $timeout i64) (result i32)
       (struct.wait $t 2 (ref.null $t) (global.get $wq) (local.get $expected) (local.get $timeout))
     )
-  ) "struct index out of bounds"
+  ) "out of bounds struct.wait field"
 )
 
 (assert_invalid
@@ -25,7 +25,29 @@
     (func (param $expected i32) (param $timeout i64) (result i32)
       (struct.wait $t 0 (global.get $g) (global.get $wq) (i64.const 0) (local.get $timeout))
     )
-  ) "struct.wait expected must be an i32"
+  ) "struct.wait expected value must match the field immediate"
+)
+
+(assert_invalid
+  (module
+    (type $t (shared (struct (field f32))))
+    (global $g (ref $t) (struct.new $t (f32.const 0)))
+    (global $wq (ref (shared waitqueue)) (waitqueue.new))
+    (func (param $expected f32) (param $timeout i64) (result i32)
+      (struct.wait $t 0 (global.get $g) (global.get $wq) (local.get $expected) (local.get $timeout))
+    )
+  ) "struct.wait control word field must be i32, i64 or a subtype of (ref null (shared eq))"
+)
+
+(assert_invalid
+  (module
+    (type $t (shared (struct (field i8))))
+    (global $g (ref $t) (struct.new $t (i32.const 0)))
+    (global $wq (ref (shared waitqueue)) (waitqueue.new))
+    (func (param $expected i32) (param $timeout i64) (result i32)
+      (struct.wait $t 0 (global.get $g) (global.get $wq) (local.get $expected) (local.get $timeout))
+    )
+  ) "struct.wait field must not be packed"
 )
 
 (assert_invalid
@@ -41,8 +63,6 @@
 
 (assert_invalid
   (module
-    (type $t (shared (struct (field i32))))
-    (global $wq (ref (shared waitqueue)) (waitqueue.new))
     (func (param $count i32) (result i32)
       (waitqueue.notify (ref.null waitqueue) (local.get $count))
     )
@@ -74,6 +94,7 @@
   )
 )
 
+;; i32 control word
 (module
   (type $t (shared (struct (field (mut i32)))))
 
@@ -119,6 +140,78 @@
 (assert_trap (invoke "struct.wait" (i32.const 0) (i64.const 0)) "null ref")
 (assert_trap (invoke "waitqueue.notify" (i32.const 0)) "null ref")
 
+;; i64 control word
+(module
+  (type $t (shared (struct (field (mut i64)))))
+
+  (global $g (mut (ref null $t)) (struct.new $t (i64.const 0)))
+  (global $wq (mut (ref (shared waitqueue))) (waitqueue.new))
+
+  (func (export "struct.wait") (param $expected i64) (param $timeout i64) (result i32)
+    (struct.wait $t 0 (global.get $g) (global.get $wq) (local.get $expected) (local.get $timeout))
+  )
+
+  (func (export "struct.set") (param $val i64)
+    (struct.set $t 0 (global.get $g) (local.get $val))
+  )
+
+  (func (export "struct.get") (result i64)
+    (struct.get $t 0 (global.get $g))
+  )
+)
+
+(invoke "struct.set" (i64.const 42))
+(assert_return (invoke "struct.get") (i64.const 42))
+(assert_return (invoke "struct.wait" (i64.const 0) (i64.const 100)) (i32.const 1))
+(assert_return (invoke "struct.wait" (i64.const 42) (i64.const 0)) (i32.const 2))
+
+;; (ref null (shared eq)) control word
+(module
+  (type $control (shared (struct)))
+
+  (type $t (shared (struct
+    (field (mut (ref null (shared eq))))
+  )))
+
+  (global $control1 (ref $control) (struct.new $control))
+  (global $control2 (ref $control) (struct.new $control))
+
+  (global $g (mut (ref null $t)) (struct.new $t
+    (global.get $control1)
+  ))
+
+  (global $wq (mut (ref null (shared waitqueue))) (waitqueue.new))
+
+  (func (export "wait_control1") (result i32)
+    (struct.wait $t 0 (global.get $g) (global.get $wq) (global.get $control1) (i64.const 0))
+  )
+
+  (func (export "wait_control2") (result i32)
+    (struct.wait $t 0 (global.get $g) (global.get $wq) (global.get $control2) (i64.const 0))
+  )
+
+  (func (export "wait_null") (result i32)
+    (struct.wait $t 0 (global.get $g) (global.get $wq) (ref.null (shared eq)) (i64.const 0))
+  )
+
+  (func (export "set_control_to_null")
+    (struct.set $t 0 (global.get $g) (ref.null (shared eq)))
+  )
+)
+
+;; $control1 is the control word, wait 0ns and return 2.
+(assert_return (invoke "wait_control1") (i32.const 2))
+;; $control2 is not the control work, don't wait and return 1.
+(assert_return (invoke "wait_control2") (i32.const 1))
+;; ditto for null.
+(assert_return (invoke "wait_null") (i32.const 1))
+
+(invoke "set_control_to_null")
+
+;; null is now the control word.
+(assert_return (invoke "wait_null") (i32.const 2))
+(assert_return (invoke "wait_control1") (i32.const 1))
+
 ;; Binary format test for waitqueue and nowaitqueue.
 (module binary
   "\00asm\01\00\00\00" ;; Wasm header