Skip to main content

luau_vm/libs/
coroutine.rs

1use crate::handle::RawHandle;
2use crate::{VmControl, VmExit, VmResult};
3
4use crate::native::{NativeCallContext, NativeCallResult, NativeFunction};
5use crate::state::{
6    THREAD_STATUS_BREAK, THREAD_STATUS_ERR_ERR, THREAD_STATUS_ERR_MEM, THREAD_STATUS_OK,
7    THREAD_STATUS_SCHEDULED_REENTRY, THREAD_STATUS_YIELD, ThreadState,
8};
9use crate::thread::{LUA_COERR, LUA_COFIN, LUA_COSUS, Thread};
10use crate::types::LUA_TFUNCTION;
11
12static CO_STATUS_NAMES: [&str; 5] = ["running", "suspended", "normal", "dead", "dead"];
13
14static CO_FUNCS: [NativeFunction; 7] = [
15    NativeFunction {
16        name: "create",
17        function: coroutine_create,
18    },
19    NativeFunction {
20        name: "running",
21        function: coroutine_running,
22    },
23    NativeFunction {
24        name: "status",
25        function: coroutine_status,
26    },
27    NativeFunction {
28        name: "wrap",
29        function: coroutine_wrap,
30    },
31    NativeFunction {
32        name: "yield",
33        function: coroutine_yield,
34    },
35    NativeFunction {
36        name: "isyieldable",
37        function: coroutine_is_yieldable,
38    },
39    NativeFunction {
40        name: "close",
41        function: coroutine_close,
42    },
43];
44
45enum CoroutineResume {
46    Values(usize),
47    Error,
48}
49
50/// `costatus`
51fn coroutine_status(ctx: NativeCallContext) -> NativeCallResult {
52    let thread = ctx.raw_thread();
53    unsafe {
54        let co = thread.to_thread(1);
55        ctx.arg(1).expected(co.is_some(), "thread")?;
56        let co = co.unwrap_unchecked();
57
58        thread.push_string(CO_STATUS_NAMES[thread.co_status(&co) as usize])?;
59        Ok(1)
60    }
61}
62
63/// `auxresume`
64unsafe fn aux_resume(thread: &Thread, co: &Thread, narg: i32) -> VmResult<CoroutineResume> {
65    unsafe {
66        if co.as_ptr().as_ref().unwrap_unchecked().status != THREAD_STATUS_YIELD {
67            let status = thread.co_status(co);
68            if status != LUA_COSUS {
69                let _ = crate::push_fstring!(
70                    thread,
71                    "cannot resume %s coroutine",
72                    CO_STATUS_NAMES[status as usize]
73                );
74                return Ok(CoroutineResume::Error);
75            }
76        }
77
78        if narg != 0 {
79            if co.check_stack(narg) == 0 {
80                return crate::error!(thread, "too many arguments to resume").map_err(Into::into);
81            }
82            thread.x_move(co, narg)?;
83        } else if co.stack_top().offset_from(co.stack_base()) as i32
84            > crate::thread::LUAI_MAX_C_STACK
85        {
86            return crate::error!(thread, "too many arguments to resume").map_err(Into::into);
87        }
88
89        co.as_ptr().as_mut().unwrap_unchecked().single_step =
90            thread.as_ptr().as_ref().unwrap_unchecked().single_step;
91
92        match co.resume(Some(thread), narg) {
93            Ok(()) | Err(VmExit::Control(VmControl::Yield)) => {
94                debug_assert_ne!(
95                    co.as_ptr().as_ref().unwrap_unchecked().status,
96                    THREAD_STATUS_SCHEDULED_REENTRY
97                );
98
99                let nres = co.stack_top().offset_from(co.stack_base()) as i32;
100                if nres != 0 {
101                    if nres + 1 > crate::thread::LUA_MIN_STACK as i32
102                        && thread.check_stack(nres + 1) == 0
103                    {
104                        return crate::error!(thread, "too many results to resume")
105                            .map_err(Into::into);
106                    }
107                    co.x_move(thread, nres)?;
108                }
109                Ok(CoroutineResume::Values(nres as usize))
110            }
111            Err(exit) => {
112                let VmExit::Error(_) = exit else {
113                    return Err(exit);
114                };
115                co.x_move(thread, 1)?;
116                Ok(CoroutineResume::Error)
117            }
118        }
119    }
120}
121
122/// `interruptThread`
123unsafe fn interrupt_thread(thread: &Thread, co: &Thread) -> NativeCallResult {
124    unsafe {
125        let global = thread.global();
126        if let Some(hook) = global.debug_interrupt_callback() {
127            thread.call_hook(
128                |thread, debug| hook(thread, debug).map_err(Into::into),
129                co.as_ptr().cast(),
130            )?;
131        }
132
133        thread.break_current()
134    }
135}
136
137/// `auxresumecont`
138unsafe fn aux_resume_cont(thread: &Thread, co: &Thread) -> VmResult<CoroutineResume> {
139    unsafe {
140        let status = co.as_ptr().as_ref().unwrap_unchecked().status;
141        if matches!(status, x if x == THREAD_STATUS_OK || x == THREAD_STATUS_YIELD) {
142            let nres = co.stack_top().offset_from(co.stack_base()) as i32;
143            if thread.check_stack(nres + 1) == 0 {
144                return crate::error!(thread, "too many results to resume").map_err(Into::into);
145            }
146            co.x_move(thread, nres)?;
147            Ok(CoroutineResume::Values(nres as usize))
148        } else {
149            thread.raw_check_stack(2)?;
150            co.x_move(thread, 1)?;
151            Ok(CoroutineResume::Error)
152        }
153    }
154}
155
156/// `coresumefinish`
157unsafe fn coroutine_resume_finish(thread: &Thread, result: CoroutineResume) -> NativeCallResult {
158    unsafe {
159        match result {
160            CoroutineResume::Values(count) => {
161                let count_i32 = i32::try_from(count).expect("coroutine result count fits in i32");
162                thread.push_boolean(1)?;
163                thread.insert(-(count_i32 + 1));
164                Ok(count + 1)
165            }
166            CoroutineResume::Error => {
167                thread.push_boolean(0)?;
168                thread.insert(-2);
169                Ok(2)
170            }
171        }
172    }
173}
174
175/// `auxwrapfinish`
176unsafe fn aux_wrap_finish(thread: &Thread, result: CoroutineResume) -> NativeCallResult {
177    unsafe {
178        match result {
179            CoroutineResume::Values(count) => Ok(count),
180            CoroutineResume::Error => {
181                if thread.is_string(-1) != 0 {
182                    thread.push_where(1)?;
183                    thread.insert(-2);
184                    thread.concat(2)?;
185                }
186                thread.error().map_err(Into::into)
187            }
188        }
189    }
190}
191
192/// `coresumey`
193fn coroutine_resume_yieldable(ctx: NativeCallContext) -> NativeCallResult {
194    let thread = ctx.raw_thread();
195    unsafe {
196        let co = thread.to_thread(1);
197        ctx.arg(1).expected(co.is_some(), "thread")?;
198        let co = co.unwrap_unchecked();
199        let narg = thread.get_top() - 1;
200        let result = match aux_resume(thread, &co, narg) {
201            Ok(result) => result,
202            Err(VmExit::Control(VmControl::Break)) => return interrupt_thread(thread, &co),
203            Err(exit) => return Err(exit),
204        };
205
206        coroutine_resume_finish(thread, result)
207    }
208}
209
210/// `coresumecont`
211fn coroutine_resume_cont(ctx: NativeCallContext, _status: i32) -> NativeCallResult {
212    let thread = ctx.raw_thread();
213    unsafe {
214        let co = thread.to_thread(1);
215        ctx.arg(1).expected(co.is_some(), "thread")?;
216        let co = co.unwrap_unchecked();
217
218        if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_BREAK {
219            return interrupt_thread(thread, &co);
220        }
221
222        let result = aux_resume_cont(thread, &co)?;
223        coroutine_resume_finish(thread, result)
224    }
225}
226
227/// `auxwrapy`
228fn coroutine_wrap_yieldable(ctx: NativeCallContext) -> NativeCallResult {
229    let thread = ctx.raw_thread();
230    unsafe {
231        let co = thread
232            .to_thread(crate::thread::upvalue_index(1))
233            .unwrap_unchecked();
234        let result = match aux_resume(thread, &co, thread.get_top()) {
235            Ok(result) => result,
236            Err(VmExit::Control(VmControl::Break)) => return interrupt_thread(thread, &co),
237            Err(exit) => return Err(exit),
238        };
239
240        aux_wrap_finish(thread, result)
241    }
242}
243
244/// `auxwrapcont`
245fn coroutine_wrap_cont(ctx: NativeCallContext, _status: i32) -> NativeCallResult {
246    let thread = ctx.raw_thread();
247    unsafe {
248        let co = thread
249            .to_thread(crate::thread::upvalue_index(1))
250            .unwrap_unchecked();
251
252        if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_BREAK {
253            return interrupt_thread(thread, &co);
254        }
255
256        let result = aux_resume_cont(thread, &co)?;
257        aux_wrap_finish(thread, result)
258    }
259}
260
261/// `cocreate`
262fn coroutine_create(ctx: NativeCallContext) -> NativeCallResult {
263    let thread = ctx.raw_thread();
264    unsafe {
265        thread.check_type(1, LUA_TFUNCTION)?;
266        let new_thread = thread.new_thread()?;
267        thread.x_push(&new_thread, 1)?;
268        Ok(1)
269    }
270}
271
272/// `cowrap`
273fn coroutine_wrap(ctx: NativeCallContext) -> NativeCallResult {
274    let thread = ctx.raw_thread();
275    coroutine_create(NativeCallContext::new(thread))?;
276    unsafe {
277        thread.push_native_closure_k(
278            coroutine_wrap_yieldable,
279            None,
280            1,
281            Some(coroutine_wrap_cont),
282        )?
283    };
284    Ok(1)
285}
286
287/// `coyield`
288fn coroutine_yield(ctx: NativeCallContext) -> NativeCallResult {
289    let thread = ctx.raw_thread();
290    let nres = unsafe { thread.get_top() };
291    unsafe { thread.yield_current(nres) }
292}
293
294/// `corunning`
295fn coroutine_running(ctx: NativeCallContext) -> NativeCallResult {
296    let thread = ctx.raw_thread();
297    if unsafe { thread.push_thread()? } != 0 {
298        unsafe { thread.push_nil()? };
299    }
300    Ok(1)
301}
302
303/// `coyieldable`
304fn coroutine_is_yieldable(ctx: NativeCallContext) -> NativeCallResult {
305    let thread = ctx.raw_thread();
306    unsafe { thread.push_boolean(thread.is_yieldable())? };
307    Ok(1)
308}
309
310/// `coclose`
311fn coroutine_close(ctx: NativeCallContext) -> NativeCallResult {
312    let thread = ctx.raw_thread();
313    unsafe {
314        let co = thread.to_thread(1);
315        ctx.arg(1).expected(co.is_some(), "thread")?;
316        let co = co.unwrap_unchecked();
317
318        let status = thread.co_status(&co);
319        if status != LUA_COFIN && status != LUA_COERR && status != LUA_COSUS {
320            return crate::error!(
321                thread,
322                "cannot close %s coroutine",
323                CO_STATUS_NAMES[status as usize]
324            )
325            .map_err(Into::into);
326        }
327
328        if matches!(
329            co.as_ptr().as_ref().unwrap_unchecked().status,
330            x if x == THREAD_STATUS_OK || x == THREAD_STATUS_YIELD
331        ) {
332            thread.push_boolean(1)?;
333            co.reset()?;
334            Ok(1)
335        } else {
336            thread.push_boolean(0)?;
337
338            if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_ERR_MEM {
339                thread.push_string(crate::state::LUA_MEMERRMSG)?;
340            } else if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_ERR_ERR {
341                thread.push_string(crate::state::LUA_ERRERRMSG)?;
342            } else if co.get_top() != 0 {
343                co.x_move(thread, 1)?;
344            }
345
346            co.reset()?;
347            Ok(2)
348        }
349    }
350}
351
352impl Thread {
353    /// `luaopen_coroutine`
354    pub unsafe fn open_coroutine(&self) -> NativeCallResult {
355        unsafe { self.register(Some(super::LUA_COLIB_NAME), &CO_FUNCS[..])? };
356        unsafe {
357            self.push_native_closure_k(
358                coroutine_resume_yieldable,
359                Some("resume"),
360                0,
361                Some(coroutine_resume_cont),
362            )?;
363            self.raw_set_field(-2, "resume")?;
364        }
365        Ok(1)
366    }
367}