Continuation objects as fat pointers (#182)

This patch makes continuation objects allocation free by representing a
continuation object as an immediate, whose 16 high bits are the witness
of the continuation reference revision counter, and the 48 low bits are
the actual continuation reference address. This gives us a cheap way of
implementing the linearity check for continuations. Using 16 bits for
the witness, means that we are currently restricted to 2^16 uses, which
suffices for our examples and benchmarks, but isn't sufficient in
general. However, we can readily generalise this implementation in a
subsequent patch.

The compile time option `unsafe_disable_continuation_linearity_check`
preserves its behaviour.

Resolves #178.

cranelift/codegen/src/ir/memflags.rs

@@ -339,9 +339,8 @@ impl MemFlags {
             0b1010 => Some(TrapCode::Interrupt),
             0b1011 => Some(TrapCode::NullReference),
             0b1100 => Some(TrapCode::NullI31Ref),
-            0b1101 => Some(TrapCode::UnhandledTag),
-            // 0b1101 => {} not allocated
-            // 0b1110 => {} not allocated
+            0b1101 => Some(TrapCode::ContinuationAlreadyConsumed),
+            0b1110 => Some(TrapCode::UnhandledTag),
             0b1111 => None,
             _ => unreachable!(),
         }
@@ -369,7 +368,8 @@ impl MemFlags {
             Some(TrapCode::Interrupt) => 0b1010,
             Some(TrapCode::NullReference) => 0b1011,
             Some(TrapCode::NullI31Ref) => 0b1100,
-            Some(TrapCode::UnhandledTag) => 0b1101,
+            Some(TrapCode::ContinuationAlreadyConsumed) => 0b1101,
+            Some(TrapCode::UnhandledTag) => 0b1110,
             None => 0b1111,
 
             Some(TrapCode::User(_)) => panic!("cannot set user trap code in mem flags"),

cranelift/codegen/src/ir/trapcode.rs

@@ -58,6 +58,10 @@ pub enum TrapCode {
     /// We are suspending to a tag for which there is no active handler.
     UnhandledTag,
 
+    /// Attempt to resume a continuation object whose revision is out
+    /// of sync with the underlying continuation reference.
+    ContinuationAlreadyConsumed,
+
     /// A null `i31ref` was encountered which was required to be non-null.
     NullI31Ref,
 }
@@ -79,6 +83,7 @@ impl TrapCode {
             TrapCode::Interrupt,
             TrapCode::NullReference,
             TrapCode::UnhandledTag,
+            TrapCode::ContinuationAlreadyConsumed,
         ]
     }
 }
@@ -102,6 +107,7 @@ impl Display for TrapCode {
             NullReference => "null_reference",
             UnhandledTag => "unhandled_tag",
             NullI31Ref => "null_i31ref",
+            ContinuationAlreadyConsumed => "continuation_already_consumed",
         };
         f.write_str(identifier)
     }
@@ -127,6 +133,7 @@ impl FromStr for TrapCode {
             "null_reference" => Ok(NullReference),
             "unhandled_tag" => Ok(UnhandledTag),
             "null_i31ref" => Ok(NullI31Ref),
+            "continuation_already_consumed" => Ok(ContinuationAlreadyConsumed),
             _ if s.starts_with("user") => s[4..].parse().map(User).map_err(|_| ()),
             _ => Err(()),
         }

crates/continuations/src/lib.rs

@@ -302,6 +302,8 @@ pub mod offsets {
         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 mod stack_limits {

crates/cranelift-shared/src/lib.rs

@@ -87,6 +87,8 @@ pub fn mach_trap_to_trap(trap: &MachTrap) -> Option<TrapInformation> {
             ir::TrapCode::User(CANNOT_ENTER_CODE) => Trap::CannotEnterComponent,
             ir::TrapCode::NullReference => Trap::NullReference,
             ir::TrapCode::NullI31Ref => Trap::NullI31Ref,
+            ir::TrapCode::UnhandledTag => Trap::UnhandledTag,
+            ir::TrapCode::ContinuationAlreadyConsumed => Trap::ContinuationAlreadyConsumed,
 
             // These do not get converted to wasmtime traps, since they
             // shouldn't ever be hit in theory. Instead of catching and handling

crates/cranelift/src/lib.rs

@@ -306,6 +306,7 @@ pub fn mach_trap_to_trap(trap: &MachTrap) -> Option<TrapInformation> {
             ir::TrapCode::NullReference => Trap::NullReference,
             ir::TrapCode::UnhandledTag => Trap::UnhandledTag,
             ir::TrapCode::NullI31Ref => Trap::NullI31Ref,
+            ir::TrapCode::ContinuationAlreadyConsumed => Trap::ContinuationAlreadyConsumed,
 
             // These do not get converted to wasmtime traps, since they
             // shouldn't ever be hit in theory. Instead of catching and handling

crates/cranelift/src/wasmfx/baseline.rs

@@ -8,10 +8,49 @@ use cranelift_codegen::ir::InstBuilder;
 use cranelift_frontend::{FunctionBuilder, Switch};
 use cranelift_wasm::FuncEnvironment;
 use cranelift_wasm::{FuncTranslationState, WasmResult, WasmValType};
-use shared::typed_continuations_cont_obj_get_cont_ref;
-use shared::typed_continuations_new_cont_obj;
+use shared::{assemble_contobj, disassemble_contobj};
 use wasmtime_environ::PtrSize;
 
+fn get_revision<'a>(
+    _env: &mut crate::func_environ::FuncEnvironment<'a>,
+    builder: &mut FunctionBuilder,
+    contref: ir::Value,
+) -> ir::Value {
+    if cfg!(feature = "unsafe_disable_continuation_linearity_check") {
+        builder.ins().iconst(I64, 0)
+    } else {
+        let mem_flags = ir::MemFlags::trusted();
+        builder.ins().load(I64, mem_flags, contref, 0)
+    }
+}
+
+fn compare_revision_and_increment<'a>(
+    env: &mut crate::func_environ::FuncEnvironment<'a>,
+    builder: &mut FunctionBuilder,
+    contref: ir::Value,
+    witness: ir::Value,
+) -> ir::Value {
+    if cfg!(feature = "unsafe_disable_continuation_linearity_check") {
+        builder.ins().iconst(I64, 0)
+    } else {
+        let mem_flags = ir::MemFlags::trusted();
+        let revision = get_revision(env, builder, contref);
+
+        let evidence = builder.ins().icmp(IntCC::Equal, witness, revision);
+        builder
+            .ins()
+            .trapz(evidence, ir::TrapCode::ContinuationAlreadyConsumed);
+
+        let revision_plus1 = builder.ins().iadd_imm(revision, 1);
+        let overflow = builder
+            .ins()
+            .icmp_imm(IntCC::UnsignedLessThan, revision_plus1, 1 << 16);
+        builder.ins().trapz(overflow, ir::TrapCode::IntegerOverflow); // TODO(dhil): Consider introducing a designated trap code.
+        builder.ins().store(mem_flags, revision_plus1, contref, 0);
+        revision_plus1
+    }
+}
+
 fn typed_continuations_load_payloads<'a>(
     env: &mut crate::func_environ::FuncEnvironment<'a>,
     builder: &mut FunctionBuilder,
@@ -194,9 +233,9 @@ pub(crate) fn translate_resume<'a>(
     let forwarding_block = builder.create_block();
 
     // Prelude: Push the continuation arguments.
-    {
-        let resumee_fiber =
-            shared::typed_continuations_cont_obj_get_cont_ref(env, builder, resumee_obj);
+    let next_revision = {
+        let (witness, resumee_fiber) = disassemble_contobj(env, builder, resumee_obj);
+        let next_revision = compare_revision_and_increment(env, builder, resumee_fiber, witness);
         if resume_args.len() > 0 {
             let nargs = builder.ins().iconst(I64, resume_args.len() as i64);
 
@@ -221,7 +260,8 @@ pub(crate) fn translate_resume<'a>(
         }
 
         builder.ins().jump(resume_block, &[resumee_fiber]);
-    }
+        next_revision
+    };
 
     // Resume block: here we continue the (suspended) resumee.
     let (tag, resumee_fiber) = {
@@ -298,7 +338,7 @@ pub(crate) fn translate_resume<'a>(
         let mut args = typed_continuations_load_payloads(env, builder, &param_types);
 
         // Create and push the continuation object.
-        let resumee_obj = shared::typed_continuations_new_cont_obj(env, builder, resumee_fiber);
+        let resumee_obj = assemble_contobj(env, builder, next_revision, resumee_fiber);
         args.push(resumee_obj);
 
         // Finally, emit the jump to `label`.
@@ -378,11 +418,11 @@ pub(crate) fn translate_cont_bind<'a>(
     args: &[ir::Value],
     remaining_arg_count: usize,
 ) -> ir::Value {
-    let contref = typed_continuations_cont_obj_get_cont_ref(env, builder, contobj);
+    let (witness, contref) = disassemble_contobj(env, builder, contobj);
+    let revision = compare_revision_and_increment(env, builder, contref, witness);
     let remaining_arg_count = builder.ins().iconst(I32, remaining_arg_count as i64);
     typed_continuations_store_resume_args(env, builder, args, remaining_arg_count, contref);
-
-    typed_continuations_new_cont_obj(env, builder, contref)
+    assemble_contobj(env, builder, revision, contref)
 }
 
 pub(crate) fn translate_cont_new<'a>(
@@ -397,7 +437,8 @@ pub(crate) fn translate_cont_new<'a>(
     let nargs = builder.ins().iconst(I64, arg_types.len() as i64);
     let nreturns = builder.ins().iconst(I64, return_types.len() as i64);
     call_builtin!(builder, env, let contref = tc_baseline_cont_new(func, nargs, nreturns));
-    let contobj = typed_continuations_new_cont_obj(env, builder, contref);
+    let revision = get_revision(env, builder, contref);
+    let contobj = assemble_contobj(env, builder, revision, contref);
     Ok(contobj)
 }
 

crates/cranelift/src/wasmfx/optimized.rs

@@ -8,8 +8,6 @@ use cranelift_codegen::ir::InstBuilder;
 use cranelift_frontend::{FunctionBuilder, Switch};
 use cranelift_wasm::FuncEnvironment;
 use cranelift_wasm::{FuncTranslationState, WasmResult, WasmValType};
-use shared::typed_continuations_cont_obj_get_cont_ref;
-use shared::typed_continuations_new_cont_obj;
 use wasmtime_environ::PtrSize;
 
 #[macro_use]
@@ -76,7 +74,7 @@ pub(crate) mod typed_continuation_helpers {
                          val: ir::Value| {
             let ty = builder.func.dfg.value_type(val);
             let val = match ty {
-                I32 => builder.ins().uextend(I64, val),
+                I8 | I32 => builder.ins().uextend(I64, val),
                 I64 => val,
                 _ => panic!("Cannot print type {}", ty),
             };
@@ -419,6 +417,57 @@ pub(crate) mod typed_continuation_helpers {
             let offset = wasmtime_continuations::offsets::vm_cont_ref::PARENT_CHAIN as i32;
             new_stack_chain.store(env, builder, self.address, offset)
         }
+
+        /// Gets the revision counter the a given continuation
+        /// reference.
+        pub fn get_revision<'a>(
+            &mut self,
+            _env: &mut crate::func_environ::FuncEnvironment<'a>,
+            builder: &mut FunctionBuilder,
+        ) -> ir::Value {
+            if cfg!(feature = "unsafe_disable_continuation_linearity_check") {
+                builder.ins().iconst(I64, 0)
+            } else {
+                let mem_flags = ir::MemFlags::trusted();
+                let offset = wasmtime_continuations::offsets::vm_cont_ref::REVISION as i32;
+                let revision = builder.ins().load(I64, mem_flags, self.address, offset);
+                revision
+            }
+        }
+
+        /// Sets the revision counter on the given continuation
+        /// reference to `revision + 1`.
+        pub fn incr_revision<'a>(
+            &mut self,
+            env: &mut crate::func_environ::FuncEnvironment<'a>,
+            builder: &mut FunctionBuilder,
+            revision: ir::Value,
+        ) -> ir::Value {
+            if cfg!(feature = "unsafe_disable_continuation_linearity_check") {
+                builder.ins().iconst(I64, 0)
+            } else {
+                if cfg!(debug_assertions) {
+                    let actual_revision = self.get_revision(env, builder);
+                    emit_debug_assert_eq!(env, builder, revision, actual_revision);
+                }
+                let mem_flags = ir::MemFlags::trusted();
+                let offset = wasmtime_continuations::offsets::vm_cont_ref::REVISION as i32;
+                let revision_plus1 = builder.ins().iadd_imm(revision, 1);
+                builder
+                    .ins()
+                    .store(mem_flags, revision_plus1, self.address, offset);
+                if cfg!(debug_assertions) {
+                    let new_revision = self.get_revision(env, builder);
+                    emit_debug_assert_eq!(env, builder, revision_plus1, new_revision);
+                }
+                let overflow =
+                    builder
+                        .ins()
+                        .icmp_imm(IntCC::UnsignedLessThan, revision_plus1, 1 << 16);
+                builder.ins().trapz(overflow, ir::TrapCode::IntegerOverflow); // TODO(dhil): Consider introducing a designated trap code.
+                revision_plus1
+            }
+        }
     }
 
     impl Payloads {
@@ -1268,11 +1317,31 @@ pub(crate) fn translate_cont_bind<'a>(
     args: &[ir::Value],
     remaining_arg_count: usize,
 ) -> ir::Value {
-    let contref = typed_continuations_cont_obj_get_cont_ref(env, builder, contobj);
+    //let contref = typed_continuations_cont_obj_get_cont_ref(env, builder, contobj);
+    let (witness, contref) = shared::disassemble_contobj(env, builder, contobj);
+    let mut vmcontref = tc::VMContRef::new(contref, env.pointer_type());
+    let revision = vmcontref.get_revision(env, builder);
+    let evidence = builder.ins().icmp(IntCC::Equal, witness, revision);
+    emit_debug_println!(
+        env,
+        builder,
+        "[cont_bind] witness = {}, revision = {}, evidence = {}",
+        witness,
+        revision,
+        evidence
+    );
+    builder
+        .ins()
+        .trapz(evidence, ir::TrapCode::ContinuationAlreadyConsumed);
+
     let remaining_arg_count = builder.ins().iconst(I32, remaining_arg_count as i64);
     typed_continuations_store_resume_args(env, builder, args, remaining_arg_count, contref);
 
-    typed_continuations_new_cont_obj(env, builder, contref)
+    let revision = vmcontref.incr_revision(env, builder, revision);
+    emit_debug_println!(env, builder, "new revision = {}", revision);
+    let contobj = shared::assemble_contobj(env, builder, revision, contref);
+    emit_debug_println!(env, builder, "[cont_bind] contobj = {:p}", contobj);
+    contobj
 }
 
 pub(crate) fn translate_cont_new<'a>(
@@ -1286,7 +1355,9 @@ pub(crate) fn translate_cont_new<'a>(
     let nargs = builder.ins().iconst(I32, arg_types.len() as i64);
     let nreturns = builder.ins().iconst(I32, return_types.len() as i64);
     call_builtin!(builder, env, let contref = tc_cont_new(func, nargs, nreturns));
-    let contobj = typed_continuations_new_cont_obj(env, builder, contref);
+    let tag = tc::VMContRef::new(contref, env.pointer_type()).get_revision(env, builder);
+    let contobj = shared::assemble_contobj(env, builder, tag, contref);
+    emit_debug_println!(env, builder, "[cont_new] contobj = {:p}", contobj);
     Ok(contobj)
 }
 
@@ -1308,9 +1379,20 @@ pub(crate) fn translate_resume<'a>(
 
     // Preamble: Part of previously active block
 
-    let (resume_contref, parent_stack_chain) = {
-        let resume_contref =
-            shared::typed_continuations_cont_obj_get_cont_ref(env, builder, contobj);
+    let (next_revision, resume_contref, parent_stack_chain) = {
+        let (witness, resume_contref) = shared::disassemble_contobj(env, builder, contobj);
+
+        let mut vmcontref = tc::VMContRef::new(resume_contref, env.pointer_type());
+
+        let revision = vmcontref.get_revision(env, builder);
+        let evidence = builder.ins().icmp(IntCC::Equal, revision, witness);
+        emit_debug_println!(env, builder, "[resume] contobj = {:p}, resume_contref = {:p} witness = {}, revision = {}, evidence = {}", contobj, resume_contref, witness, revision, evidence);
+
+        builder
+            .ins()
+            .trapz(evidence, ir::TrapCode::ContinuationAlreadyConsumed);
+        let next_revision = vmcontref.incr_revision(env, builder, revision);
+        emit_debug_println!(env, builder, "[resume] new revision = {}", next_revision);
 
         if resume_args.len() > 0 {
             // We store the arguments in the `VMContRef` to be resumed.
@@ -1322,14 +1404,10 @@ pub(crate) fn translate_resume<'a>(
         let original_stack_chain =
             tc::VMContext::new(vmctx, env.pointer_type()).load_stack_chain(env, builder);
         original_stack_chain.assert_not_absent(env, builder);
-        tc::VMContRef::new(resume_contref, env.pointer_type()).set_parent_stack_chain(
-            env,
-            builder,
-            &original_stack_chain,
-        );
+        vmcontref.set_parent_stack_chain(env, builder, &original_stack_chain);
 
         builder.ins().jump(resume_block, &[]);
-        (resume_contref, original_stack_chain)
+        (next_revision, resume_contref, original_stack_chain)
     };
 
     // Resume block: actually resume the fiber corresponding to the
@@ -1499,12 +1577,19 @@ pub(crate) fn translate_resume<'a>(
         // link to `StackChain::Absent`.
         let pointer_type = env.pointer_type();
         let chain = tc::StackChain::absent(builder, pointer_type);
-        tc::VMContRef::new(resume_contref, pointer_type)
-            .set_parent_stack_chain(env, builder, &chain);
+        let mut vmcontref = tc::VMContRef::new(resume_contref, pointer_type);
+        vmcontref.set_parent_stack_chain(env, builder, &chain);
 
         // Create and push the continuation object. We only create
         // them here because we don't need them when forwarding.
-        let contobj = typed_continuations_new_cont_obj(env, builder, resume_contref);
+        let contobj = shared::assemble_contobj(env, builder, next_revision, resume_contref);
+        emit_debug_println!(
+            env,
+            builder,
+            "[resume] revision = {}, contobj = {:p}",
+            next_revision,
+            contobj
+        );
 
         args.push(contobj);