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
69extern "C-unwind" fn new_index(mut lua: LuaState) -> AmountReturned {
70	let i = lua.check_integer(2);
71	let new = lua.check_number(3);
72	let vec = check_vec_mut(&mut lua);
73
74	let Some(value) = vec.get_mut(i.cast_unsigned()) else { lua.error(c"invalid index"); };
75	*value = new;
76
77	lua.return_none()
78}
79
80extern "C-unwind" fn constructor(mut lua: LuaState) -> AmountReturned {
81	let mut vec = Vec::new();
82
83	// Initialize the vector if the called passed some numbers.
84	for i in 1..=lua.get_top() {
85		vec.push(lua.to_number(i));
86	}
87
88	// Push the vector to the stack as userdata.
89	lua.new_full_userdata::<Vec<f64>>(vec);
90
91	// Set the vector's metatable, initializing the metatable itself if's been first created.
92	if lua.new_userdata_metatable(NUM_VECTOR_TYPE_NAME) {
93		// Vectors allocate their own memory, so we need to have them be dropped when garbage collected.
94		// SAFETY: The metatable is being used by a valid full userdata of the same type.
95		unsafe { lua.set_drop_gc_metamethod::<Vec<f64>>(); }
96
97		// __tostring metamethod.
98		lua.push_cfunction(to_string);
99		lua.set_field(-2, c"__tostring");
100
101		// __index metamethod.
102		lua.push_cfunction(index);
103		lua.set_field(-2, c"__index");
104
105		// __newindex metamethod.
106		lua.push_cfunction(new_index);
107		lua.set_field(-2, c"__newindex");
108	}
109	lua.set_metatable(-2);
110
111	lua.return_any(1)
112}
113
114#[unsafe(no_mangle)]
115extern "C-unwind" fn luaopen_num_vector(lua: LuaState) -> AmountReturned {
116	lua.return_cfunction(constructor)
117}