[workerd-cxx] Add kj::Own ffi type to workerd-cxx (#26)

Extends the cxx bridge module to allow using the `Own` smart pointer in Rust from C++ code. 
Using a static disposer is not supported in Rust.
Passing a null Own to Rust, through returns or function arguments, will result in an exception. 
Tests contain both tests that were hand-written and generated by Claude Code.

Author: LordGoatius

gen/src/write.rs

@@ -223,7 +223,7 @@ fn pick_includes_and_builtins(out: &mut OutFile, apis: &[Api]) {
             Type::SliceRef(_) => out.builtin.rust_slice = true,
             Type::Array(_) => out.include.array = true,
             Type::Ref(_) | Type::Void(_) | Type::Ptr(_) => {}
-            Type::Future(_) => out.include.kj_rs = true,
+            Type::Future(_) | Type::Own(_) => out.include.kj_rs = true,
         }
     }
 }
@@ -1030,6 +1030,12 @@ fn write_rust_function_shim_impl(
         return;
     }
     writeln!(out, " {{");
+    sig.args.iter()
+        .filter(|arg| {
+            matches!(arg.ty, Type::Own(_))
+        }).for_each(|arg_own| {
+            writeln!(out, "  KJ_ASSERT({}.get() != nullptr, \"Cannot pass a null Own to Rust\");", arg_own.name.cxx);
+        });
     for arg in &sig.args {
         if arg.ty != RustString && out.types.needs_indirect_abi(&arg.ty) {
             out.include.utility = true;
@@ -1228,6 +1234,11 @@ fn write_type(out: &mut OutFile, ty: &Type) {
             write_type(out, &ptr.inner);
             write!(out, ">");
         }
+        Type::Own(ptr) => {
+            write!(out, "::kj::Own<");
+            write_type(out, &ptr.inner);
+            write!(out, ">");
+        }
         Type::SharedPtr(ptr) => {
             write!(out, "::std::shared_ptr<");
             write_type(out, &ptr.inner);
@@ -1328,6 +1339,7 @@ fn write_space_after_type(out: &mut OutFile, ty: &Type) {
         Type::Ident(_)
         | Type::RustBox(_)
         | Type::UniquePtr(_)
+        | Type::Own(_)
         | Type::SharedPtr(_)
         | Type::WeakPtr(_)
         | Type::Str(_)
@@ -1404,6 +1416,7 @@ fn write_generic_instantiations(out: &mut OutFile) {
             ImplKey::RustBox(ident) => write_rust_box_extern(out, ident),
             ImplKey::RustVec(ident) => write_rust_vec_extern(out, ident),
             ImplKey::UniquePtr(ident) => write_unique_ptr(out, ident),
+            ImplKey::Own(ident) => write_kj_own(out, ident),
             ImplKey::SharedPtr(ident) => write_shared_ptr(out, ident),
             ImplKey::WeakPtr(ident) => write_weak_ptr(out, ident),
             ImplKey::CxxVector(ident) => write_cxx_vector(out, ident),
@@ -1615,6 +1628,34 @@ fn write_rust_vec_impl(out: &mut OutFile, key: NamedImplKey) {
     writeln!(out, "}}");
 }
 
+// Writes static assertion that we do not use an Own with a static disposer
+fn write_kj_own(out: &mut OutFile, key: NamedImplKey) {
+    let ident = key.rust;
+    let resolve = out.types.resolve(ident);
+    let inner = resolve.name.to_fully_qualified();
+    let instance = resolve.name.to_symbol();
+    
+    out.include.utility = true;
+    out.include.kj_rs = true;
+
+    // Static disposers are not supported, only Owns containing 2 pointers are allowed
+    writeln!(
+        out,
+        "static_assert(sizeof(::kj::Own<{}>) == 2 * sizeof(void *), \"Static disposers for Own are not supported in workerd-cxx\");",
+        inner,
+    );
+    writeln!(
+        out,
+        "static_assert(alignof(::kj::Own<{}>) == sizeof(void *), \"Static disposers for Own are not supported in workerd-cxx\");",
+        inner,
+    );
+    writeln!(
+        out,
+        "static_assert(!::std::is_base_of<::kj::Refcounted, {}>::value, \"Value must not inherit from kj::Refcounted\");",
+        inner
+    );
+}
+
 fn write_unique_ptr(out: &mut OutFile, key: NamedImplKey) {
     let ty = UniquePtr::Ident(key.rust);
     write_unique_ptr_common(out, ty);

kj-rs/lib.rs

@@ -14,9 +14,11 @@ mod awaiter;
 mod future;
 mod promise;
 mod waker;
+mod own;
 
 pub mod repr {
     pub use crate::future::repr::*;
+    pub use crate::own::repr::*;
 }
 
 pub type Result<T> = std::io::Result<T>;

kj-rs/own.c++

@@ -0,0 +1,11 @@
+#include "own.h"
+
+extern "C" {
+// As of right now, this function works and passes all tests, destroying all tested objects correctly.
+// This works because disposers work by virtual call, and the disposer is created during creation of
+// the Own, so the type of the Own at destruction doesn't actually matter for Owns that do not use a
+// static disposer, which are not currently supported by workerd-cxx.
+void cxxbridge$kjrs$own$drop(void *own) {
+  reinterpret_cast<kj::Own<void> *>(own)->~Own();
+}
+}

kj-rs/own.h

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <kj/memory.h>
+
+extern "C" {
+    void cxxbridge$kjrs$own$drop(void* own);
+}

kj-rs/own.rs

@@ -0,0 +1,169 @@
+//! The `workerd-cxx` module containing the [`Own<T>`] type, which is bindings to the `kj::Own<T>` C++ type
+
+use static_assertions::{assert_eq_align, assert_eq_size};
+
+assert_eq_size!(repr::Own<()>, [*const (); 2]);
+assert_eq_align!(repr::Own<()>, *const ());
+
+pub mod repr {
+    use std::ffi::c_void;
+    use std::fmt::{self, Debug, Display};
+    use std::hash::{Hash, Hasher};
+    use std::ops::Deref;
+    use std::ops::DerefMut;
+    use std::pin::Pin;
+    use std::ptr::NonNull;
+
+    /// A [`Own<T>`] represents the `kj::Own<T>`. It is a smart pointer to an opaque C++ type.
+    /// Safety:
+    /// - Passing a null `kj::Own` to rust is considered unsafe from the C++ side,
+    ///   and it is required that this invariant is upheld in C++ code.
+    /// - Currently, it is runtime asserted in the bridge macro that no null Own can be passed
+    ///   to Rust
+    #[repr(C)]
+    pub struct Own<T> {
+        disposer: *const c_void,
+        ptr: NonNull<T>,
+    }
+
+    /// Public-facing Own api
+    impl<T> Own<T> {
+        /// Returns a mutable pinned reference to the object owned by this [`Own`]
+        /// if any, otherwise None.
+        pub fn as_mut(&mut self) -> Pin<&mut T> {
+            // Safety: Passing a null kj::Own to Rust from C++ is not supported.
+            unsafe {
+                let mut_reference = self.ptr.as_mut();
+                Pin::new_unchecked(mut_reference)
+            }
+        }
+
+        /// Returns a mutable pinned reference to the object owned by this
+        /// [`Own`].
+        ///
+        /// ```compile_fail
+        /// let mut own = ffi::cxx_kj_own();
+        /// let pin1 = own.pin_mut();
+        /// let pin2 = own.pin_mut();
+        /// pin1.set_data(12); // Causes a compile fail, because we invalidated the first borrow
+        /// ```
+        ///
+        /// ```compile_fail
+        ///
+        /// let mut own = ffi::cxx_kj_own();
+        /// let pin = own.pin_mut();
+        /// let moved  = own;
+        /// own.set_data(143); // Compile fail, because we tried using a moved object
+        /// ```
+        pub fn pin_mut(&mut self) -> Pin<&mut T> {
+            self.as_mut()
+        }
+
+        /// Returns a raw const pointer to the object owned by this [`Own`]
+        #[must_use]
+        pub fn as_ptr(&self) -> *const T {
+            self.ptr.as_ptr().cast()
+        }
+    }
+
+    impl<T> AsRef<T> for Own<T> {
+        /// Returns a reference to the object owned by this [`Own`] if any,
+        /// otherwise None.
+        fn as_ref(&self) -> &T {
+            // Safety: Passing a null kj::Own to Rust from C++ is not supported.
+            unsafe { self.ptr.as_ref() }
+        }
+    }
+
+    unsafe impl<T> Send for Own<T> where T: Send {}
+
+    unsafe impl<T> Sync for Own<T> where T: Sync {}
+
+    impl<T> Deref for Own<T> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            self.as_ref()
+        }
+    }
+
+    impl<T> DerefMut for Own<T>
+    where
+        T: Unpin,
+    {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            Pin::into_inner(self.as_mut())
+        }
+    }
+
+    // Own<T> is safe to implement Unpin because moving the Own doesn't move the pointee, and
+    // the drop implentation doesn't depend on the Own's location, because it's handed by virtual dispatch
+    impl<T> Unpin for Own<T> {}
+
+    impl<T> Debug for Own<T> {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(f, "Own(ptr: {:p}, disposer: {:p})", self.ptr, self.disposer)
+        }
+    }
+
+    impl<T> Display for Own<T>
+    where
+        T: Display,
+    {
+        fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+            Display::fmt(self.as_ref(), formatter)
+        }
+    }
+
+    impl<T> PartialEq for Own<T>
+    where
+        T: PartialEq,
+    {
+        fn eq(&self, other: &Self) -> bool {
+            self.as_ref() == other.as_ref()
+        }
+    }
+
+    impl<T> Eq for Own<T> where T: Eq {}
+
+    impl<T> PartialOrd for Own<T>
+    where
+        T: PartialOrd,
+    {
+        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+            PartialOrd::partial_cmp(&self.as_ref(), &other.as_ref())
+        }
+    }
+
+    impl<T> Ord for Own<T>
+    where
+        T: Ord,
+    {
+        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+            Ord::cmp(&self.as_ref(), &other.as_ref())
+        }
+    }
+
+    impl<T> Hash for Own<T>
+    where
+        T: Hash,
+    {
+        fn hash<H: Hasher>(&self, state: &mut H) {
+            self.as_ref().hash(state);
+        }
+    }
+
+    impl<T> Drop for Own<T> {
+        fn drop(&mut self) {
+            unsafe extern "C" {
+                #[link_name = "cxxbridge$kjrs$own$drop"]
+                fn __drop(this: *mut c_void);
+            }
+
+            let this = std::ptr::from_mut::<Self>(self).cast::<c_void>();
+            unsafe {
+                __drop(this);
+            }
+        }
+    }
+}

kj-rs/tests/BUILD.bazel

@@ -45,6 +45,7 @@ rust_cxx_bridge(
     src = "lib.rs",
     hdrs = [
         "test-promises.h",
+        "test-own.h"
     ],
     include_prefix = "kj-rs-demo",
     deps = [
@@ -56,6 +57,7 @@ cc_library(
     name = "test-promises",
     srcs = [
         "test-promises.c++",
+        "test-own.c++"
     ],
     linkstatic = select({
         "@platforms//os:windows": True,

kj-rs/tests/lib.rs

@@ -1,8 +1,13 @@
 #![allow(clippy::needless_lifetimes)]
 #![allow(clippy::missing_errors_doc)]
 #![allow(clippy::unused_async)]
+#![allow(clippy::must_use_candidate)]
+#![allow(clippy::cast_possible_truncation)]
+#![allow(clippy::should_panic_without_expect)]
+#![allow(clippy::missing_panics_doc)]
 
 mod test_futures;
+mod test_own;
 
 use test_futures::{
     new_awaiting_future_i32, new_error_handling_future_void_infallible, new_errored_future_void,
@@ -11,6 +16,8 @@ use test_futures::{
     new_waking_future_void, new_wrapped_waker_future_void,
 };
 
+use kj_rs::repr::Own;
+
 type Result<T> = std::io::Result<T>;
 type Error = std::io::Error;
 
@@ -32,6 +39,37 @@ mod ffi {
         async fn new_ready_promise_shared_type() -> Shared;
     }
 
+    // Helper functions to test `kj_rs::Own`
+    unsafe extern "C++" {
+        include!("kj-rs-demo/test-own.h");
+        type OpaqueCxxClass;
+
+        #[cxx_name = "getData"]
+        fn get_data(&self) -> u64;
+        #[cxx_name = "setData"]
+        fn set_data(self: Pin<&mut OpaqueCxxClass>, val: u64);
+
+        fn cxx_kj_own() -> Own<OpaqueCxxClass>;
+        fn null_kj_own() -> Own<OpaqueCxxClass>;
+        fn give_own_back(own: Own<OpaqueCxxClass>);
+        fn modify_own_return_test();
+        fn breaking_things() -> Own<OpaqueCxxClass>;
+
+        fn own_integer() -> Own<i64>;
+        fn own_integer_attached() -> Own<i64>;
+
+        fn null_exception_test_driver_1() -> String;
+        fn null_exception_test_driver_2() -> String;
+        fn rust_take_own_driver();
+    }
+
+    // Helper function to test moving `Own` to C++
+    extern "Rust" {
+        fn modify_own_return(cpp_own: Own<OpaqueCxxClass>) -> Own<OpaqueCxxClass>;
+        fn take_own(cpp_own: Own<OpaqueCxxClass>);
+        fn get_null() -> Own<OpaqueCxxClass>;
+    }
+
     enum CloningAction {
         None,
         CloneSameThread,
@@ -75,6 +113,22 @@ mod ffi {
     }
 }
 
+pub fn modify_own_return(mut own: Own<ffi::OpaqueCxxClass>) -> Own<ffi::OpaqueCxxClass> {
+    own.pin_mut().set_data(72);
+    own
+}
+
+pub fn get_null() -> Own<ffi::OpaqueCxxClass> {
+    ffi::null_kj_own()
+}
+
+pub fn take_own(cpp_own: Own<ffi::OpaqueCxxClass>) {
+    assert_eq!(cpp_own.get_data(), 14);
+    // The point of this function is to drop the [`Own`] from rust and this makes
+    // it explicit, while avoiding a clippy lint
+    std::mem::drop(cpp_own);
+}
+
 pub async fn lifetime_arg_void<'a>(_buf: &'a [u8]) {}
 
 pub async fn lifetime_arg_result<'a>(_buf: &'a [u8]) -> Result<()> {