1use std::{cell::RefCell, collections::HashMap, future::Future, mem::take, pin::Pin, rc::Rc, task::{Context, Poll}};
2
3use r2g_mlua::prelude::*;
4use super::from_lua_clone_impl;
5use crate::core::{get_state, get_state_with_rwlock, get_task_scheduler_from_lua, FastFlag, LuauState, ParallelDispatch, Trc, TrcReadLock, TrcWriteLock, Weak};
6pub type ManagedRBXScriptSignal = Trc<RBXScriptSignal>;
7
8#[derive(Debug, Clone)]
9pub struct RBXScriptConnection {
10 id: usize,
11 signal: ManagedRBXScriptSignal
12}
13#[derive(Debug, Clone)]
14struct SignalCallback {
15 func: LuaFunction,
16 state: Trc<LuauState>,
17 once: bool,
18 parallel: ParallelDispatch
19}
20#[derive(Debug)]
21pub struct RBXScriptSignal {
22 callbacks: HashMap<usize, SignalCallback>,
23 this_ptr: Option<Weak<RBXScriptSignal>>,
24 id: usize
25}
26struct InnerRBXScriptSignalFuture {
27 event: ManagedRBXScriptSignal,
28 lua: Lua,
29 resolved: bool,
30 waiting: bool,
31 values: LuaMultiValue
32}
33pub struct RBXScriptSignalFuture {
34 future: Rc<RefCell<InnerRBXScriptSignalFuture>>
35}
36
37impl RBXScriptSignal {
38 pub fn new() -> Trc<RBXScriptSignal> {
39 Trc::new_cyclic(|x| RBXScriptSignal { callbacks: HashMap::default(), this_ptr: Some(x.clone()), id: 0 })
40 }
41 pub fn connect(&mut self, lua: &Lua, func: LuaFunction, parallel: ParallelDispatch) -> LuaResult<RBXScriptConnection> {
42 let id = self.id;
43 self.id += 1;
44 self.callbacks.insert(id, SignalCallback {
45 func,
46 state: get_state_with_rwlock(lua).clone(),
47 once: false,
48 parallel: parallel
49 });
50 Ok(RBXScriptConnection {
51 id,
52 signal: self.this_ptr.as_ref().unwrap().upgrade().unwrap()
53 })
54 }
55 #[inline]
56 pub fn connect_parallel(&mut self, lua: &Lua, func: LuaFunction) -> LuaResult<RBXScriptConnection> {
57 self.connect(lua, func, ParallelDispatch::Desynchronized)
58 }
59 pub fn once(&mut self, lua: &Lua, func: LuaFunction, parallel: ParallelDispatch) -> LuaResult<RBXScriptConnection> {
60 let id = self.id;
61 self.id += 1;
62 self.callbacks.insert(id, SignalCallback {
63 func,
64 state: get_state_with_rwlock(lua).clone(),
65 once: true,
66 parallel: parallel
67 });
68 Ok(RBXScriptConnection {
69 id,
70 signal: self.this_ptr.as_ref().unwrap().upgrade().unwrap()
71 })
72 }
73 pub fn fire(mut self: TrcWriteLock<'_, RBXScriptSignal>, lua: &Lua, args: impl IntoLuaMulti) -> LuaResult<()> {
74 let args = args.into_lua_multi(lua)?;
75 let mut to_remove = Vec::new();
76 let task = get_task_scheduler_from_lua(unsafe {(lua as *const Lua).as_ref().unwrap_unchecked()});
77 let callbacks_clone = self.callbacks.clone();
78 let release = self.guard_release();
79 let fire_behavior = get_state(lua).flags().get_int(FastFlag::SignalBehavior);
80 for (id, callback) in callbacks_clone {
81 let _ = match fire_behavior {
82 0 | 1 | 3 => task.spawn_func(lua, callback.func, args.clone()),
83 2 => task.defer_func(lua, callback.func, args.clone(), callback.parallel),
84 _ => unreachable!()
85 };
86 if callback.once {
87 to_remove.push(id);
88 }
89 }
90 drop(release);
91 if !to_remove.is_empty() {
92 for i in to_remove {
93 self.callbacks.remove(&i);
94 }
95 }
96 Ok(())
97 }
98 pub fn fire_ancestry(mut self: TrcWriteLock<'_, RBXScriptSignal>, lua: &Lua, args: impl IntoLuaMulti) -> LuaResult<()> {
99 let args = args.into_lua_multi(lua)?;
100 let mut to_remove = Vec::new();
101 let task = get_task_scheduler_from_lua(unsafe {(lua as *const Lua).as_ref().unwrap_unchecked()});
102 let callbacks_clone = self.callbacks.clone();
103 let release = self.guard_release();
104 let fire_behavior = get_state(lua).flags().get_int(FastFlag::SignalBehavior);
105 for (id, callback) in callbacks_clone {
106 let _ = match fire_behavior {
107 0 | 1 => task.spawn_func(lua, callback.func, args.clone()),
108 2 | 3 => task.defer_func(lua, callback.func, args.clone(), callback.parallel),
109 _ => unreachable!()
110 };
111 if callback.once {
112 to_remove.push(id);
113 }
114 }
115 drop(release);
116 if !to_remove.is_empty() {
117 for i in to_remove {
118 self.callbacks.remove(&i);
119 }
120 }
121 Ok(())
122 }
123 pub fn wait<'a>(self: TrcReadLock<'_, Self>, lua: &'a Lua) -> RBXScriptSignalFuture {
124 RBXScriptSignalFuture {
125 future: Rc::new(RefCell::new(InnerRBXScriptSignalFuture {
126 event: self.this_ptr.as_ref().unwrap().upgrade().unwrap(),
127 lua: lua.clone(),
128 resolved: false,
129 waiting: false,
130 values: LuaMultiValue::new()
131 }))
132 }
133 }
134}
135impl RBXScriptConnection {
136 pub fn is_connected(&self) -> bool {
137 self.signal.read().callbacks.contains_key(&self.id)
138 }
139 pub fn disconnect(&self) {
140 let mut lock = self.signal.write();
141 lock.callbacks.remove(&self.id);
142 }
143}
144
145impl Future for RBXScriptSignalFuture {
146 type Output = LuaResult<LuaMultiValue>;
147
148 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
149 let immut_borrow = self.future.borrow();
150 if immut_borrow.resolved {
151 drop(immut_borrow);
152 let mut borrow = self.future.borrow_mut();
153 borrow.resolved = false;
154 borrow.waiting = false;
155 let values = take(&mut borrow.values);
156 Poll::Ready(Ok(values))
157 } else {
158 if !immut_borrow.waiting {
159 let waker = cx.waker().clone();
160 let clone = self.future.clone();
161 let func = immut_borrow.lua.create_function_mut(move |_, mv: LuaMultiValue| {
162 let mut borrow = clone.borrow_mut();
163 borrow.resolved = true;
164 borrow.waiting = false;
165 borrow.values = mv;
166 waker.clone().wake();
167 Ok(())
168 })?;
169 immut_borrow.event.write().once(&immut_borrow.lua, func, ParallelDispatch::Default)?;
170 }
171 Poll::Pending
172 }
173 }
174}
175impl LuaUserData for ManagedRBXScriptSignal {
176 fn add_methods<M: LuaUserDataMethods<Self>>(methods: &mut M) {
177 methods.add_method_mut("Connect", |lua, this, func: LuaFunction| {
178 this.write().connect(lua, func, ParallelDispatch::Synchronized)
179 });
180 methods.add_method_mut("ConnectParallel", |lua, this, func: LuaFunction| {
181 this.write().connect_parallel(lua, func)
182 });
183 methods.add_method_mut("Once", |lua, this, func: LuaFunction| {
184 this.write().once(lua, func, ParallelDispatch::Synchronized)
185 });
186 methods.add_async_method_mut("Wait", async |lua, this, ()| {
187 this.read().wait(&lua).await
188 });
189 }
190}
191impl LuaUserData for RBXScriptConnection {
192 fn add_methods<M: LuaUserDataMethods<Self>>(methods: &mut M) {
193 methods.add_method_mut("Disconnect",
194 |_, this, ()|
195 Ok(this.disconnect())
196 );
197 }
198 fn add_fields<F: LuaUserDataFields<Self>>(fields: &mut F) {
199 fields.add_field_method_get("Connected",
200 |_, this|
201 Ok(this.is_connected())
202 );
203 }
204}
205
206from_lua_clone_impl!(RBXScriptConnection);
207from_lua_clone_impl!(ManagedRBXScriptSignal);