More fine-grained tracking of continuations' states (#228)

This PR is part of an effort to implement more efficient effect
forwarding.

This particular PR makes our tracking of the state of continuations more
fine-grained.
Currently, the `State` enum only allows us to distinguish three cases:
1. A continuation was allocated, but never invoked (`State::Allocated`).
2. A continuation has been invoked at some point, and not returned, yet
(`State::Invoked`).
3. A continuation has returned  (`State::Returned`).

This PR replaces `State::Invoked` with three new states, indicating more
closely what's going on:
- `State::Running` indicates that the continuation is the one currently
executing code
- `State::Parent` indicates that a continuation is suspended because it
`resumed` another one.
- `State::Suspended` indicates that a continuation was suspended
directly (i.e., by calling `suspend`, or in the future `switch`)

This PR incurrs a few additional stores per stack switch, for example
because on `resume` we now need to update the state of the resumee *and*
its new parent. On our usual benchmarks (in particular, `state`), this
has no measurable effect. However, the two benchmarks in the `micro`
subfolder see a 5-10% regression. Interestingly, this doesn't seem to be
caused by the instructions responsible for the additional state updates
(I experimented with *not* updating the parent states), but just some
result of instructions being scheduled slightly differently after
applying this PR. In any case, this regression seems acceptable anyway.

This PR also extends the information we track about the main stack to
carry a `State`. Otherwise we would need to check whenever we want to
update the `State` of the parent if the parent is the main stack or
another continuation.

crates/continuations/src/lib.rs

@@ -71,6 +71,27 @@ pub struct StackLimits {
  pub last_wasm_entry_sp: usize,
 }
 
+/// This type represents "common" information that we need to save both for the
+/// main stack and each continuation.
+#[repr(C)]
+#[derive(Debug, Clone)]
+pub struct CommonStackInformation {
+ pub limits: StackLimits,
+ /// For the main stack, this field must only have one of the following values:
+ /// - Running
+ /// - Parent
+ pub state: State,
+}
+
+impl CommonStackInformation {
+ pub fn running_default() -> Self {
+ Self {
+ limits: StackLimits::default(),
+ state: State::Running,
+ }
+ }
+}
+
 impl StackLimits {
  pub fn with_stack_limit(stack_limit: usize) -> Self {
  Self {
@@ -171,12 +192,17 @@ pub const STACK_CHAIN_CONTINUATION_DISCRIMINANT: usize = 2;
 pub enum State {
  /// The `VMContRef` has been created, but `resume` has never been
  /// called on it. During this stage, we may add arguments using `cont.bind`.
- Allocated,
- /// `resume` has been invoked at least once on the `VMContRef`,
- /// meaning that the function passed to `cont.new` has started executing.
- /// Note that this state does not indicate whether the execution of this
- /// function is currently suspended or not.
- Invoked,
+ Fresh,
+ /// The continuation is running, meaning that it is the one currently
+ /// executing code.
+ Running,
+ /// The continuation is suspended because it executed a resume instruction
+ /// that has not finished yet. In other words, it became the parent of
+ /// another continuation (which may itself be `Running`, a `Parent`, or
+ /// `Suspended`).
+ Parent,
+ /// The continuation was suspended by a `suspend` instruction.
+ Suspended,
  /// The function originally passed to `cont.new` has returned normally.
  /// Note that there is no guarantee that a VMContRef will ever
  /// reach this status, as it may stay suspended until being dropped.
@@ -214,22 +240,21 @@ pub mod offsets {
  /// Offsets of fields in `wasmtime_runtime::continuation::VMContRef`.
  /// We uses tests there to ensure these values are correct.
  pub mod vm_cont_ref {
- use crate::Payloads;
+ use crate::{CommonStackInformation, Payloads};
 
- /// Offset of `limits` field
- pub const LIMITS: usize = 0;
+ /// Offset of `common_stack_information` field
+ pub const COMMON_STACK_INFORMATION: usize = 0;
  /// Offset of `parent_chain` field
- pub const PARENT_CHAIN: usize = LIMITS + 4 * core::mem::size_of::<usize>();
+ pub const PARENT_CHAIN: usize =
+ COMMON_STACK_INFORMATION + core::mem::size_of::<CommonStackInformation>();
  /// Offset of `stack` field
  pub const STACK: usize = PARENT_CHAIN + 2 * core::mem::size_of::<usize>();
  /// Offset of `args` field
  pub const ARGS: usize = STACK + super::FIBER_STACK_SIZE;
  /// Offset of `tag_return_values` field
  pub const TAG_RETURN_VALUES: usize = ARGS + core::mem::size_of::<Payloads>();
- /// Offset of `state` field
- pub const STATE: usize = TAG_RETURN_VALUES + core::mem::size_of::<Payloads>();
  /// Offset of `revision` field
- pub const REVISION: usize = STATE + core::mem::size_of::<usize>();
+ pub const REVISION: usize = TAG_RETURN_VALUES + core::mem::size_of::<Payloads>();
  }
 
  pub mod stack_limits {
@@ -242,6 +267,14 @@ pub mod offsets {
  pub const LAST_WASM_ENTRY_SP: usize = offset_of!(StackLimits, last_wasm_entry_sp);
  }
 
+ pub mod common_stack_information {
+ use crate::CommonStackInformation;
+ use memoffset::offset_of;
+
+ pub const LIMITS: usize = offset_of!(CommonStackInformation, limits);
+ pub const STATE: usize = offset_of!(CommonStackInformation, state);
+ }
+
  /// Size of wasmtime_runtime::continuation::FiberStack.
  /// We test there that this value is correct.
  pub const FIBER_STACK_SIZE: usize = 3 * core::mem::size_of::<usize>();