Reentrant context for script functions

TitanNano · TitanNano · commit bb0fa8aafe79 · 2024-06-01T15:45:43.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -11,6 +11,7 @@ edition = "2021"
 
 [workspace.dependencies]
 godot = { git = "https://github.com/godot-rust/gdext", rev = "22fd33d4d5213a3fe5db9a58547888cebe35c647" }
+godot-cell = { git = "https://github.com/godot-rust/gdext", rev = "22fd33d4d5213a3fe5db9a58547888cebe35c647" }
 itertools = "0.10.3"
 rand = "0.8.5"
 darling = { version = "0.20.3" }
diff --git a/derive/src/impl_attribute.rs b/derive/src/impl_attribute.rs
@@ -9,7 +9,7 @@ use quote::{quote, quote_spanned, ToTokens};
 use syn::{parse_macro_input, spanned::Spanned, FnArg, ImplItem, ItemImpl, ReturnType, Type};
 
 use crate::{
-    rust_to_variant_type,
+    is_context_type, rust_to_variant_type,
     type_paths::{godot_types, property_hints, string_name_ty, variant_ty},
 };
 
@@ -56,29 +56,37 @@ pub fn godot_script_impl(
                     let arg_name = arg.pat.as_ref();
                     let arg_type = rust_to_variant_type(arg.ty.as_ref()).unwrap();
 
-                    (
-                        quote_spanned! {
-                            arg.span() =>
-                            ::godot_rust_script::RustScriptPropDesc {
-                                name: stringify!(#arg_name),
-                                ty: #arg_type,
-                                exported: false,
-                                hint: #property_hints::NONE,
-                                hint_string: "",
-                                description: "",
+                    is_context_type(arg.ty.as_ref()).then(|| {
+                        (
+                            quote!(),
+
+                            quote_spanned!(arg.span() => ctx,)
+                        )
+                    }).unwrap_or_else(|| {
+                        (
+                            quote_spanned! {
+                                arg.span() =>
+                                ::godot_rust_script::RustScriptPropDesc {
+                                    name: stringify!(#arg_name),
+                                    ty: #arg_type,
+                                    exported: false,
+                                    hint: #property_hints::NONE,
+                                    hint_string: "",
+                                    description: "",
+                                },
                             },
-                        },
-
-                        quote_spanned! {
-                            arg.span() =>
-                            #godot_types::prelude::FromGodot::try_from_variant(
-                                args.get(#index).ok_or(#godot_types::sys::GDEXTENSION_CALL_ERROR_TOO_FEW_ARGUMENTS)?
-                            ).map_err(|err| {
-                                #godot_types::log::godot_error!("failed to convert variant for argument {} of {}: {}", stringify!(#arg_name), #fn_name_str,  err);
-                                #godot_types::sys::GDEXTENSION_CALL_ERROR_INVALID_ARGUMENT
-                            })?
-                        }
-                    )
+
+                            quote_spanned! {
+                                arg.span() =>
+                                #godot_types::prelude::FromGodot::try_from_variant(
+                                    args.get(#index).ok_or(#godot_types::sys::GDEXTENSION_CALL_ERROR_TOO_FEW_ARGUMENTS)?
+                                ).map_err(|err| {
+                                    #godot_types::log::godot_error!("failed to convert variant for argument {} of {}: {}", stringify!(#arg_name), #fn_name_str,  err);
+                                    #godot_types::sys::GDEXTENSION_CALL_ERROR_INVALID_ARGUMENT
+                                })?,
+                            }
+                        )
+                    })
                 })
                 .collect();
 
@@ -143,7 +151,8 @@ pub fn godot_script_impl(
     let trait_impl = quote_spanned! {
         current_type.span() =>
         impl ::godot_rust_script::GodotScriptImpl for #current_type {
-            fn call_fn(&mut self, name: #string_name_ty, args: &[&#variant_ty]) -> ::std::result::Result<#variant_ty, #call_error_ty> {
+            #[allow(unused_variables)]
+            fn call_fn(&mut self, name: #string_name_ty, args: &[&#variant_ty], ctx: ::godot_rust_script::Context) -> ::std::result::Result<#variant_ty, #call_error_ty> {
                 match name.to_string().as_str() {
                     #method_dispatch
 
diff --git a/derive/src/lib.rs b/derive/src/lib.rs
@@ -132,8 +132,8 @@ pub fn derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
 
             #set_fields_impl
 
-            fn call(&mut self, name: #string_name_ty, args: &[&#variant_ty]) -> ::std::result::Result<#variant_ty, #call_error_ty> {
-                ::godot_rust_script::GodotScriptImpl::call_fn(self, name, args)
+            fn call(&mut self, name: #string_name_ty, args: &[&#variant_ty], ctx: ::godot_rust_script::Context) -> ::std::result::Result<#variant_ty, #call_error_ty> {
+                ::godot_rust_script::GodotScriptImpl::call_fn(self, name, args, ctx)
             }
 
             fn to_string(&self) -> String {
@@ -205,6 +205,14 @@ fn rust_to_variant_type(ty: &syn::Type) -> Result<TokenStream, TokenStream> {
     }
 }
 
+fn is_context_type(ty: &syn::Type) -> bool {
+    let syn::Type::Path(path) = ty else {
+        return false;
+    };
+
+    path.path.segments.last().map(|segment| segment.ident == "Context").unwrap_or(false)
+}
+
 fn derive_default_with_base(field_opts: &[FieldOpts]) -> TokenStream {
     let godot_types = godot_types();
     let fields: TokenStream = field_opts
diff --git a/rust-script/Cargo.toml b/rust-script/Cargo.toml
@@ -7,6 +7,7 @@ edition.workspace = true
 
 [dependencies]
 godot.workspace = true
+godot-cell.workspace = true
 
 itertools = { workspace = true, optional = true }
 rand = { workspace = true, optional = true }
diff --git a/rust-script/src/runtime/mod.rs b/rust-script/src/runtime/mod.rs
@@ -31,6 +31,8 @@ use crate::{
 
 use self::rust_script_language::RustScriptLanguage;
 
+pub use rust_script_instance::Context;
+
 #[macro_export]
 macro_rules! setup {
     ($lib_crate:tt) => {
diff --git a/rust-script/src/runtime/rust_script_instance.rs b/rust-script/src/runtime/rust_script_instance.rs
@@ -4,19 +4,21 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/.
  */
 
-use std::{collections::HashMap, rc::Rc};
+use core::panic;
+use std::{collections::HashMap, fmt::Debug, ops::DerefMut, pin::Pin};
 
 use godot::{
     builtin::meta::{MethodInfo, PropertyInfo},
     engine::{Script, ScriptInstance, SiMut},
     prelude::{GString, Gd, Object, StringName, Variant, VariantType},
 };
+use godot_cell::GdCell;
 
 use crate::script_registry::GodotScriptObject;
 
 use super::{rust_script::RustScript, rust_script_language::RustScriptLanguage, SCRIPT_REGISTRY};
 
-fn script_method_list(script: &Gd<RustScript>) -> Rc<[MethodInfo]> {
+fn script_method_list(script: &Gd<RustScript>) -> Box<[MethodInfo]> {
     let rs = script.bind();
     let class_name = rs.str_class_name();
 
@@ -25,7 +27,7 @@ fn script_method_list(script: &Gd<RustScript>) -> Rc<[MethodInfo]> {
         .expect("script registry is inaccessible")
         .get(&class_name)
         .map(|meta| meta.methods().iter().map(MethodInfo::from).collect())
-        .unwrap_or_else(|| Rc::new([]) as Rc<[MethodInfo]>);
+        .unwrap_or_else(|| Box::new([]) as Box<[MethodInfo]>);
 
     methods
 }
@@ -34,7 +36,7 @@ fn script_class_name(script: &Gd<RustScript>) -> GString {
     script.bind().get_class_name()
 }
 
-fn script_property_list(script: &Gd<RustScript>) -> Rc<[PropertyInfo]> {
+fn script_property_list(script: &Gd<RustScript>) -> Box<[PropertyInfo]> {
     let rs = script.bind();
     let class_name = rs.str_class_name();
 
@@ -43,18 +45,67 @@ fn script_property_list(script: &Gd<RustScript>) -> Rc<[PropertyInfo]> {
         .expect("script registry is inaccessible")
         .get(&class_name)
         .map(|meta| meta.properties().iter().map(PropertyInfo::from).collect())
-        .unwrap_or_else(|| Rc::new([]) as Rc<[PropertyInfo]>);
+        .unwrap_or_else(|| Box::new([]) as Box<[PropertyInfo]>);
 
     props
 }
 
+pub struct Context<'a> {
+    cell: Pin<&'a GdCell<Box<dyn GodotScriptObject>>>,
+    data_ptr: *mut Box<dyn GodotScriptObject>,
+}
+
+impl<'a> Debug for Context<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("Context { <Call Context> }")
+    }
+}
+
+impl<'a> Context<'a> {
+    unsafe fn new(
+        cell: Pin<&'a GdCell<Box<dyn GodotScriptObject>>>,
+        data_ptr: *mut Box<dyn GodotScriptObject>,
+    ) -> Self {
+        Self { cell, data_ptr }
+    }
+
+    pub fn reentrant_scope<T: GodotScriptObject + 'static, R>(
+        &mut self,
+        self_ref: &mut T,
+        cb: impl FnOnce() -> R,
+    ) -> R {
+        let known_ptr = unsafe {
+            let any = (*self.data_ptr).as_any_mut();
+
+            any.downcast_mut::<T>().unwrap() as *mut T
+        };
+
+        let self_ptr = self_ref as *mut _;
+
+        if known_ptr != self_ptr {
+            panic!("unable to create reentrant scope with unrelated self reference!");
+        }
+
+        let guard = self
+            .cell
+            .make_inaccessible(unsafe { &mut *self.data_ptr })
+            .unwrap();
+
+        let result = cb();
+
+        drop(guard);
+
+        result
+    }
+}
+
 pub(super) struct RustScriptInstance {
-    data: Box<dyn GodotScriptObject>,
+    data: Pin<Box<GdCell<Box<dyn GodotScriptObject>>>>,
 
     script: Gd<RustScript>,
     generic_script: Gd<Script>,
-    property_list: Rc<[PropertyInfo]>,
-    method_list: Rc<[MethodInfo]>,
+    property_list: Box<[PropertyInfo]>,
+    method_list: Box<[MethodInfo]>,
 }
 
 impl RustScriptInstance {
@@ -64,7 +115,7 @@ impl RustScriptInstance {
         script: Gd<RustScript>,
     ) -> Self {
         Self {
-            data,
+            data: GdCell::new(data),
             generic_script: script.clone().upcast(),
             property_list: script_property_list(&script),
             method_list: script_method_list(&script),
@@ -80,12 +131,17 @@ impl ScriptInstance for RustScriptInstance {
         script_class_name(&self.script)
     }
 
-    fn set_property(mut this: SiMut<Self>, name: StringName, value: &Variant) -> bool {
-        this.data.set(name, value.to_owned())
+    fn set_property(this: SiMut<Self>, name: StringName, value: &Variant) -> bool {
+        let cell_ref = this.data.as_ref();
+        let mut mut_data = cell_ref.borrow_mut().unwrap();
+
+        mut_data.set(name, value.to_owned())
     }
 
     fn get_property(&self, name: StringName) -> Option<Variant> {
-        self.data.get(name)
+        let guard = self.data.as_ref().borrow().unwrap();
+
+        guard.get(name)
     }
 
     fn get_property_list(&self) -> Vec<PropertyInfo> {
@@ -97,11 +153,18 @@ impl ScriptInstance for RustScriptInstance {
     }
 
     fn call(
-        mut this: SiMut<Self>,
+        this: SiMut<Self>,
         method: StringName,
         args: &[&Variant],
     ) -> Result<Variant, godot::sys::GDExtensionCallErrorType> {
-        this.data.call(method, args)
+        let cell = this.data.as_ref();
+        let mut data_guard = cell.borrow_mut().unwrap();
+        let data = data_guard.deref_mut();
+        let data_ptr = data as *mut _;
+
+        let context = unsafe { Context::new(cell, data_ptr) };
+
+        data.call(method, args, context)
     }
 
     fn get_script(&self) -> &Gd<Script> {
@@ -127,7 +190,8 @@ impl ScriptInstance for RustScriptInstance {
     }
 
     fn to_string(&self) -> GString {
-        self.data.to_string().into()
+        // self.data.to_string().into()
+        GString::new()
     }
 
     fn get_property_state(&self) -> Vec<(StringName, Variant)> {
@@ -166,8 +230,8 @@ pub(super) struct RustScriptPlaceholder {
     script: Gd<RustScript>,
     generic_script: Gd<Script>,
     properties: HashMap<StringName, Variant>,
-    property_list: Rc<[PropertyInfo]>,
-    method_list: Rc<[MethodInfo]>,
+    property_list: Box<[PropertyInfo]>,
+    method_list: Box<[MethodInfo]>,
 }
 
 impl RustScriptPlaceholder {
diff --git a/rust-script/src/script_registry.rs b/rust-script/src/script_registry.rs
@@ -4,7 +4,7 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/.
  */
 
-use std::{collections::HashMap, fmt::Debug, sync::Arc};
+use std::{any::Any, collections::HashMap, fmt::Debug, sync::Arc};
 
 use godot::{
     builtin::{
@@ -17,13 +17,16 @@ use godot::{
     sys::VariantType,
 };
 
+use crate::Context;
+
 pub trait GodotScript: Debug + GodotScriptImpl {
     fn set(&mut self, name: StringName, value: Variant) -> bool;
     fn get(&self, name: StringName) -> Option<Variant>;
     fn call(
         &mut self,
         method: StringName,
         args: &[&Variant],
+        context: Context,
     ) -> Result<Variant, godot::sys::GDExtensionCallErrorType>;
 
     fn to_string(&self) -> String;
@@ -37,6 +40,7 @@ pub trait GodotScriptImpl {
         &mut self,
         name: StringName,
         args: &[&Variant],
+        context: Context,
     ) -> Result<Variant, godot::sys::GDExtensionCallErrorType>;
 }
 
@@ -47,12 +51,15 @@ pub trait GodotScriptObject {
         &mut self,
         method: StringName,
         args: &[&Variant],
+        context: Context,
     ) -> Result<Variant, godot::sys::GDExtensionCallErrorType>;
     fn to_string(&self) -> String;
     fn property_state(&self) -> HashMap<StringName, Variant>;
+
+    fn as_any_mut(&mut self) -> &mut dyn Any;
 }
 
-impl<T: GodotScript> GodotScriptObject for T {
+impl<T: GodotScript + 'static> GodotScriptObject for T {
     fn set(&mut self, name: StringName, value: Variant) -> bool {
         GodotScript::set(self, name, value)
     }
@@ -65,8 +72,9 @@ impl<T: GodotScript> GodotScriptObject for T {
         &mut self,
         method: StringName,
         args: &[&Variant],
+        context: Context,
     ) -> Result<Variant, godot::sys::GDExtensionCallErrorType> {
-        GodotScript::call(self, method, args)
+        GodotScript::call(self, method, args, context)
     }
 
     fn to_string(&self) -> String {
@@ -76,6 +84,10 @@ impl<T: GodotScript> GodotScriptObject for T {
     fn property_state(&self) -> HashMap<StringName, Variant> {
         GodotScript::property_state(self)
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
 }
 
 #[derive(Debug, Clone)]
diff --git a/rust-script/tests/script_derive.rs b/rust-script/tests/script_derive.rs
