Add Base::to_init_gd() to hand out early pointers

For manually-managed objects, the implementation is simple: just create a copy
of the internal Gd pointer.

For ref-counted objects, we need to deal with the problem that `Base<T>` stores
a weak Gd pointer, and thus dropping handed-out pointers will trigger object
destruction. To deal with this, we keep an extra strong pointer around, and
defer its destruction to the next frame.

It would be even better if we could immediately destroy the strong ref after
object construction, however we could only do so for Rust paths `new_gd()`,
`Gd::from_init_fn()`, not for Godot-side construction. Unfortunately, Godot
doesn't offer a hook to run code immediately after creation. The POSTINITIALIZE
notification needs to be dispatched by godot-rust, and thus suffers from the
same issue.

Author: Bromeon

godot-core/src/obj/base.rs

@@ -5,12 +5,45 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
+#[cfg(debug_assertions)]
+use std::cell::Cell;
+use std::cell::RefCell;
 use std::fmt::{Debug, Display, Formatter, Result as FmtResult};
 use std::mem::ManuallyDrop;
+use std::rc::Rc;
 
-use crate::obj::{Gd, GodotClass};
+use crate::builtin::{Callable, Variant};
+use crate::obj::{bounds, Gd, GodotClass};
 use crate::{classes, sys};
 
+/// Represents the initialization state of a `Base<T>` object.
+#[cfg(debug_assertions)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+enum InitState {
+    /// Object is being constructed (inside `I*::init()` or `Gd::from_init_fn()`).
+    ObjectConstructing,
+    /// Object construction is complete.
+    ObjectInitialized,
+    /// `ScriptInstance` context - always considered initialized (bypasses lifecycle checks).
+    Script,
+}
+
+#[cfg(debug_assertions)]
+macro_rules! base_from_obj {
+    ($obj:expr, $state:expr) => {
+        Base::from_obj($obj, $state)
+    };
+}
+
+#[cfg(not(debug_assertions))]
+macro_rules! base_from_obj {
+    ($obj:expr, $state:expr) => {
+        Base::from_obj($obj)
+    };
+}
+
+// ----------------------------------------------------------------------------------------------------------------------------------------------
+
 /// Restricted version of `Gd`, to hold the base instance inside a user's `GodotClass`.
 ///
 /// Behaves similarly to [`Gd`][crate::obj::Gd], but is more constrained. Cannot be constructed by the user.
@@ -35,6 +68,15 @@ pub struct Base<T: GodotClass> {
     // 1.   Gd<T>  -- triggers InstanceStorage destruction
     // 2.
     obj: ManuallyDrop<Gd<T>>,
+
+    /// Additional strong ref, needed to prevent destruction if [`Self::to_init_gd()`] is called on ref-counted objects.
+    extra_strong_ref: Rc<RefCell<Option<Gd<T>>>>,
+
+    /// Tracks the initialization state of this `Base<T>` in Debug mode.
+    ///
+    /// Rc allows to "copy-construct" the base from an existing one, while still affecting the user-instance through the original `Base<T>`.
+    #[cfg(debug_assertions)]
+    init_state: Rc<Cell<InitState>>,
 }
 
 impl<T: GodotClass> Base<T> {
@@ -47,7 +89,15 @@ impl<T: GodotClass> Base<T> {
     /// If `base` is destroyed while the returned `Base<T>` is in use, that constitutes a logic error, not a safety issue.
     pub(crate) unsafe fn from_base(base: &Base<T>) -> Base<T> {
         debug_assert!(base.obj.is_instance_valid());
-        Base::from_obj(Gd::from_obj_sys_weak(base.obj.obj_sys()))
+
+        let obj = Gd::from_obj_sys_weak(base.obj.obj_sys());
+
+        Self {
+            obj: ManuallyDrop::new(obj),
+            extra_strong_ref: Rc::clone(&base.extra_strong_ref), // Before user init(), no handing out of Gd pointers occurs.
+            #[cfg(debug_assertions)]
+            init_state: Rc::clone(&base.init_state),
+        }
     }
 
     /// Create base from existing object (used in script instances).
@@ -57,9 +107,11 @@ impl<T: GodotClass> Base<T> {
     /// # Safety
     /// `gd` must be alive at the time of invocation. If it is destroyed while the returned `Base<T>` is in use, that constitutes a logic
     /// error, not a safety issue.
-    pub(crate) unsafe fn from_gd(gd: &Gd<T>) -> Self {
+    pub(crate) unsafe fn from_script_gd(gd: &Gd<T>) -> Self {
         debug_assert!(gd.is_instance_valid());
-        Base::from_obj(Gd::from_obj_sys_weak(gd.obj_sys()))
+
+        let obj = Gd::from_obj_sys_weak(gd.obj_sys());
+        base_from_obj!(obj, InitState::Script)
     }
 
     /// Create new base from raw Godot object.
@@ -72,20 +124,31 @@ impl<T: GodotClass> Base<T> {
     pub(crate) unsafe fn from_sys(base_ptr: sys::GDExtensionObjectPtr) -> Self {
         assert!(!base_ptr.is_null(), "instance base is null pointer");
 
-        // Initialize only as weak pointer (don't increment reference count)
+        // Initialize only as weak pointer (don't increment reference count).
         let obj = Gd::from_obj_sys_weak(base_ptr);
 
         // This obj does not contribute to the strong count, otherwise we create a reference cycle:
         // 1. RefCounted (dropped in GDScript)
         // 2. holds user T (via extension instance and storage)
         // 3. holds Base<T> RefCounted (last ref, dropped in T destructor, but T is never destroyed because this ref keeps storage alive)
         // Note that if late-init never happened on self, we have the same behavior (still a raw pointer instead of weak Gd)
-        Base::from_obj(obj)
+        base_from_obj!(obj, InitState::ObjectConstructing)
     }
 
+    #[cfg(debug_assertions)]
+    fn from_obj(obj: Gd<T>, init_state: InitState) -> Self {
+        Self {
+            obj: ManuallyDrop::new(obj),
+            extra_strong_ref: Rc::new(RefCell::new(None)),
+            init_state: Rc::new(Cell::new(init_state)),
+        }
+    }
+
+    #[cfg(not(debug_assertions))]
     fn from_obj(obj: Gd<T>) -> Self {
         Self {
             obj: ManuallyDrop::new(obj),
+            extra_strong_ref: Rc::new(RefCell::new(None)),
         }
     }
 
@@ -94,10 +157,112 @@ impl<T: GodotClass> Base<T> {
     /// Using this method to call methods on the base field of a Rust object is discouraged, instead use the
     /// methods from [`WithBaseField`](super::WithBaseField) when possible.
     #[doc(hidden)]
+    #[deprecated = "Private API. Use `Base::to_init_gd()` or `WithBaseField::to_gd()` instead."] // TODO(v0.4): remove.
     pub fn to_gd(&self) -> Gd<T> {
         (*self.obj).clone()
     }
 
+    /// Returns a [`Gd`] referencing the base object, for exclusive use during object initialization.
+    ///
+    /// Can be used during an initialization function [`I*::init()`][crate::classes::IObject::init] or [`Gd::from_init_fn()`].
+    ///
+    /// The base pointer is only pointing to a base object; you cannot yet downcast it to the object being constructed.
+    /// The instance ID is the same as the one the in-construction object will have.
+    ///
+    /// # Lifecycle for ref-counted classes
+    /// If `T: Inherits<RefCounted>`, then the ref-counted object is not yet fully-initialized at the time of the `init` function running.
+    /// Accessing the base object without further measures would be dangerous. Here, godot-rust employs a workaround: the `Base` object (which
+    /// holds a weak pointer to the actual instance) is temporarily upgraded to a strong pointer, preventing use-after-free.
+    ///
+    /// This additional reference is automatically dropped at an implementation-defined point in time (which may change, and technically delay
+    /// destruction of your object as soon as you use `Base::to_init_gd()`). Right now, this refcount-decrement is deferred to the next frame.
+    ///
+    /// For now, ref-counted bases can only use `to_init_gd()` on the main thread.
+    ///
+    /// # Panics (Debug)
+    /// If called outside an initialization function, or for ref-counted objects on a non-main thread.
+    #[cfg(since_api = "4.2")]
+    pub fn to_init_gd(&self) -> Gd<T> {
+        #[cfg(debug_assertions)] // debug_assert! still checks existence of symbols.
+        assert!(
+            self.is_initializing(),
+            "Base::to_init_gd() can only be called during object initialization, inside I*::init() or Gd::from_init_fn()"
+        );
+
+        // For manually-managed objects, regular clone is fine.
+        // Only static type matters, because this happens immediately after initialization, so T is both static and dynamic type.
+        if !<T::Memory as bounds::Memory>::IS_REF_COUNTED {
+            return Gd::clone(&self.obj);
+        }
+
+        debug_assert!(
+            sys::is_main_thread(),
+            "Base::to_init_gd() can only be called on the main thread for ref-counted objects (for now)"
+        );
+
+        // First time handing out a Gd<T>, we need to take measures to temporarily upgrade the Base's weak pointer to a strong one.
+        // During the initialization phase (derived object being constructed), increment refcount by 1.
+        if self.extra_strong_ref.borrow().is_none() {
+            let strong_ref = unsafe { Gd::from_obj_sys(self.obj.obj_sys()) };
+            *self.extra_strong_ref.borrow_mut() = Some(strong_ref);
+        }
+
+        // Can't use Gd::apply_deferred(), as that implicitly borrows &mut self, causing a "destroyed while bind was active" panic.
+        let name = format!("Base<{}> deferred unref", T::class_name());
+        let rc = Rc::clone(&self.extra_strong_ref);
+        let callable = Callable::from_once_fn(&name, move |_args| {
+            Self::drop_strong_ref(rc);
+            Ok(Variant::nil())
+        });
+        callable.call_deferred(&[]);
+
+        (*self.obj).clone()
+    }
+
+    /// Drops any extra strong references, possibly causing object destruction.
+    fn drop_strong_ref(extra_strong_ref: Rc<RefCell<Option<Gd<T>>>>) {
+        let mut r = extra_strong_ref.borrow_mut();
+        assert!(r.is_some());
+
+        *r = None; // Triggers RawGd::drop() -> dec-ref -> possibly object destruction.
+    }
+
+    /// Finalizes the initialization of this `Base<T>` and returns whether
+    pub(crate) fn mark_initialized(&mut self) -> bool {
+        #[cfg(debug_assertions)]
+        {
+            assert_eq!(
+                self.init_state.get(),
+                InitState::ObjectConstructing,
+                "Base<T> is already initialized, or holds a script instance"
+            );
+
+            self.init_state.set(InitState::ObjectInitialized);
+        }
+
+        self.extra_strong_ref.borrow().is_some()
+    }
+
+    /// Returns a [`Gd`] referencing the base object, assuming the derived object is fully constructed.
+    #[doc(hidden)]
+    pub fn __fully_constructed_gd(&self) -> Gd<T> {
+        #[cfg(debug_assertions)] // debug_assert! still checks existence of symbols.
+        assert!(
+            !self.is_initializing(),
+            "WithBaseField::to_gd(), base(), base_mut() can only be called on fully-constructed objects, after I*::init() or Gd::from_init_fn()"
+        );
+
+        (*self.obj).clone()
+    }
+
+    /// Returns a [`Gd`] referencing the base object, for use in script contexts only.
+    #[doc(hidden)]
+    pub fn __script_gd(&self) -> Gd<T> {
+        // Used internally by `SiMut::base()` and `SiMut::base_mut()` for script re-entrancy.
+        // Could maybe add debug validation to ensure script context in the future.
+        (*self.obj).clone()
+    }
+
     // Currently only used in outbound virtual calls (for scripts); search for: base_field(self).obj_sys().
     #[doc(hidden)]
     pub fn obj_sys(&self) -> sys::GDExtensionObjectPtr {
@@ -109,6 +274,36 @@ impl<T: GodotClass> Base<T> {
     pub(crate) fn debug_instance_id(&self) -> crate::obj::InstanceId {
         self.obj.instance_id()
     }
+
+    /// Returns a [`Gd`] referencing the base object, for use in script contexts only.
+    pub(crate) fn to_script_gd(&self) -> Gd<T> {
+        #[cfg(debug_assertions)]
+        assert_eq!(
+            self.init_state.get(),
+            InitState::Script,
+            "to_script_gd() can only be called on script-context Base objects"
+        );
+
+        (*self.obj).clone()
+    }
+
+    /// Returns `true` if this `Base<T>` is currently in the initializing state.
+    #[cfg(debug_assertions)]
+    fn is_initializing(&self) -> bool {
+        self.init_state.get() == InitState::ObjectConstructing
+    }
+
+    /// Returns a [`Gd`] referencing the base object, assuming the derived object is fully constructed.
+    #[doc(hidden)]
+    pub fn __constructed_gd(&self) -> Gd<T> {
+        #[cfg(debug_assertions)] // debug_assert! still checks existence of symbols.
+        assert!(
+            !self.is_initializing(),
+            "WithBaseField::to_gd(), base(), base_mut() can only be called on fully-constructed objects, after I*::init() or Gd::from_init_fn()"
+        );
+
+        (*self.obj).clone()
+    }
 }
 
 impl<T: GodotClass> Debug for Base<T> {

godot-core/src/obj/call_deferred.rs

@@ -20,6 +20,8 @@ use crate::registry::signal::ToSignalObj;
 #[cfg(before_api = "4.2")]
 pub trait WithDeferredCall<T: GodotClass> {}
 
+// TODO(v0.4): seal this and similar traits.
+
 /// Enables `Gd::apply_deferred()` for type-safe deferred calls.
 ///
 /// The trait is automatically available for all engine-defined Godot classes and user classes containing a `Base<T>` field.

godot-core/src/obj/casts.rs

@@ -48,6 +48,7 @@ impl<T: GodotClass, U: GodotClass> CastSuccess<T, U> {
     }
 
     /// Access shared reference to destination, without consuming object.
+    #[cfg(debug_assertions)]
     pub fn as_dest_ref(&self) -> &RawGd<U> {
         self.check_validity();
         &self.dest

godot-core/src/obj/script.rs

@@ -319,7 +319,7 @@ pub unsafe fn create_script_instance<T: ScriptInstance>(
         method_lists: BoundedPtrList::new(),
         // SAFETY: The script instance is always freed before the base object is destroyed. The weak reference should therefore never be
         // accessed after it has been freed.
-        base: unsafe { Base::from_gd(&for_object) },
+        base: unsafe { Base::from_script_gd(&for_object) },
     };
 
     let data_ptr = Box::into_raw(Box::new(data));
@@ -454,7 +454,7 @@ impl<'a, T: ScriptInstance> SiMut<'a, T> {
     /// }
     /// ```
     pub fn base(&self) -> ScriptBaseRef<'_, T> {
-        ScriptBaseRef::new(self.base_ref.to_gd(), self.mut_ref)
+        ScriptBaseRef::new(self.base_ref.to_script_gd(), self.mut_ref)
     }
 
     /// Returns a mutable reference suitable for calling engine methods on this object.
@@ -518,7 +518,7 @@ impl<'a, T: ScriptInstance> SiMut<'a, T> {
     pub fn base_mut(&mut self) -> ScriptBaseMut<'_, T> {
         let guard = self.cell.make_inaccessible(self.mut_ref).unwrap();
 
-        ScriptBaseMut::new(self.base_ref.to_gd(), guard)
+        ScriptBaseMut::new(self.base_ref.to_script_gd(), guard)
     }
 }
 

godot-core/src/obj/traits.rs

@@ -355,9 +355,13 @@ pub trait WithBaseField: GodotClass + Bounds<Declarer = bounds::DeclUser> {
     ///
     /// This is intended to be stored or passed to engine methods. You cannot call `bind()` or `bind_mut()` on it, while the method
     /// calling `to_gd()` is still running; that would lead to a double borrow panic.
+    ///
+    /// # Panics
+    /// If called during initialization (the `init()` function or `Gd::from_init_fn()`). Use [`Base::to_init_gd()`] instead.
     fn to_gd(&self) -> Gd<Self>;
 
     /// Returns a reference to the `Base` stored by this object.
+    #[doc(hidden)]
     fn base_field(&self) -> &Base<Self::Base>;
 
     /// Returns a shared reference suitable for calling engine methods on this object.
@@ -419,7 +423,7 @@ pub trait WithBaseField: GodotClass + Bounds<Declarer = bounds::DeclUser> {
     ///
     /// For this, use [`base_mut()`](WithBaseField::base_mut()) instead.
     fn base(&self) -> BaseRef<'_, Self> {
-        let gd = self.base_field().to_gd();
+        let gd = self.base_field().__constructed_gd();
 
         BaseRef::new(gd, self)
     }
@@ -490,7 +494,7 @@ pub trait WithBaseField: GodotClass + Bounds<Declarer = bounds::DeclUser> {
     /// ```
     #[allow(clippy::let_unit_value)]
     fn base_mut(&mut self) -> BaseMut<'_, Self> {
-        let base_gd = self.base_field().to_gd();
+        let base_gd = self.base_field().__constructed_gd();
 
         let gd = self.to_gd();
         // SAFETY:

godot-core/src/registry/callbacks.rs

@@ -124,7 +124,6 @@ where
     F: FnOnce(Base<T::Base>) -> T,
 {
     let class_name = T::class_name();
-
     //out!("create callback: {}", class_name.backing);
 
     let base = unsafe { Base::from_sys(base_ptr) };
@@ -133,12 +132,15 @@ where
     let context = || format!("panic during {class_name}::init() constructor");
     let code = || make_user_instance(unsafe { Base::from_base(&base) });
     let user_instance = handle_panic(context, std::panic::AssertUnwindSafe(code))?;
+
     // Print shouldn't be necessary as panic itself is printed. If this changes, re-enable in error case:
     // godot_error!("failed to create instance of {class_name}; Rust init() panicked");
 
+    let mut base_copy = unsafe { Base::from_base(&base) };
+
     let instance = InstanceStorage::<T>::construct(user_instance, base);
-    let instance_ptr = instance.into_raw();
-    let instance_ptr = instance_ptr as sys::GDExtensionClassInstancePtr;
+    let instance_rust_ptr = instance.into_raw();
+    let instance_ptr = instance_rust_ptr as sys::GDExtensionClassInstancePtr;
 
     let binding_data_callbacks = crate::storage::nop_instance_callbacks();
     unsafe {
@@ -151,7 +153,14 @@ where
         );
     }
 
-    // std::mem::forget(class_name);
+    // Mark initialization as complete, now that user constructor has finished.
+    if base_copy.mark_initialized() {
+        // If an extra RefCounted reference was handed out, notify storage about it, so it can unreference later.
+        let instance_ref = unsafe { &*instance_rust_ptr };
+        instance_ref.mark_surplus_ref();
+    }
+
+    // No std::mem::forget(base_copy) here, since Base may stores other fields that need deallocation.
     Ok(instance_ptr)
 }