1use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
19use std::sync::{Arc, Mutex};
20
21use crate::state_stub::{lua_CFunction, upvalue_index, LuaState, LuaStateStubExt as _};
22use lua_types::{error::LuaError, gc::GcRef, value::LuaValue, LuaStatus, LuaThreadClose, LuaType};
23
24const COS_RUN: i32 = 0;
28
29const COS_DEAD: i32 = 1;
31
32const COS_YIELD: i32 = 2;
34
35const COS_NORM: i32 = 3;
37
38const STAT_NAMES: [&[u8]; 4] = [b"running", b"dead", b"suspended", b"normal"];
42
43pub const CO_FUNCS: &[(&[u8], lua_CFunction)] = &[
51 (b"create", co_create),
52 (b"resume", co_resume),
53 (b"running", co_running),
54 (b"status", co_status),
55 (b"wrap", co_wrap),
56 (b"yield", co_yield),
57 (b"isyieldable", co_isyieldable),
58 (b"close", co_close),
59];
60
61fn get_co(state: &mut LuaState) -> Result<GcRef<lua_types::value::LuaThread>, LuaError> {
67 let co = state.to_thread(1);
68 if co.is_none() {
69 let got = state.arg(1);
70 return Err(LuaError::type_arg_error(1, "thread", &got));
71 }
72 Ok(co.expect("checked above"))
73}
74
75fn get_opt_co(state: &mut LuaState) -> Result<GcRef<lua_types::value::LuaThread>, LuaError> {
76 if matches!(state.global().lua_version, lua_types::LuaVersion::V55)
77 && state.type_at(1) == LuaType::None
78 {
79 let id = state.global().current_thread_id;
80 return state
81 .global()
82 .thread_value_for(id)
83 .ok_or_else(|| LuaError::runtime(format_args!("current thread is not registered")));
84 }
85 get_co(state)
86}
87
88fn aux_status(state: &mut LuaState, co: &GcRef<lua_types::value::LuaThread>) -> i32 {
100 let co_id = co.id;
101 let entry_rc = {
102 let g = state.global();
103 if co_id == g.current_thread_id {
104 return COS_RUN;
105 }
106 if co_id == g.main_thread_id {
107 return COS_NORM;
108 }
109 match g.threads.get(&co_id) {
110 Some(e) => e.state.clone(),
111 None => return COS_DEAD,
112 }
113 };
114 let co_state = match entry_rc.try_borrow() {
115 Ok(state) => state,
116 Err(_) => {
117 return COS_NORM;
122 }
123 };
124 let raw_status = co_state.status;
125 if raw_status == LuaStatus::Yield as u8 {
126 return COS_YIELD;
127 }
128 if raw_status != LuaStatus::Ok as u8 {
129 return COS_DEAD;
130 }
131 let has_frames = co_state.ci.as_usize() > 0;
132 if has_frames {
133 return COS_NORM;
134 }
135 let ci_func = co_state.call_info[0].func.0;
136 let top = co_state.top.0;
137 let lua_gettop = top as i64 - ci_func as i64 - 1;
138 if lua_gettop == 0 {
139 COS_DEAD
140 } else {
141 COS_YIELD
142 }
143}
144
145fn aux_resume(state: &mut LuaState, co: GcRef<lua_types::value::LuaThread>, narg: i32) -> i32 {
162 let co_id = co.id;
163 let entry_rc = {
164 let g = state.global();
165 match g.threads.get(&co_id) {
166 Some(e) => e.state.clone(),
167 None => {
168 drop(g);
169 push_lit_or_nil(state, b"cannot resume dead coroutine");
170 return -1;
171 }
172 }
173 };
174 let parent_thread_id = state.global().current_thread_id;
175 let top_before = state.get_top();
176 if top_before < narg {
177 push_lit_or_nil(state, b"not enough arguments to resume");
178 return -1;
179 }
180 let first_arg_idx = top_before - narg + 1;
181 let args: Vec<LuaValue> = (first_arg_idx..=top_before)
182 .map(|i| state.value_at(i))
183 .collect();
184 lua_vm::api::set_top(state, (top_before - narg) as i32).ok();
185
186 let parent_open_upval_slots: Vec<(u64, lua_vm::state::StackIdx)> = state
187 .openupval
188 .iter()
189 .filter_map(|uv| match &*uv.slot() {
190 lua_types::UpValState::Open { thread_id, idx } => Some((*thread_id as u64, *idx)),
191 lua_types::UpValState::Closed(_) => None,
192 })
193 .collect();
194 {
195 let mut g = state.global_mut();
196 for (tid, idx) in &parent_open_upval_slots {
197 let val = state.get_at(*idx);
198 g.cross_thread_upvals.insert((*tid, *idx), val);
199 }
200 }
201
202 push_parent_gc_snapshot(state);
203
204 let (status, results_or_err): (LuaStatus, Vec<LuaValue>) = {
205 let mut co_state = match entry_rc.try_borrow_mut() {
206 Ok(b) => b,
207 Err(_) => {
208 pop_parent_gc_snapshot(state);
209 let mut g = state.global_mut();
210 for (tid, idx) in &parent_open_upval_slots {
211 g.cross_thread_upvals.remove(&(*tid, *idx));
212 }
213 drop(g);
214 push_lit_or_nil(state, b"cannot resume non-suspended coroutine");
215 return -1;
216 }
217 };
218 if co_state.check_stack(narg + 1).is_err() {
219 drop(co_state);
220 pop_parent_gc_snapshot(state);
221 let mut g = state.global_mut();
222 for (tid, idx) in &parent_open_upval_slots {
223 g.cross_thread_upvals.remove(&(*tid, *idx));
224 }
225 drop(g);
226 push_lit_or_nil(state, b"too many arguments to resume");
227 return -1;
228 }
229 for v in args {
230 co_state.push(v);
231 }
232 co_state.global_mut().current_thread_id = co_id;
233 let mut nres: i32 = 0;
234 let previous_hook = Arc::new(Mutex::new(Some(std::panic::take_hook())));
235 let previous_for_hook = Arc::clone(&previous_hook);
236 std::panic::set_hook(Box::new(move |info| {
237 if info.payload().downcast_ref::<LuaThreadClose>().is_none() {
238 if let Ok(guard) = previous_for_hook.lock() {
239 if let Some(hook) = guard.as_ref() {
240 hook(info);
241 }
242 }
243 }
244 }));
245 let resume_result = catch_unwind(AssertUnwindSafe(|| {
246 lua_vm::do_::lua_resume(&mut *co_state, Some(state), narg, &mut nres)
247 }));
248 let _installed_hook = std::panic::take_hook();
249 if let Some(hook) = previous_hook.lock().ok().and_then(|mut h| h.take()) {
250 std::panic::set_hook(hook);
251 }
252 co_state.global_mut().current_thread_id = parent_thread_id;
253 let status = match resume_result {
254 Ok(status) => status,
255 Err(payload) => {
256 if let Some(close) = payload.downcast_ref::<LuaThreadClose>() {
257 close.0
258 } else {
259 resume_unwind(payload);
260 }
261 }
262 };
263 let co_top = co_state.top_idx().0 as i32;
264 let ci_func = co_state.current_call_info().func.0 as i32;
265 let count = if status == LuaStatus::Ok || status == LuaStatus::Yield {
266 nres
267 } else {
268 1
269 };
270 let start = co_top - count;
271 let vals: Vec<LuaValue> = (start..co_top)
272 .map(|i| co_state.get_at(lua_vm::state::StackIdx(i as u32)))
273 .collect();
274 let new_co_top = if status == LuaStatus::Ok || status == LuaStatus::Yield {
275 (co_top - count).max(ci_func + 1)
276 } else {
277 co_top - count
278 };
279 co_state.set_top(lua_vm::state::StackIdx(new_co_top.max(0) as u32));
280 (status, vals)
281 };
282
283 pop_parent_gc_snapshot(state);
285
286 {
287 let mut g = state.global_mut();
288 let mut flush: Vec<(lua_vm::state::StackIdx, LuaValue)> = Vec::new();
289 for (tid, idx) in &parent_open_upval_slots {
290 if let Some(v) = g.cross_thread_upvals.remove(&(*tid, *idx)) {
291 flush.push((*idx, v));
292 }
293 }
294 drop(g);
295 for (idx, v) in flush {
296 state.set_at(idx, v);
297 }
298 }
299
300 match status {
301 LuaStatus::Ok | LuaStatus::Yield => {
302 if state.check_stack(results_or_err.len() as i32 + 1).is_err() {
303 push_lit_or_nil(state, b"too many results to resume");
304 return -1;
305 }
306 let n = results_or_err.len();
307 for v in results_or_err {
308 state.push(v);
309 }
310 n as i32
311 }
312 _ => {
313 for v in results_or_err {
314 state.push(v);
315 }
316 -1
317 }
318 }
319}
320
321fn push_parent_gc_snapshot(state: &mut LuaState) {
322 let top = (state.top_idx().0 as usize).min(state.stack.len());
323 let (mut stack_snapshot, mut upval_snapshot) = {
324 let mut g = state.global_mut();
325 (
326 g.snapshot_stack_pool.pop().unwrap_or_default(),
327 g.snapshot_upval_pool.pop().unwrap_or_default(),
328 )
329 };
330 stack_snapshot.extend(state.stack[..top].iter().map(|sv| sv.val));
331 upval_snapshot.extend(state.openupval.iter().cloned());
332 let mut g = state.global_mut();
333 g.suspended_parent_stacks.push(stack_snapshot);
334 g.suspended_parent_open_upvals.push(upval_snapshot);
335}
336
337fn pop_parent_gc_snapshot(state: &mut LuaState) {
338 let mut g = state.global_mut();
339 if let Some(mut v) = g.suspended_parent_open_upvals.pop() {
340 v.clear();
341 g.snapshot_upval_pool.push(v);
342 }
343 if let Some(mut v) = g.suspended_parent_stacks.pop() {
344 v.clear();
345 g.snapshot_stack_pool.push(v);
346 }
347}
348
349pub(crate) struct RootedThreadBorrow<'a> {
367 inner: std::cell::RefMut<'a, LuaState>,
368}
369
370impl RootedThreadBorrow<'_> {
371 pub(crate) fn resnapshot(&mut self) {
375 let top = (self.inner.top_idx().0 as usize).min(self.inner.stack.len());
376 let stack_copy: Vec<LuaValue> = self.inner.stack[..top].iter().map(|sv| sv.val).collect();
377 let upval_copy: Vec<GcRef<lua_types::UpVal>> = self.inner.openupval.to_vec();
378 let mut g = self.inner.global_mut();
379 if let Some(slot) = g.suspended_parent_stacks.last_mut() {
380 slot.clear();
381 slot.extend(stack_copy);
382 }
383 if let Some(slot) = g.suspended_parent_open_upvals.last_mut() {
384 slot.clear();
385 slot.extend(upval_copy);
386 }
387 }
388}
389
390impl std::ops::Deref for RootedThreadBorrow<'_> {
391 type Target = LuaState;
392 fn deref(&self) -> &LuaState {
393 &self.inner
394 }
395}
396
397impl std::ops::DerefMut for RootedThreadBorrow<'_> {
398 fn deref_mut(&mut self) -> &mut LuaState {
399 &mut self.inner
400 }
401}
402
403impl Drop for RootedThreadBorrow<'_> {
404 fn drop(&mut self) {
405 let mut g = self.inner.global_mut();
406 if let Some(mut v) = g.suspended_parent_open_upvals.pop() {
407 v.clear();
408 g.snapshot_upval_pool.push(v);
409 }
410 if let Some(mut v) = g.suspended_parent_stacks.pop() {
411 v.clear();
412 g.snapshot_stack_pool.push(v);
413 }
414 }
415}
416
417pub(crate) fn borrow_thread_rooted<'a>(
421 state: &mut LuaState,
422 cell: &'a std::cell::RefCell<LuaState>,
423) -> RootedThreadBorrow<'a> {
424 let inner = cell.borrow_mut();
425 let top = (inner.top_idx().0 as usize).min(inner.stack.len());
426 let (mut stack_snapshot, mut upval_snapshot) = {
427 let mut g = state.global_mut();
428 (
429 g.snapshot_stack_pool.pop().unwrap_or_default(),
430 g.snapshot_upval_pool.pop().unwrap_or_default(),
431 )
432 };
433 stack_snapshot.extend(inner.stack[..top].iter().map(|sv| sv.val));
434 upval_snapshot.extend(inner.openupval.iter().cloned());
435 let mut g = state.global_mut();
436 g.suspended_parent_stacks.push(stack_snapshot);
437 g.suspended_parent_open_upvals.push(upval_snapshot);
438 drop(g);
439 RootedThreadBorrow { inner }
440}
441
442fn push_lit_or_nil(state: &mut LuaState, bytes: &[u8]) {
444 match state.intern_str(bytes) {
445 Ok(s) => state.push(LuaValue::Str(s)),
446 Err(_) => state.push(LuaValue::Nil),
447 }
448}
449
450pub fn co_resume(state: &mut LuaState) -> Result<usize, LuaError> {
458 let co = get_co(state)?;
459 let narg = state.get_top() - 1;
462 let r = aux_resume(state, co, narg);
463 if r < 0 {
464 if state.sandbox_aborting() {
468 let top = state.get_top();
469 let err_val = state.value_at(top);
470 return Err(LuaError::from_value(err_val));
471 }
472 state.push(LuaValue::Bool(false));
473 state.insert(-2)?;
474 Ok(2)
475 } else {
476 state.push(LuaValue::Bool(true));
477 state.insert(-(r + 1))?;
478 Ok((r + 1) as usize)
479 }
480}
481
482fn aux_wrap(state: &mut LuaState) -> Result<usize, LuaError> {
490 let up = state.value_at(upvalue_index(1));
491 let co = match up {
492 LuaValue::Thread(t) => t,
493 _ => {
494 return Err(LuaError::runtime(format_args!(
495 "coroutine.wrap: upvalue is not a thread"
496 )))
497 }
498 };
499 let narg = state.get_top();
500 let r = aux_resume(state, co.clone(), narg);
501 if r < 0 {
502 let top = state.get_top();
503 let mut err_val = state.value_at(top);
504 if aux_status(state, &co) == COS_DEAD {
505 let old_err = state.pop();
506 let nclose = close_suspended_or_dead(state, co)?;
507 err_val = if nclose >= 2 {
508 let top = state.get_top();
509 state.value_at(top)
510 } else {
511 old_err
512 };
513 state.pop_n(nclose);
514 }
515 Err(LuaError::from_value(err_val))
516 } else {
517 Ok(r as usize)
518 }
519}
520
521pub fn co_create(state: &mut LuaState) -> Result<usize, LuaError> {
534 state.check_arg_type(1, LuaType::Function)?;
535 let body = state.value_at(1);
536 let _nl = state.new_thread(Some(body))?;
537 Ok(1)
538}
539
540pub fn co_wrap(state: &mut LuaState) -> Result<usize, LuaError> {
548 co_create(state)?;
549 state.push_cclosure(aux_wrap, 1)?;
550 Ok(1)
551}
552
553pub fn co_yield(state: &mut LuaState) -> Result<usize, LuaError> {
560 let n = state.get_top();
561 let r = lua_vm::do_::lua_yieldk(state, n, 0, None)?;
562 Ok(r as usize)
563}
564
565pub fn co_status(state: &mut LuaState) -> Result<usize, LuaError> {
570 let co = get_co(state)?;
571 let idx = aux_status(state, &co) as usize;
572 let name: &[u8] = STAT_NAMES[idx];
573 let interned = state.intern_str(name)?;
574 state.push(LuaValue::Str(interned));
575 Ok(1)
576}
577
578pub fn co_isyieldable(state: &mut LuaState) -> Result<usize, LuaError> {
582 let is_yieldable = if matches!(state.type_at(1), LuaType::None) {
583 state.is_yieldable()
584 } else {
585 let co = get_co(state)?;
586 let co_id = co.id;
587 let (is_main, is_current) = {
588 let g = state.global();
589 (co_id == g.main_thread_id, co_id == g.current_thread_id)
590 };
591 if is_main {
592 false
593 } else if is_current {
594 state.is_yieldable()
595 } else {
596 let entry_rc = {
597 let g = state.global();
598 g.threads
599 .get(&co_id)
600 .expect("thread value carries an id that must resolve in GlobalState::threads")
601 .state
602 .clone()
603 };
604 let target_is_yieldable = match entry_rc.try_borrow() {
605 Ok(b) => b.is_yieldable(),
606 Err(_) => false,
607 };
608 target_is_yieldable
609 }
610 };
611 state.push(LuaValue::Bool(is_yieldable));
612 Ok(1)
613}
614
615pub fn co_running(state: &mut LuaState) -> Result<usize, LuaError> {
620 let is_main = state.push_thread()?;
623 if matches!(state.global().lua_version, lua_types::LuaVersion::V51) {
629 if is_main {
630 state.pop_n(1);
631 state.push(LuaValue::Nil);
632 }
633 return Ok(1);
634 }
635 state.push(LuaValue::Bool(is_main));
636 Ok(2)
637}
638
639pub fn co_close(state: &mut LuaState) -> Result<usize, LuaError> {
647 lua_vm::state::inc_c_stack(state)?;
648 let result = (|| {
649 let co = get_opt_co(state)?;
650 let status = aux_status(state, &co);
651 match status {
652 COS_DEAD | COS_YIELD => close_suspended_or_dead(state, co),
653 _ => {
654 if matches!(state.global().lua_version, lua_types::LuaVersion::V55)
655 && status == COS_RUN
656 && state.global().closing_thread_id == Some(co.id)
657 {
658 state.push(LuaValue::Bool(true));
659 return Ok(1);
660 }
661 if matches!(state.global().lua_version, lua_types::LuaVersion::V55)
662 && status == COS_RUN
663 && co.id == state.global().main_thread_id
664 {
665 return Err(LuaError::runtime(format_args!("cannot close main thread")));
666 }
667 if matches!(state.global().lua_version, lua_types::LuaVersion::V55)
668 && status == COS_RUN
669 && co.id == state.global().current_thread_id
670 {
671 state.global_mut().closing_thread_id = Some(co.id);
672 let in_status = state.status as i32;
673 let s = lua_vm::state::reset_thread(state, in_status);
674 state.global_mut().closing_thread_id = None;
675 state.n_ccalls = state.n_ccalls.saturating_sub(1);
676 std::panic::panic_any(LuaThreadClose(LuaStatus::from_raw(s)));
677 }
678 let name = if status == COS_RUN {
679 "running"
680 } else {
681 "normal"
682 };
683 Err(LuaError::runtime(format_args!(
684 "cannot close a {} coroutine",
685 name
686 )))
687 }
688 }
689 })();
690 state.n_ccalls -= 1;
691 result
692}
693
694fn close_suspended_or_dead(
696 state: &mut LuaState,
697 co: GcRef<lua_types::value::LuaThread>,
698) -> Result<usize, LuaError> {
699 let co_id = co.id;
700 let entry_rc_opt = {
701 let g = state.global();
702 g.threads.get(&co_id).map(|e| e.state.clone())
703 };
704 let entry_rc = match entry_rc_opt {
705 Some(rc) => rc,
706 None => {
707 state.push(LuaValue::Bool(true));
708 return Ok(1);
709 }
710 };
711 let parent_thread_id = state.global().current_thread_id;
712 let caller_c_calls = state.c_calls();
713
714 let parent_open_upval_slots: Vec<(u64, lua_vm::state::StackIdx)> = state
715 .openupval
716 .iter()
717 .filter_map(|uv| match &*uv.slot() {
718 lua_types::UpValState::Open { thread_id, idx } => Some((*thread_id as u64, *idx)),
719 lua_types::UpValState::Closed(_) => None,
720 })
721 .collect();
722 {
723 let mut g = state.global_mut();
724 for (tid, idx) in &parent_open_upval_slots {
725 let val = state.get_at(*idx);
726 g.cross_thread_upvals.insert((*tid, *idx), val);
727 }
728 }
729
730 push_parent_gc_snapshot(state);
731
732 let (status, err_value): (i32, Option<LuaValue>) = {
733 let mut co_state = entry_rc.borrow_mut();
734 co_state.global_mut().current_thread_id = co_id;
735 co_state.global_mut().closing_thread_id = Some(co_id);
736 co_state.n_ccalls = caller_c_calls;
737 let in_status = co_state.status as i32;
738 let s = lua_vm::state::reset_thread(&mut *co_state, in_status);
739 co_state.global_mut().closing_thread_id = None;
740 co_state.global_mut().current_thread_id = parent_thread_id;
741 if s == LuaStatus::Ok as i32 {
742 (s, None)
743 } else {
744 let top = co_state.top_idx().0;
745 if top > 0 {
746 let err = co_state.get_at(lua_vm::state::StackIdx(top - 1));
747 co_state.set_top(lua_vm::state::StackIdx(top - 1));
748 (s, Some(err))
749 } else {
750 (s, Some(LuaValue::Nil))
751 }
752 }
753 };
754
755 pop_parent_gc_snapshot(state);
756
757 {
758 let mut g = state.global_mut();
759 let mut flush: Vec<(lua_vm::state::StackIdx, LuaValue)> = Vec::new();
760 for (tid, idx) in &parent_open_upval_slots {
761 if let Some(v) = g.cross_thread_upvals.remove(&(*tid, *idx)) {
762 flush.push((*idx, v));
763 }
764 }
765 drop(g);
766 for (idx, v) in flush {
767 state.set_at(idx, v);
768 }
769 }
770
771 if status == LuaStatus::Ok as i32 {
772 state.push(LuaValue::Bool(true));
773 Ok(1)
774 } else {
775 state.push(LuaValue::Bool(false));
776 if let Some(v) = err_value {
777 state.push(v);
778 } else {
779 state.push(LuaValue::Nil);
780 }
781 Ok(2)
782 }
783}
784
785pub fn open_coroutine(state: &mut LuaState) -> Result<usize, LuaError> {
791 use lua_types::LuaVersion;
795 let version = state.global().lua_version;
796 let has_close = !matches!(version, LuaVersion::V51 | LuaVersion::V52 | LuaVersion::V53);
797 let has_isyieldable = !matches!(version, LuaVersion::V51 | LuaVersion::V52);
801 if has_close && has_isyieldable {
802 state.new_lib(CO_FUNCS)?;
803 } else {
804 let filtered: Vec<(&[u8], lua_CFunction)> = CO_FUNCS
805 .iter()
806 .filter(|(name, _)| {
807 (has_close || *name != b"close".as_slice())
808 && (has_isyieldable || *name != b"isyieldable".as_slice())
809 })
810 .copied()
811 .collect();
812 state.new_lib(&filtered)?;
813 }
814 Ok(1)
815}
816
817