[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,
+        "bool cxxbridge1$kj_rs$rc${}$is_shared(::kj::Refcounted *ptr) noexcept {{",
+        instance,
+    );
+    writeln!(out, "  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,
+        "bool cxxbridge1$kj_rs$arc${}$is_shared(::kj::AtomicRefcounted *ptr) noexcept {{",
+        instance,
+    );
+    writeln!(out, "  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>,
     }
@@ -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,72 @@
+//! 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)]
+    pub struct KjRc<T: Refcounted + ?Sized> {
+        own: Own<T>,
+    }
+
+    // TODO: `Send` and `Sync`
+    #[repr(C)]
+    pub struct KjArc<T: AtomicRefcounted + ?Sized> {
+        own: Own<T>,
+    }
+
+    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()
+        }
+    }
+
+    impl<T: Refcounted> Deref for KjRc<T> {
+        type Target = Own<T>;
+
+        fn deref(&self) -> &Self::Target {
+            &self.own
+        }
+    }
+
+    impl<T: AtomicRefcounted> Deref for KjArc<T> {
+        type Target = Own<T>;
+
+        fn deref(&self) -> &Self::Target {
+            &self.own
+        }
+    }
+
+    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)
+        }
+    }
+}

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,29 @@ 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;
+    }
+
+    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;
+    }
     // Helper function to test moving `Own` to C++
     extern "Rust" {
         fn modify_own_return(cpp_own: Own<OpaqueCxxClass>) -> Own<OpaqueCxxClass>;

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);
+}
+}

kj-rs/tests/test-refcount.h

@@ -0,0 +1,42 @@
+#pragma once
+
+#include "kj/refcount.h"
+
+#include <cstdint>
+
+namespace kj_rs_demo {
+
+class OpaqueRefcountedClass: public kj::Refcounted {
+ public:
+  OpaqueRefcountedClass(uint64_t data): data(data) {}
+  ~OpaqueRefcountedClass() {}
+  uint64_t getData() const {
+    return this->data;
+  }
+  void setData(uint64_t val) {
+    this->data = val;
+  }
+
+ private:
+  uint64_t data;
+};
+
+class OpaqueAtomicRefcountedClass: public kj::AtomicRefcounted {
+ public:
+  OpaqueAtomicRefcountedClass(uint64_t data): data(data) {}
+  ~OpaqueAtomicRefcountedClass() {}
+  uint64_t getData() const {
+    return this->data;
+  }
+  void setData(uint64_t val) {
+    this->data = val;
+  }
+
+ private:
+  uint64_t data;
+};
+
+kj::Rc<OpaqueRefcountedClass> get_rc();
+kj::Arc<OpaqueAtomicRefcountedClass> get_arc();
+
+}  // namespace kj_rs_demo

kj-rs/tests/test_futures.rs

@@ -149,10 +149,7 @@ pub async fn new_layered_ready_future_void() -> Result<()> {
 }
 
 // From example at https://doc.rust-lang.org/std/future/fn.poll_fn.html#capturing-a-pinned-state
-async fn naive_select<T>(
-    a: impl Future<Output = T>,
-    b: impl Future<Output = T>,
-) -> T {
+async fn naive_select<T>(a: impl Future<Output = T>, b: impl Future<Output = T>) -> T {
     let (mut a, mut b) = (pin!(a), pin!(b));
     future::poll_fn(move |cx| {
         if let Poll::Ready(r) = a.as_mut().poll(cx) {