Skip to content

Commit

Permalink
Add Scope::create_any_userdata to create Lua objects from any non-s…
Browse files Browse the repository at this point in the history
…tatic Rust types.
  • Loading branch information
khvzak committed Nov 7, 2024
1 parent a7d0691 commit c7094d4
Show file tree
Hide file tree
Showing 2 changed files with 87 additions and 26 deletions.
86 changes: 68 additions & 18 deletions src/scope.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
let _sg = StackGuard::new(state);
check_stack(state, 3)?;

// // We don't write the data to the userdata until pushing the metatable
// We don't write the data to the userdata until pushing the metatable
let protect = !self.lua.unlikely_memory_error();
#[cfg(feature = "luau")]
let ud_ptr = {
Expand All @@ -194,29 +194,79 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
ffi::lua_setmetatable(state, -2);

let ud = AnyUserData(self.lua.pop_ref());
self.attach_destructor::<T>(&ud);

let destructor: DestructorCallback = Box::new(|rawlua, vref| {
let state = rawlua.state();
let _sg = StackGuard::new(state);
assert_stack(state, 2);
Ok(ud)
}
}

// Check that userdata is valid (very likely)
if rawlua.push_userdata_ref(&vref).is_err() {
return vec![];
}
/// Creates a Lua userdata object from a custom Rust type.
///
/// Since the Rust type is not required to be static and implement [`UserData`] trait,
/// you need to provide a function to register fields or methods for the object.
///
/// See also [`Scope::create_userdata`] for more details about non-static limitations.
pub fn create_any_userdata<T>(
&'scope self,
data: T,
register: impl FnOnce(&mut UserDataRegistry<T>),
) -> Result<AnyUserData>
where
T: 'env,
{
let state = self.lua.state();
let ud = unsafe {
let _sg = StackGuard::new(state);
check_stack(state, 3)?;

// Deregister metatable
let mt_ptr = get_metatable_ptr(state, -1);
rawlua.deregister_userdata_metatable(mt_ptr);
// We don't write the data to the userdata until pushing the metatable
let protect = !self.lua.unlikely_memory_error();
#[cfg(feature = "luau")]
let ud_ptr = {
let data = UserDataStorage::new_scoped(data);
util::push_userdata::<UserDataStorage<T>>(state, data, protect)?
};
#[cfg(not(feature = "luau"))]
let ud_ptr = util::push_uninit_userdata::<UserDataStorage<T>>(state, protect)?;

let ud = take_userdata::<UserDataStorage<T>>(state);
// Push the metatable and register it with no TypeId
let mut registry = UserDataRegistry::new_unique(ud_ptr as *mut _);
register(&mut registry);
self.lua.push_userdata_metatable(registry)?;
let mt_ptr = ffi::lua_topointer(state, -1);
self.lua.register_userdata_metatable(mt_ptr, None);

vec![Box::new(move || drop(ud))]
});
self.destructors.0.borrow_mut().push((ud.0.clone(), destructor));
// Write data to the pointer and attach metatable
#[cfg(not(feature = "luau"))]
std::ptr::write(ud_ptr, UserDataStorage::new_scoped(data));
ffi::lua_setmetatable(state, -2);

Ok(ud)
}
AnyUserData(self.lua.pop_ref())
};
self.attach_destructor::<T>(&ud);
Ok(ud)
}

fn attach_destructor<T: 'env>(&'scope self, ud: &AnyUserData) {
let destructor: DestructorCallback = Box::new(|rawlua, vref| unsafe {
let state = rawlua.state();
let _sg = StackGuard::new(state);
assert_stack(state, 2);

// Check that userdata is valid (very likely)
if rawlua.push_userdata_ref(&vref).is_err() {
return vec![];
}

// Deregister metatable
let mt_ptr = get_metatable_ptr(state, -1);
rawlua.deregister_userdata_metatable(mt_ptr);

let ud = take_userdata::<UserDataStorage<T>>(state);

vec![Box::new(move || drop(ud))]
});
self.destructors.0.borrow_mut().push((ud.0.clone(), destructor));
}

/// Adds a destructor function to be run when the scope ends.
Expand Down
27 changes: 19 additions & 8 deletions tests/scope.rs
Original file line number Diff line number Diff line change
Expand Up @@ -410,15 +410,26 @@ fn test_scope_userdata_ref_mut() -> Result<()> {
fn test_scope_any_userdata() -> Result<()> {
let lua = Lua::new();

lua.register_userdata_type::<StdString>(|reg| {
reg.add_meta_method("__tostring", |_, data, ()| Ok(data.clone()));
})?;
fn register(reg: &mut UserDataRegistry<&mut StdString>) {
reg.add_method_mut("push", |_, this, s: String| {
this.push_str(&s.to_str()?);
Ok(())
});
reg.add_meta_method("__tostring", |_, data, ()| Ok((*data).clone()));
}

let data = StdString::from("foo");
let mut data = StdString::from("foo");
lua.scope(|scope| {
let ud = scope.create_any_userdata_ref(&data)?;
let ud = scope.create_any_userdata(&mut data, register)?;
lua.globals().set("ud", ud)?;
lua.load("assert(tostring(ud) == 'foo')").exec()
lua.load(
r#"
assert(tostring(ud) == "foo")
ud:push("bar")
assert(tostring(ud) == "foobar")
"#,
)
.exec()
})?;

// Check that userdata is destructed
Expand Down Expand Up @@ -498,7 +509,7 @@ fn test_scope_destructors() -> Result<()> {
let ud = lua.create_any_userdata(arc_str.clone())?;
lua.scope(|scope| {
scope.add_destructor(|| {
assert!(ud.take::<Arc<StdString>>().is_ok());
assert!(ud.destroy().is_ok());
});
Ok(())
})?;
Expand All @@ -510,7 +521,7 @@ fn test_scope_destructors() -> Result<()> {
assert_eq!(arc_str.as_str(), "foo");
lua.scope(|scope| {
scope.add_destructor(|| {
assert!(ud.take::<Arc<StdString>>().is_err());
assert!(ud.destroy().is_err());
});
Ok(())
})
Expand Down

0 comments on commit c7094d4

Please sign in to comment.