Merge pull request #1270 from godot-rust/qol/refcount-tracking

Properly register GDExtension `reference`/`unreference` callbacks

Author: Bromeon

godot-core/src/builtin/rect2i.rs

@@ -83,17 +83,12 @@ impl Rect2i {
     }
 
     /// The end of the `Rect2i` calculated as `position + size`.
-    ///
-    /// _Godot equivalent: `Rect2i.size` property_
-    #[doc(alias = "size")]
     #[inline]
     pub const fn end(self) -> Vector2i {
         Vector2i::new(self.position.x + self.size.x, self.position.y + self.size.y)
     }
 
-    /// Set size based on desired end-point.
-    ///
-    /// _Godot equivalent: `Rect2i.size` property_
+    /// Set end of the `Rect2i`, updating `size = end - position`.
     #[inline]
     pub fn set_end(&mut self, end: Vector2i) {
         self.size = end - self.position
@@ -110,8 +105,7 @@ impl Rect2i {
 
     /// Returns `true` if this `Rect2i` completely encloses another one.
     ///
-    /// Any `Rect2i` encloses itself, i.e. an enclosed `Rect2i` does is not required to be a
-    /// proper sub-rect.
+    /// Any `Rect2i` encloses itself, i.e. an enclosed `Rect2i` does is not required to be a proper sub-rect.
     #[inline]
     pub const fn encloses(self, other: Self) -> bool {
         self.assert_nonnegative();

godot-core/src/obj/bounds.rs

@@ -152,7 +152,11 @@ pub use crate::implement_godot_bounds;
 // Memory bounds
 
 /// Specifies the memory strategy of the static type.
-pub trait Memory: Sealed {}
+pub trait Memory: Sealed {
+    /// True for everything inheriting `RefCounted`, false for `Object` and all other classes.
+    #[doc(hidden)]
+    const IS_REF_COUNTED: bool;
+}
 
 /// Specifies the memory strategy of the dynamic type.
 ///
@@ -205,7 +209,9 @@ pub trait DynMemory: Sealed {
 /// This is used for `RefCounted` classes and derived.
 pub struct MemRefCounted {}
 impl Sealed for MemRefCounted {}
-impl Memory for MemRefCounted {}
+impl Memory for MemRefCounted {
+    const IS_REF_COUNTED: bool = true;
+}
 impl DynMemory for MemRefCounted {
     fn maybe_init_ref<T: GodotClass>(obj: &mut RawGd<T>) {
         out!("  MemRefc::init:  {obj:?}");
@@ -327,7 +333,9 @@ impl DynMemory for MemDynamic {
 /// This is used for all `Object` derivates, which are not `RefCounted`. `Object` itself is also excluded.
 pub struct MemManual {}
 impl Sealed for MemManual {}
-impl Memory for MemManual {}
+impl Memory for MemManual {
+    const IS_REF_COUNTED: bool = false;
+}
 impl DynMemory for MemManual {
     fn maybe_init_ref<T: GodotClass>(_obj: &mut RawGd<T>) {}
     fn maybe_inc_ref<T: GodotClass>(_obj: &mut RawGd<T>) {}

godot-core/src/obj/gd.rs

@@ -294,10 +294,11 @@ impl<T: GodotClass> Gd<T> {
 
     /// Returns the dynamic class name of the object as `StringName`.
     ///
-    /// This method retrieves the class name of the object at runtime, which can be different from [`T::class_name()`] if derived
-    /// classes are involved.
+    /// This method retrieves the class name of the object at runtime, which can be different from [`T::class_name()`][GodotClass::class_name]
+    /// if derived classes are involved.
     ///
-    /// Unlike [`Object::get_class()`], this returns `StringName` instead of `GString` and needs no `Inherits<Object>` bound.
+    /// Unlike [`Object::get_class()`][crate::classes::Object::get_class], this returns `StringName` instead of `GString` and needs no
+    /// `Inherits<Object>` bound.
     pub(crate) fn dynamic_class_string(&self) -> StringName {
         unsafe {
             StringName::new_with_string_uninit(|ptr| {

godot-core/src/registry/class.rs

@@ -136,8 +136,8 @@ impl ClassRegistrationInfo {
 }
 
 /// Registers a class with static type information.
-// Currently dead code, but will be needed for builder API. Don't remove.
-pub fn register_class<
+#[expect(dead_code)] // Will be needed for builder API. Don't remove.
+pub(crate) fn register_class<
     T: cap::GodotDefault
         + cap::ImplementsGodotVirtual
         + cap::GodotToString
@@ -427,6 +427,8 @@ fn fill_class_info(item: PluginItem, c: &mut ClassRegistrationInfo) {
             is_instantiable,
             #[cfg(all(since_api = "4.3", feature = "register-docs"))]
                 docs: _,
+            reference_fn,
+            unreference_fn,
         }) => {
             c.parent_class_name = Some(base_class_name);
             c.default_virtual_fn = default_get_virtual_fn;
@@ -443,6 +445,8 @@ fn fill_class_info(item: PluginItem, c: &mut ClassRegistrationInfo) {
             // See also: https://github.com/godotengine/godot/pull/58972
             c.godot_params.is_abstract = sys::conv::bool_to_sys(!is_instantiable);
             c.godot_params.free_instance_func = Some(free_fn);
+            c.godot_params.reference_func = reference_fn;
+            c.godot_params.unreference_func = unreference_fn;
 
             fill_into(
                 &mut c.godot_params.create_instance_func,

godot-core/src/registry/plugin.rs

@@ -167,6 +167,12 @@ pub struct Struct {
     /// Callback to library-generated function which registers properties in the `struct` definition.
     pub(crate) register_properties_fn: ErasedRegisterFn,
 
+    /// Callback on refc-increment. Only for `RefCounted` classes.
+    pub(crate) reference_fn: sys::GDExtensionClassReference,
+
+    /// Callback on refc-decrement. Only for `RefCounted` classes.
+    pub(crate) unreference_fn: sys::GDExtensionClassUnreference,
+
     /// Function called by Godot when an object of this class is freed.
     ///
     /// Always implemented as [`callbacks::free`].
@@ -200,6 +206,8 @@ impl Struct {
     pub fn new<T: GodotClass + cap::ImplementsGodotExports>(
         #[cfg(all(since_api = "4.3", feature = "register-docs"))] docs: StructDocs,
     ) -> Self {
+        let refcounted = <T::Memory as bounds::Memory>::IS_REF_COUNTED;
+
         Self {
             base_class_name: T::Base::class_name(),
             generated_create_fn: None,
@@ -215,6 +223,9 @@ impl Struct {
             is_instantiable: false,
             #[cfg(all(since_api = "4.3", feature = "register-docs"))]
             docs,
+            // While Godot doesn't do anything with these callbacks for non-RefCounted classes, we can avoid instantiating them in Rust.
+            reference_fn: refcounted.then_some(callbacks::reference::<T>),
+            unreference_fn: refcounted.then_some(callbacks::unreference::<T>),
         }
     }
 

godot-core/src/storage/instance_storage.rs

@@ -154,9 +154,6 @@ pub unsafe trait Storage {
 
 /// An internal trait for keeping track of reference counts for a storage.
 pub(crate) trait StorageRefCounted: Storage {
-    #[allow(dead_code)] // used in `debug-log` feature. Don't micromanage.
-    fn godot_ref_count(&self) -> u32;
-
     fn on_inc_ref(&self);
 
     fn on_dec_ref(&self);

godot-core/src/storage/mod.rs

@@ -65,24 +65,22 @@ mod log_active {
 
     pub fn log_construct<T: GodotClass>(base: &Base<T::Base>) {
         out!(
-            "    Storage::construct:             {base:?}  (T={ty})",
+            "    Storage::construct:   {base:?}  (T={ty})",
             ty = type_name::<T>()
         );
     }
 
     pub fn log_inc_ref<T: StorageRefCounted>(storage: &T) {
         out!(
-            "    Storage::on_inc_ref (rc={rc}):  {base:?}  (T={ty})",
-            rc = T::godot_ref_count(storage),
+            "    Storage::on_inc_ref:  {base:?}  (T={ty})",
             base = storage.base(),
             ty = type_name::<T>(),
         );
     }
 
     pub fn log_dec_ref<T: StorageRefCounted>(storage: &T) {
         out!(
-            "  | Storage::on_dec_ref (rc={rc}):  {base:?}  (T={ty})",
-            rc = T::godot_ref_count(storage),
+            "  | Storage::on_dec_ref:  {base:?}  (T={ty})",
             base = storage.base(),
             ty = type_name::<T>(),
         );
@@ -103,8 +101,7 @@ mod log_active {
         // Do not Debug-fmt `self.base()` object here, see above.
 
         out!(
-            "    Storage::drop (rc={rc}):        {base_id}",
-            rc = storage.godot_ref_count(),
+            "    Storage::drop:        {base_id}",
             base_id = storage.base().debug_instance_id(),
         );
     }

godot-core/src/storage/multi_threaded.rs

@@ -5,8 +5,6 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
-use std::sync::atomic::{AtomicU32, Ordering};
-
 #[cfg(not(feature = "experimental-threads"))]
 use godot_cell::panicking::{GdCell, InaccessibleGuard, MutGuard, RefGuard};
 
@@ -22,7 +20,6 @@ pub struct InstanceStorage<T: GodotClass> {
 
     // Declared after `user_instance`, is dropped last
     pub(super) lifecycle: AtomicLifecycle,
-    godot_ref_count: AtomicU32,
 
     // No-op in Release mode.
     borrow_tracker: DebugBorrowTracker,
@@ -49,7 +46,6 @@ unsafe impl<T: GodotClass> Storage for InstanceStorage<T> {
             user_instance: GdCell::new(user_instance),
             base,
             lifecycle: AtomicLifecycle::new(Lifecycle::Alive),
-            godot_ref_count: AtomicU32::new(1),
             borrow_tracker: DebugBorrowTracker::new(),
         }
     }
@@ -104,19 +100,11 @@ unsafe impl<T: GodotClass> Storage for InstanceStorage<T> {
 }
 
 impl<T: GodotClass> StorageRefCounted for InstanceStorage<T> {
-    fn godot_ref_count(&self) -> u32 {
-        self.godot_ref_count.load(Ordering::Relaxed)
-    }
-
     fn on_inc_ref(&self) {
-        self.godot_ref_count.fetch_add(1, Ordering::Relaxed);
-
         super::log_inc_ref(self);
     }
 
     fn on_dec_ref(&self) {
-        self.godot_ref_count.fetch_sub(1, Ordering::Relaxed);
-
         super::log_dec_ref(self);
     }
 }

godot-core/src/storage/single_threaded.rs

@@ -22,7 +22,6 @@ pub struct InstanceStorage<T: GodotClass> {
 
     // Declared after `user_instance`, is dropped last
     pub(super) lifecycle: cell::Cell<Lifecycle>,
-    godot_ref_count: cell::Cell<u32>,
 
     // No-op in Release mode.
     borrow_tracker: DebugBorrowTracker,
@@ -49,7 +48,6 @@ unsafe impl<T: GodotClass> Storage for InstanceStorage<T> {
             user_instance: GdCell::new(user_instance),
             base,
             lifecycle: cell::Cell::new(Lifecycle::Alive),
-            godot_ref_count: cell::Cell::new(1),
             borrow_tracker: DebugBorrowTracker::new(),
         }
     }
@@ -101,21 +99,11 @@ unsafe impl<T: GodotClass> Storage for InstanceStorage<T> {
 }
 
 impl<T: GodotClass> StorageRefCounted for InstanceStorage<T> {
-    fn godot_ref_count(&self) -> u32 {
-        self.godot_ref_count.get()
-    }
-
     fn on_inc_ref(&self) {
-        let refc = self.godot_ref_count.get() + 1;
-        self.godot_ref_count.set(refc);
-
         super::log_inc_ref(self);
     }
 
     fn on_dec_ref(&self) {
-        let refc = self.godot_ref_count.get() - 1;
-        self.godot_ref_count.set(refc);
-
         super::log_dec_ref(self);
     }
 }

godot-macros/src/itest.rs

@@ -64,10 +64,9 @@ pub fn attribute_itest(input_item: venial::Item) -> ParseResult<TokenStream> {
         quote! { __unused_context: &crate::framework::TestContext }
     };
 
+    let return_ty = func.return_ty.as_ref();
     if is_async
-        && func
-            .return_ty
-            .as_ref()
+        && return_ty
             .and_then(extract_typename)
             .is_none_or(|segment| segment.ident != "TaskHandle")
     {
@@ -78,7 +77,8 @@ pub fn attribute_itest(input_item: venial::Item) -> ParseResult<TokenStream> {
 
     let (return_tokens, test_case_ty, plugin_name);
     if is_async {
-        return_tokens = quote! { -> TaskHandle };
+        let [arrow, arrow_head] = func.tk_return_arrow.unwrap();
+        return_tokens = quote! { #arrow #arrow_head #return_ty }; // retain span.
         test_case_ty = quote! { crate::framework::AsyncRustTestCase };
         plugin_name = ident("__GODOT_ASYNC_ITEST");
     } else {