Change how async stream/future results are encoded (#1273)

Updating to WebAssembly/component-model#498

Author: alexcrichton

Cargo.lock

@@ -287,6 +287,41 @@ const STATUS_STARTING: u32 = 1;
 const STATUS_STARTED: u32 = 2;
 const STATUS_RETURNED: u32 = 3;
 
+const BLOCKED: u32 = 0xffff_ffff;
+const COMPLETED: u32 = 0x0;
+const CLOSED: u32 = 0x1;
+const CANCELLED: u32 = 0x2;
+
+/// Return code of stream/future operations.
+#[derive(PartialEq, Debug)]
+enum ReturnCode {
+    /// The operation is blocked and has not completed.
+    Blocked,
+    /// The operation completed with the specified number of items.
+    Completed(u32),
+    /// The other end is closed, but before that the specified number of items
+    /// were transferred.
+    Closed(u32),
+    /// The operation was cancelled, but before that the specified number of
+    /// items were transferred.
+    Cancelled(u32),
+}
+
+impl ReturnCode {
+    fn decode(val: u32) -> ReturnCode {
+        if val == BLOCKED {
+            return ReturnCode::Blocked;
+        }
+        let amt = val >> 4;
+        match val & 0xf {
+            COMPLETED => ReturnCode::Completed(amt),
+            CLOSED => ReturnCode::Closed(amt),
+            CANCELLED => ReturnCode::Cancelled(amt),
+            _ => panic!("unknown return code {val:#x}"),
+        }
+    }
+}
+
 /// Starts execution of the `task` provided, an asynchronous computation.
 ///
 /// This is used for async-lifted exports at their definition site. The
@@ -313,13 +348,6 @@ pub fn start_task(task: impl Future<Output = ()> + 'static) -> u32 {
     }
 }
 
-/// stream/future read/write results defined by the Component Model ABI.
-mod results {
-    pub const BLOCKED: u32 = 0xffff_ffff;
-    pub const CLOSED: u32 = 0x8000_0000;
-    pub const CANCELED: u32 = 0;
-}
-
 /// Handle a progress notification from the host regarding either a call to an
 /// async-lowered import or a stream/future read/write operation.
 ///

crates/guest-rust/rt/src/async_support.rs

@@ -88,8 +88,8 @@
 //! changing its layout. Currently this is par-for-the-course with bindings.
 
 use {
-    super::results,
     super::waitable::{WaitableOp, WaitableOperation},
+    super::ReturnCode,
     crate::Cleanup,
     std::{
         alloc::Layout,
@@ -306,8 +306,8 @@ where
             .as_ref()
             .map(|c| c.ptr.as_ptr())
             .unwrap_or(ptr::null_mut());
-        match code {
-            results::BLOCKED => Err((writer, cleanup)),
+        match ReturnCode::decode(code) {
+            ReturnCode::Blocked => Err((writer, cleanup)),
 
             // The other end has closed its end.
             //
@@ -316,11 +316,11 @@ where
             // instance of `T` which takes ownership of pointers and resources
             // and such. The allocation of `ptr` is then cleaned up naturally
             // when `cleanup` goes out of scope.
-            results::CLOSED | results::CANCELED => {
+            c @ ReturnCode::Closed(0) | c @ ReturnCode::Cancelled(0) => {
                 // SAFETY: we're the ones managing `ptr` so we know it's safe to
                 // pass here.
                 let value = unsafe { (writer.vtable.lift)(ptr) };
-                let status = if code == results::CLOSED {
+                let status = if c == ReturnCode::Closed(0) {
                     WriteComplete::Closed(value)
                 } else {
                     WriteComplete::Cancelled(value)
@@ -337,7 +337,7 @@ where
             //
             // Afterwards the `cleanup` itself is naturally dropped and cleaned
             // up.
-            1 => {
+            ReturnCode::Completed(1) | ReturnCode::Closed(1) | ReturnCode::Cancelled(1) => {
                 // SAFETY: we're the ones managing `ptr` so we know it's safe to
                 // pass here.
                 unsafe {
@@ -346,7 +346,7 @@ where
                 Ok((WriteComplete::Written, writer))
             }
 
-            other => unreachable!("unexpected code {other:#x}"),
+            other => unreachable!("unexpected code {other:?}"),
         }
     }
 
@@ -566,22 +566,22 @@ where
         (reader, cleanup): Self::InProgress,
         code: u32,
     ) -> Result<Self::Result, Self::InProgress> {
-        match code {
-            results::BLOCKED => Err((reader, cleanup)),
+        match ReturnCode::decode(code) {
+            ReturnCode::Blocked => Err((reader, cleanup)),
 
             // The read didn't complete, so `cleanup` is still uninitialized, so
             // let it fall out of scope.
-            results::CLOSED => Ok((ReadComplete::Closed, reader)),
+            ReturnCode::Closed(0) => Ok((ReadComplete::Closed, reader)),
 
             // Like `in_progress_closed` the read operation has finished but
             // without a value, so let `cleanup` fall out of scope to clean up
             // its allocation.
-            results::CANCELED => Ok((ReadComplete::Cancelled, reader)),
+            ReturnCode::Cancelled(0) => Ok((ReadComplete::Cancelled, reader)),
 
             // The read has completed, so lift the value from the stored memory and
             // `cleanup` naturally falls out of scope after transferring ownership of
             // everything to the returned `value`.
-            1 => {
+            ReturnCode::Completed(1) | ReturnCode::Closed(1) | ReturnCode::Cancelled(1) => {
                 let ptr = cleanup
                     .as_ref()
                     .map(|c| c.ptr.as_ptr())
@@ -593,7 +593,7 @@ where
                 Ok((ReadComplete::Value(value), reader))
             }
 
-            other => panic!("unexpected code {other:#x}"),
+            other => panic!("unexpected code {other:?}"),
         }
     }
 

crates/guest-rust/rt/src/async_support/future_support.rs

@@ -2,7 +2,7 @@
 //! module documentation in `future_support.rs`.
 
 use crate::async_support::waitable::{WaitableOp, WaitableOperation};
-use crate::async_support::{results, AbiBuffer};
+use crate::async_support::{AbiBuffer, ReturnCode};
 use {
     crate::Cleanup,
     std::{
@@ -258,11 +258,11 @@ where
         (writer, mut buf): Self::InProgress,
         code: u32,
     ) -> Result<Self::Result, Self::InProgress> {
-        match code {
-            results::BLOCKED => Err((writer, buf)),
-            results::CLOSED => Ok((StreamResult::Closed, buf)),
-            results::CANCELED => Ok((StreamResult::Cancelled, buf)),
-            amt => {
+        match ReturnCode::decode(code) {
+            ReturnCode::Blocked => Err((writer, buf)),
+            ReturnCode::Closed(0) => Ok((StreamResult::Closed, buf)),
+            ReturnCode::Cancelled(0) => Ok((StreamResult::Cancelled, buf)),
+            ReturnCode::Completed(amt) | ReturnCode::Closed(amt) | ReturnCode::Cancelled(amt) => {
                 let amt = amt.try_into().unwrap();
                 buf.advance(amt);
                 Ok((StreamResult::Complete(amt), buf))
@@ -480,20 +480,20 @@ where
         (reader, mut buf, cleanup): Self::InProgress,
         code: u32,
     ) -> Result<Self::Result, Self::InProgress> {
-        match code {
-            results::BLOCKED => Err((reader, buf, cleanup)),
+        match ReturnCode::decode(code) {
+            ReturnCode::Blocked => Err((reader, buf, cleanup)),
 
             // Note that the `cleanup`, if any, is discarded here.
-            results::CLOSED => Ok((StreamResult::Closed, buf)),
+            ReturnCode::Closed(0) => Ok((StreamResult::Closed, buf)),
 
             // When an in-progress read is successfully cancelled then the
             // allocation that was being read into, if any, is just discarded.
             //
             // TODO: should maybe thread this around like `AbiBuffer` to cache
             // the read allocation?
-            results::CANCELED => Ok((StreamResult::Cancelled, buf)),
+            ReturnCode::Cancelled(0) => Ok((StreamResult::Cancelled, buf)),
 
-            amt => {
+            ReturnCode::Completed(amt) | ReturnCode::Closed(amt) | ReturnCode::Cancelled(amt) => {
                 let amt = usize::try_from(amt).unwrap();
                 let cur_len = buf.len();
                 assert!(amt <= buf.capacity() - cur_len);