[workerd-cxx] Add kj::Rc and kj::Arc bindings to kj-rs

Author: LordGoatius

gen/src/write.rs

@@ -233,7 +233,9 @@ 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(_) | Type::Own(_) => out.include.kj_rs = true,
+            Type::Future(_) | Type::Own(_) | Type::KjRc(_) | Type::KjArc(_) => {
+                out.include.kj_rs = true;
+            }
         }
     }
 }
@@ -1253,6 +1255,16 @@ fn write_type(out: &mut OutFile, ty: &Type) {
             write_type(out, &ptr.inner);
             write!(out, ">");
         }
+        Type::KjRc(ptr) => {
+            write!(out, "::kj::Rc<");
+            write_type(out, &ptr.inner);
+            write!(out, ">");
+        }
+        Type::KjArc(ptr) => {
+            write!(out, "::kj::Arc<");
+            write_type(out, &ptr.inner);
+            write!(out, ">");
+        }
         Type::SharedPtr(ptr) => {
             write!(out, "::std::shared_ptr<");
             write_type(out, &ptr.inner);
@@ -1354,6 +1366,8 @@ fn write_space_after_type(out: &mut OutFile, ty: &Type) {
         | Type::RustBox(_)
         | Type::UniquePtr(_)
         | Type::Own(_)
+        | Type::KjRc(_)
+        | Type::KjArc(_)
         | Type::SharedPtr(_)
         | Type::WeakPtr(_)
         | Type::Str(_)
@@ -1431,6 +1445,8 @@ fn write_generic_instantiations(out: &mut OutFile) {
             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::KjRc(ident) => write_kj_rc(out, ident),
+            ImplKey::KjArc(ident) => write_kj_arc(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),
@@ -1669,6 +1685,106 @@ fn write_kj_own(out: &mut OutFile, key: NamedImplKey) {
     );
 }
 
+// Writes static assertion that we do not use an Own with a static disposer
+fn write_kj_rc(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::Rc<{}>) == 2 * sizeof(void *), \"Static disposers for Own are not supported in workerd-cxx\");",
+        inner,
+    );
+    writeln!(
+        out,
+        "static_assert(alignof(::kj::Rc<{}>) == 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 inherit kj::Refcounted\");",
+        inner
+    );
+
+    begin_function_definition(out);
+    writeln!(
+        out,
+        "void cxxbridge1$kj_rs$rc${}$is_shared(::kj::Refcounted *ptr, bool *ret) noexcept {{",
+        instance,
+    );
+    writeln!(out, "  *ret = ptr->isShared();");
+    writeln!(out, "}}");
+
+    begin_function_definition(out);
+    writeln!(
+        out,
+        "void cxxbridge1$kj_rs$rc${}$add_ref(::kj::Rc<{}> *refcounted, ::kj::Rc<{}> *ptr) noexcept {{",
+        instance, inner, inner
+    );
+    writeln!(
+        out,
+        "  ::new (ptr) ::kj::Rc<{}>(refcounted->addRef());",
+        inner
+    );
+    writeln!(out, "}}");
+}
+
+// Writes static assertion that we do not use an Own with a static disposer
+fn write_kj_arc(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::Arc<{}>) == 2 * sizeof(void *), \"Static disposers for Own are not supported in workerd-cxx\");",
+        inner,
+    );
+    writeln!(
+        out,
+        "static_assert(alignof(::kj::Arc<{}>) == sizeof(void *), \"Static disposers for Own are not supported in workerd-cxx\");",
+        inner,
+    );
+    writeln!(
+        out,
+        "static_assert(::std::is_base_of<::kj::AtomicRefcounted, {}>::value, \"Value must inherit kj::AtomicRefcounted\");",
+        inner
+    );
+
+    begin_function_definition(out);
+    writeln!(
+        out,
+        "void cxxbridge1$kj_rs$arc${}$is_shared(::kj::AtomicRefcounted *ptr, bool *ret) noexcept {{",
+        instance,
+    );
+    writeln!(out, "  *ret = ptr->isShared();");
+    writeln!(out, "}}");
+
+    begin_function_definition(out);
+    writeln!(
+        out,
+        "void cxxbridge1$kj_rs$arc${}$add_ref(::kj::Arc<{}> *refcounted, ::kj::Arc<{}> *ptr) noexcept {{",
+        instance, inner, inner
+    );
+    writeln!(
+        out,
+        "  ::new (ptr) ::kj::Arc<{}>(refcounted->addRef());",
+        inner
+    );
+    writeln!(out, "}}");
+}
+
 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

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

kj-rs/own.rs

@@ -21,7 +21,7 @@ pub mod repr {
     /// - Currently, it is runtime asserted in the bridge macro that no null Own can be passed
     ///   to Rust
     #[repr(C)]
-    pub struct Own<T> {
+    pub struct Own<T: ?Sized> {
         disposer: *const c_void,
         ptr: NonNull<T>,
     }
@@ -62,7 +62,7 @@ pub mod repr {
         /// 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()
+            self.ptr.as_ptr().cast_const()
         }
     }
 
@@ -153,7 +153,7 @@ pub mod repr {
         }
     }
 
-    impl<T> Drop for Own<T> {
+    impl<T: ?Sized> Drop for Own<T> {
         fn drop(&mut self) {
             unsafe extern "C" {
                 #[link_name = "cxxbridge$kjrs$own$drop"]

kj-rs/refcount.rs

@@ -0,0 +1,97 @@
+//! Module for both [`KjRc`] and [`KjArc`], since they're nearly identical types
+
+pub mod repr {
+    use crate::Own;
+    use std::ops::Deref;
+
+    /// # Safety
+    /// Should only be automatically implemented by the bridge macro
+    pub unsafe trait Refcounted {
+        fn is_shared(&self) -> bool;
+        fn add_ref(rc: &KjRc<Self>) -> KjRc<Self>;
+    }
+    /// # Safety
+    /// Should only be automatically implemented by the bridge macro
+    pub unsafe trait AtomicRefcounted {
+        fn is_shared(&self) -> bool;
+        fn add_ref(arc: &KjArc<Self>) -> KjArc<Self>;
+    }
+
+    #[repr(C)]
+    /// Bindings to the kj type `kj::Rc`. Represents and owned and reference counted type,
+    /// like Rust's [`std::rc::Rc`].
+    pub struct KjRc<T: Refcounted + ?Sized> {
+        own: Own<T>,
+    }
+
+    #[repr(C)]
+    /// Bindings to the kj type `kj::Arc`. Represents and owned and atomically reference
+    /// counted type, like Rust's [`std::sync::Arc`].
+    pub struct KjArc<T: AtomicRefcounted + ?Sized> {
+        own: Own<T>,
+    }
+
+    unsafe impl<T: AtomicRefcounted> Send for KjArc<T> where T: Send {}
+    unsafe impl<T: AtomicRefcounted> Sync for KjArc<T> where T: Sync {}
+
+    impl<T: Refcounted> KjRc<T> {
+        #[must_use]
+        pub fn get(&self) -> *const T {
+            self.own.as_ptr()
+        }
+    }
+
+    impl<T: AtomicRefcounted> KjArc<T> {
+        #[must_use]
+        pub fn get(&self) -> *const T {
+            self.own.as_ptr()
+        }
+
+        // The return value here represents exclusive access to the internal `Own`.
+        // This allows for exclusive mutation of the inner value.
+        pub fn get_mut(&mut self) -> Option<&mut Own<T>> {
+            if self.own.is_shared() {
+                None
+            } else {
+                Some(&mut self.own)
+            }
+        }
+    }
+
+    impl<T: Refcounted> Deref for KjRc<T> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            &self.own
+        }
+    }
+
+    impl<T: AtomicRefcounted> Deref for KjArc<T> {
+        type Target = T;
+
+        fn deref(&self) -> &Self::Target {
+            &self.own
+        }
+    }
+
+    // TODO: DerefMut? Is exposing mutable access to these types really the
+    // ideal behavior? We shouldn't expose traditional mutable access at
+    // least, we don't want to break ownership. Convert into Rust Rc? Idk.
+    // Maybe just making any mutable API unsafe is sufficient for now.
+
+    /// Using clone to create another count, like how Rust does it.
+    impl<T: Refcounted> Clone for KjRc<T> {
+        fn clone(&self) -> Self {
+            T::add_ref(self)
+        }
+    }
+
+    impl<T: AtomicRefcounted> Clone for KjArc<T> {
+        fn clone(&self) -> Self {
+            T::add_ref(self)
+        }
+    }
+
+    // No `Drop` needs to be implemented for `KjRc` or `KjArc`, because the
+    // internal `Own` `Drop` is sufficient.
+}

kj-rs/tests/BUILD.bazel

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

kj-rs/tests/lib.rs

@@ -8,6 +8,7 @@
 
 mod test_futures;
 mod test_own;
+mod test_refcount;
 
 use test_futures::{
     new_awaiting_future_i32, new_error_handling_future_void_infallible, new_errored_future_void,
@@ -72,6 +73,36 @@ mod ffi {
         fn rust_take_own_driver();
     }
 
+    unsafe extern "C++" {
+        include!("kj-rs-demo/test-refcount.h");
+
+        type OpaqueRefcountedClass;
+
+        #[allow(dead_code)]
+        fn get_rc() -> KjRc<OpaqueRefcountedClass>;
+        #[allow(dead_code)]
+        #[cxx_name = "getData"]
+        fn get_data(&self) -> u64;
+        #[allow(dead_code)]
+        #[cxx_name = "setData"]
+        fn set_data(self: Pin<&mut OpaqueRefcountedClass>, data: u64);
+    }
+
+    unsafe extern "C++" {
+        include!("kj-rs-demo/test-refcount.h");
+
+        type OpaqueAtomicRefcountedClass;
+
+        #[allow(dead_code)]
+        fn get_arc() -> KjArc<OpaqueAtomicRefcountedClass>;
+
+        #[allow(dead_code)]
+        #[cxx_name = "getData"]
+        fn get_data(&self) -> u64;
+        #[allow(dead_code)]
+        #[cxx_name = "setData"]
+        fn set_data(self: Pin<&mut OpaqueAtomicRefcountedClass>, data: u64);
+    }
     // Helper function to test moving `Own` to C++
     extern "Rust" {
         fn modify_own_return(cpp_own: Own<OpaqueCxxClass>) -> Own<OpaqueCxxClass>;
@@ -122,6 +153,9 @@ mod ffi {
     }
 }
 
+unsafe impl Send for ffi::OpaqueAtomicRefcountedClass {}
+unsafe impl Sync for ffi::OpaqueAtomicRefcountedClass {}
+
 pub fn modify_own_return(mut own: Own<ffi::OpaqueCxxClass>) -> Own<ffi::OpaqueCxxClass> {
     own.pin_mut().set_data(72);
     own

kj-rs/tests/test-refcount.c++

@@ -0,0 +1,11 @@
+#include "test-refcount.h"
+
+namespace kj_rs_demo {
+kj::Rc<OpaqueRefcountedClass> get_rc() {
+	return kj::rc<OpaqueRefcountedClass>(15);
+}
+
+kj::Arc<OpaqueAtomicRefcountedClass> get_arc() {
+	return kj::arc<OpaqueAtomicRefcountedClass>(16);
+}
+}