remove obsolete post-return functions and fields

Now that post-return calls are handled internally without requiring explicit
action by the embedder, we can avoid unnecessary bookkeeping.

Author: dicej

crates/environ/src/component/vmcomponent_offsets.rs

@@ -29,10 +29,6 @@ pub const VMCOMPONENT_MAGIC: u32 = u32::from_le_bytes(*b"comp");
 /// canonical ABI flag `may_leave`
 pub const FLAG_MAY_LEAVE: i32 = 1 << 0;
 
-/// Flag for the `VMComponentContext::flags` field which is set whenever a
-/// function is called to indicate that `post_return` must be called next.
-pub const FLAG_NEEDS_POST_RETURN: i32 = 1 << 2;
-
 /// Runtime offsets within a `VMComponentContext` for a specific component.
 #[derive(Debug, Clone, Copy)]
 pub struct VMComponentOffsets<P> {

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -1518,14 +1518,14 @@ impl StoreOpaque {
     /// - `self` has not been poisoned due to a trap.
     pub(crate) fn may_enter(&mut self, instance: RuntimeInstance) -> bool {
         if !self.concurrency_support() {
-            return self.may_enter_at_all(instance);
+            return !self.trapped();
         }
         let state = self.concurrent_state_mut();
         if let Some(caller) = state.guest_thread {
             instance != state.get_mut(caller.task).unwrap().instance
                 && self.may_enter_from_caller(caller.task, instance)
         } else {
-            self.may_enter_at_all(instance)
+            !self.trapped()
         }
     }
 
@@ -1543,7 +1543,7 @@ impl StoreOpaque {
         mut guest_task: TableId<GuestTask>,
         instance: RuntimeInstance,
     ) -> bool {
-        self.may_enter_at_all(instance) && {
+        !self.trapped() && {
             let state = self.concurrent_state_mut();
             let guest_instance = instance.instance;
             loop {

crates/wasmtime/src/runtime/component/concurrent_disabled.rs

@@ -172,7 +172,7 @@ impl StoreOpaque {
         Ok(())
     }
 
-    pub(crate) fn may_enter(&mut self, instance: RuntimeInstance) -> bool {
-        self.may_enter_at_all(instance)
+    pub(crate) fn may_enter(&mut self, _instance: RuntimeInstance) -> bool {
+        !self.trapped()
     }
 }

crates/wasmtime/src/runtime/component/func.rs

@@ -439,7 +439,7 @@ impl Func {
             MAX_FLAT_PARAMS,
             false,
             move |func, store, params_out| {
-                func.with_lower_context(store, true, |cx, ty| {
+                func.with_lower_context(store, |cx, ty| {
                     Self::lower_args(cx, &params, ty, params_out)
                 })
             },
@@ -490,7 +490,7 @@ impl Func {
         //   safe in Rust, however, due to `ValRaw` being a `union`. The
         //   contents should dynamically not be read due to the type of the
         //   function used here matching the actual lift.
-        unsafe {
+        let (_, post_return_arg) = unsafe {
             self.call_raw(
                 store.as_context_mut(),
                 |cx, ty, dst: &mut MaybeUninit<[MaybeUninit<ValRaw>; MAX_FLAT_PARAMS]>| {
@@ -510,9 +510,9 @@ impl Func {
                     Ok(())
                 },
             )?
-        }
+        };
 
-        self.post_return_impl(store)
+        self.post_return_impl(store, post_return_arg)
     }
 
     pub(crate) fn lifted_core_func(&self, store: &mut StoreOpaque) -> NonNull<VMFuncRef> {
@@ -587,7 +587,7 @@ impl Func {
             &mut MaybeUninit<LowerParams>,
         ) -> Result<()>,
         lift: impl FnOnce(&mut LiftContext<'_>, InterfaceType, &LowerReturn) -> Result<Return>,
-    ) -> Result<Return>
+    ) -> Result<(Return, ValRaw)>
     where
         LowerParams: Copy,
         LowerReturn: Copy,
@@ -631,7 +631,7 @@ impl Func {
         assert!(mem::align_of_val(map_maybe_uninit!(space.params)) == val_align);
         assert!(mem::align_of_val(map_maybe_uninit!(space.ret)) == val_align);
 
-        self.with_lower_context(store.as_context_mut(), false, |cx, ty| {
+        self.with_lower_context(store.as_context_mut(), |cx, ty| {
             cx.enter_call();
             lower(cx, ty, map_maybe_uninit!(space.params))
         })?;
@@ -671,23 +671,20 @@ impl Func {
         // is currently required to be 0 or 1 values according to the
         // canonical ABI, is saved within the `Store`'s `FuncData`. This'll
         // later get used in post-return.
-        // flags.set_needs_post_return(true);
         let val = self.with_lift_context(store.0, |cx, ty| lift(cx, ty, ret))?;
 
         // SAFETY: it's a contract of this function that `LowerReturn` is an
         // appropriate representation of the result of this function.
         let ret_slice = unsafe { storage_as_slice(ret) };
 
-        self.instance.id().get_mut(store.0).post_return_arg_set(
-            self.index,
+        Ok((
+            val,
             match ret_slice.len() {
                 0 => ValRaw::i32(0),
                 1 => ret_slice[0],
                 _ => unreachable!(),
             },
-        );
-
-        return Ok(val);
+        ))
     }
 
     #[doc(hidden)]
@@ -703,7 +700,7 @@ impl Func {
         Ok(())
     }
 
-    pub(crate) fn post_return_impl(&self, mut store: impl AsContextMut) -> Result<()> {
+    pub(crate) fn post_return_impl(&self, mut store: impl AsContextMut, arg: ValRaw) -> Result<()> {
         let mut store = store.as_context_mut();
 
         let index = self.index;
@@ -712,33 +709,9 @@ impl Func {
         let (_ty, _def, options) = component.export_lifted_function(index);
         let post_return = self.post_return_core_func(store.0);
         let flags = vminstance.instance_flags(component.env_component().options[options].instance);
-        let mut instance = self.instance.id().get_mut(store.0);
-        let post_return_arg = instance.as_mut().post_return_arg_take(index);
 
         unsafe {
-            // First assert that the instance is in a "needs post return" state.
-            // This will ensure that the previous action on the instance was a
-            // function call above. This flag is only set after a component
-            // function returns so this also can't be called (as expected)
-            // during a host import for example.
-            //
-            // Note, though, that this assert is not sufficient because it just
-            // means some function on this instance needs its post-return
-            // called. We need a precise post-return for a particular function
-            // which is the second assert here (the `.expect`). That will assert
-            // that this function itself needs to have its post-return called.
-            //
-            // The theory at least is that these two asserts ensure component
-            // model semantics are upheld where the host properly calls
-            // `post_return` on the right function despite the call being a
-            // separate step in the API.
-            assert!(
-                flags.needs_post_return(),
-                "post_return can only be called after a function has previously been called",
-            );
-            let post_return_arg = post_return_arg.expect("calling post_return on wrong function");
-
-            call_post_return(&mut store, post_return, post_return_arg, flags)?;
+            call_post_return(&mut store, post_return, arg, flags)?;
 
             let (calls, host_table, _, instance) = store
                 .0
@@ -866,11 +839,9 @@ impl Func {
     fn with_lower_context<T>(
         self,
         mut store: StoreContextMut<T>,
-        call_post_return_automatically: bool,
         lower: impl FnOnce(&mut LowerContext<T>, InterfaceType) -> Result<()>,
     ) -> Result<()> {
-        let (options_idx, mut flags, ty, options) = self.abi_info(store.0);
-        let async_ = options.async_;
+        let (options_idx, mut flags, ty, _) = self.abi_info(store.0);
 
         // Perform the actual lowering, where while this is running the
         // component is forbidden from calling imports.
@@ -882,17 +853,7 @@ impl Func {
         let param_ty = InterfaceType::Tuple(cx.types[ty].params);
         let result = lower(&mut cx, param_ty);
         unsafe { flags.set_may_leave(true) };
-        result?;
-
-        // If needed, flag a post-return call being required as we're about to
-        // enter wasm and afterwards need a post-return.
-        unsafe {
-            if !(call_post_return_automatically && async_) {
-                flags.set_needs_post_return(true);
-            }
-        }
-
-        Ok(())
+        result
     }
 
     /// Creates a `LiftContext` using the configuration values with this lifted
@@ -951,11 +912,6 @@ pub(crate) unsafe fn call_post_return(
     mut flags: InstanceFlags,
 ) -> Result<()> {
     unsafe {
-        // Unset the "needs post return" flag now that post-return is being
-        // processed. This will cause future invocations of this method to
-        // panic, even if the function call below traps.
-        flags.set_needs_post_return(false);
-
         // Post return functions are forbidden from calling imports or
         // intrinsics.
         flags.set_may_leave(false);
@@ -966,7 +922,7 @@ pub(crate) unsafe fn call_post_return(
             crate::Func::call_unchecked_raw(
                 &mut store.as_context_mut(),
                 func,
-                std::slice::from_ref(&arg).into(),
+                core::slice::from_ref(&arg).into(),
             )?;
         }
 

crates/wasmtime/src/runtime/component/func/typed.rs

@@ -380,7 +380,7 @@ where
             param_count,
             host_future_present,
             move |func, store, params_out| {
-                func.with_lower_context(store, true, |cx, ty| lower(cx, ty, params_out))
+                func.with_lower_context(store, |cx, ty| lower(cx, ty, params_out))
             },
             move |func, store, results| {
                 let result = if Return::flatten_count() <= max_results {
@@ -439,7 +439,7 @@ where
         // safety requirements of `Lift` and `Lower` on `Params` and `Return` in
         // combination with checking the various possible branches here and
         // dispatching to appropriately typed functions.
-        let result = unsafe {
+        let (result, post_return_arg) = unsafe {
             // This type is used as `LowerParams` for `call_raw` which is either
             // `Params::Lower` or `ValRaw` representing it's either on the stack
             // or it's on the heap. This allocates 1 extra `ValRaw` on the stack
@@ -473,7 +473,7 @@ where
             }
         }?;
 
-        self.func.post_return_impl(store)?;
+        self.func.post_return_impl(store, post_return_arg)?;
 
         Ok(result)
     }

crates/wasmtime/src/runtime/component/store.rs

@@ -34,23 +34,6 @@ impl StoreOpaque {
     pub(crate) fn set_trapped(&mut self) {
         self.store_data_mut().components.trapped = true;
     }
-
-    /// Returns `false` if the specified instance may not be entered, regardless
-    /// of what's on a task's call stack.
-    ///
-    /// If this returns `true`, the instance may be entered as long as it isn't
-    /// on the task's call stack, if applicable.
-    pub(crate) fn may_enter_at_all(&mut self, instance: RuntimeInstance) -> bool {
-        if self.trapped() {
-            return false;
-        }
-
-        let flags = self
-            .component_instance(instance.instance)
-            .instance_flags(instance.index);
-
-        unsafe { !flags.needs_post_return() }
-    }
 }
 
 impl StoreData {

crates/wasmtime/src/runtime/vm/component.rs

@@ -148,12 +148,6 @@ pub struct ComponentInstance {
     /// Self-pointer back to `Store<T>` and its functions.
     store: VMStoreRawPtr,
 
-    /// Cached ABI return value from the last-invoked function call along with
-    /// the function index that was invoked.
-    ///
-    /// Used in `post_return_arg_set` and `post_return_arg_take` below.
-    post_return_arg: Option<(ExportIndex, ValRaw)>,
-
     /// Required by `InstanceLayout`, also required to be the last field (with
     /// repr(C))
     vmctx: OwnedVMContext<VMComponentContext>,
@@ -341,7 +335,6 @@ impl ComponentInstance {
             resource_types,
             imports: imports.clone(),
             store: VMStoreRawPtr(store),
-            post_return_arg: None,
             vmctx: OwnedVMContext::new(),
         })?;
         unsafe {
@@ -954,50 +947,6 @@ impl ComponentInstance {
         }
     }
 
-    /// Sets the cached argument for the canonical ABI option `post-return` to
-    /// the `arg` specified.
-    ///
-    /// This function is used in conjunction with function calls to record,
-    /// after a function call completes, the optional ABI return value. This
-    /// return value is cached within this instance for future use when the
-    /// `post_return` Rust-API-level function is invoked.
-    ///
-    /// Note that `index` here is the index of the export that was just
-    /// invoked, and this is used to ensure that `post_return` is called on the
-    /// same function afterwards. This restriction technically isn't necessary
-    /// though and may be one we want to lift in the future.
-    ///
-    /// # Panics
-    ///
-    /// This function will panic if `post_return_arg` is already set to `Some`.
-    pub fn post_return_arg_set(self: Pin<&mut Self>, index: ExportIndex, arg: ValRaw) {
-        assert!(self.post_return_arg.is_none());
-        *self.post_return_arg_mut() = Some((index, arg));
-    }
-
-    /// Re-acquires the value originally saved via `post_return_arg_set`.
-    ///
-    /// This function will take a function `index` that's having its
-    /// `post_return` function called. If an argument was previously stored and
-    /// `index` matches the index that was stored then `Some(arg)` is returned.
-    /// Otherwise `None` is returned.
-    pub fn post_return_arg_take(self: Pin<&mut Self>, index: ExportIndex) -> Option<ValRaw> {
-        let post_return_arg = self.post_return_arg_mut();
-        let (expected_index, arg) = post_return_arg.take()?;
-        if index != expected_index {
-            *post_return_arg = Some((expected_index, arg));
-            None
-        } else {
-            Some(arg)
-        }
-    }
-
-    fn post_return_arg_mut(self: Pin<&mut Self>) -> &mut Option<(ExportIndex, ValRaw)> {
-        // SAFETY: we've chosen the `Pin` guarantee of `Self` to not apply to
-        // the map returned.
-        unsafe { &mut self.get_unchecked_mut().post_return_arg }
-    }
-
     pub(crate) fn check_may_leave(
         &self,
         instance: RuntimeComponentInstanceIndex,
@@ -1124,22 +1073,6 @@ impl InstanceFlags {
         }
     }
 
-    #[inline]
-    pub unsafe fn needs_post_return(&self) -> bool {
-        unsafe { *self.as_raw().as_ref().as_i32() & FLAG_NEEDS_POST_RETURN != 0 }
-    }
-
-    #[inline]
-    pub unsafe fn set_needs_post_return(&mut self, val: bool) {
-        unsafe {
-            if val {
-                *self.as_raw().as_mut().as_i32_mut() |= FLAG_NEEDS_POST_RETURN;
-            } else {
-                *self.as_raw().as_mut().as_i32_mut() &= !FLAG_NEEDS_POST_RETURN;
-            }
-        }
-    }
-
     #[inline]
     pub fn as_raw(&self) -> NonNull<VMGlobalDefinition> {
         self.0.as_non_null()