Skip to main content

num_vector/
num_vector.rs

1use std::ffi::CStr;
2use low_luajit::{AmountReturned, LuaState, LuaType, reg};
3
4const NUM_VECTOR_TYPE_NAME: &CStr =  c"numerical vector";
5
6fn check_vec(lua: &LuaState) -> &Vec<f64> {
7	// SAFETY: "numerical vector" userdata should always be valid and of this type, and there should be no other mutable reference to it.
8	unsafe { lua.check_full_userdata::<Vec<f64>>(1, NUM_VECTOR_TYPE_NAME) }
9}
10
11fn check_vec_mut(lua: &mut LuaState) -> &mut Vec<f64> {
12	// SAFETY: "numerical vector" userdata should always be valid and of this type, and there should be no other reference to it.
13	unsafe { lua.check_full_userdata::<Vec<f64>>(1, NUM_VECTOR_TYPE_NAME) }
14}
15
16extern "C-unwind" fn to_string(lua: LuaState) -> AmountReturned {
17	let vec = check_vec(&lua);
18
19	let string = format!("{vec:?}");
20	lua.return_string(string)
21}
22
23extern "C-unwind" fn push(mut lua: LuaState) -> AmountReturned {
24	let n = lua.check_number(2);
25	let vec = check_vec_mut(&mut lua);
26
27	vec.push(n);
28
29	lua.return_none()
30}
31
32extern "C-unwind" fn pop(mut lua: LuaState) -> AmountReturned {
33	let vec = check_vec_mut(&mut lua);
34
35	match vec.pop() {
36		Some(n) => lua.return_number(n),
37		None => lua.return_nil(),
38	}
39}
40
41extern "C-unwind" fn index(mut lua: LuaState) -> AmountReturned {
42	let vec = check_vec(&lua);
43
44	match lua.type_of(2) {
45		// The index is a string, so it's probably trying to access the vector's methods.
46		// Let's quickly generate the methods' table and get the requested method.
47		// (It would be better to cache it somewhere, but for this example this is fine enough.)
48		LuaType::String => {
49			// We can create the table as a "module" for convenience. This way we can quickly add more methods if needed.
50			lua.new_module(reg![push, pop]);
51
52			// The key needs to be the top element, so we move the table to the 2nd element instead.
53			lua.insert(2);
54
55			// Index the method table.
56			lua.get_table(2);
57		},
58
59		// If the index is just a number, attempt to index the vector itself
60		LuaType::Number if let Some(n) = vec.get(lua.to_integer(2).cast_unsigned()) => lua.push_number(*n),
61
62		// Invalid index, return nil.
63		_ => lua.push_nil(),
64	}
65
66	lua.return_any(1)
67}
68
69#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
70extern "C-unwind" fn new_index(mut lua: LuaState) -> AmountReturned {
71	let i = lua.check_number(2);
72	let new = lua.check_number(3);
73	let vec = check_vec_mut(&mut lua);
74
75	let Some(value) = vec.get_mut(i as usize) else { lua.error(c"invalid index"); };
76	*value = new;
77
78	lua.return_none()
79}
80
81extern "C-unwind" fn constructor(mut lua: LuaState) -> AmountReturned {
82	let mut vec = Vec::new();
83
84	// Initialize the vector if the called passed some numbers.
85	for i in 1..=lua.get_top() {
86		vec.push(lua.to_number(i));
87	}
88
89	// Push the vector to the stack as userdata.
90	lua.new_full_userdata::<Vec<f64>>(vec);
91
92	// Set the vector's metatable, initializing the metatable itself if's been first created.
93	if lua.new_userdata_metatable(NUM_VECTOR_TYPE_NAME) {
94		// Vectors allocate their own memory, so we need to have them be dropped when garbage collected.
95		// SAFETY: The metatable is being used by a valid full userdata of the same type.
96		unsafe { lua.set_drop_gc_metamethod::<Vec<f64>>(); }
97
98		// __tostring metamethod.
99		lua.push_cfunction(to_string);
100		lua.set_field(-2, c"__tostring");
101
102		// __index metamethod.
103		lua.push_cfunction(index);
104		lua.set_field(-2, c"__index");
105
106		// __newindex metamethod.
107		lua.push_cfunction(new_index);
108		lua.set_field(-2, c"__newindex");
109	}
110	lua.set_metatable(-2);
111
112	lua.return_any(1)
113}
114
115#[unsafe(no_mangle)]
116extern "C-unwind" fn luaopen_num_vector(lua: LuaState) -> AmountReturned {
117	lua.return_cfunction(constructor)
118}