add test

Author: loongs-zhang

core/src/coroutine/korosensei.rs

@@ -324,7 +324,28 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
                 r
             });
         }
-        DefaultStack::new(stack_size).map(|stack| corosensei::on_stack(stack, callback))
+        // in thread
+        thread_local! {
+            #[allow(clippy::missing_const_for_thread_local)]
+            static STACK_INFOS: RefCell<VecDeque<StackInfo>> = const { RefCell::new(VecDeque::new()) };
+        }
+        if let Some(last_stack_info) = STACK_INFOS.with(|s| s.borrow().back().copied()) {
+            let remaining_stack = psm::stack_pointer() as usize - last_stack_info.stack_bottom;
+            if remaining_stack >= red_zone {
+                return Ok(callback());
+            }
+        }
+        DefaultStack::new(stack_size).map(|stack| {
+            STACK_INFOS.with(|s| {
+                s.borrow_mut().push_back(StackInfo {
+                    stack_top: stack.base().get(),
+                    stack_bottom: stack.limit().get(),
+                });
+            });
+            let r = corosensei::on_stack(stack, callback);
+            _ = STACK_INFOS.with(|s| s.borrow_mut().pop_back());
+            r
+        })
     }
 }
 

core/tests/coroutine.rs

@@ -70,6 +70,47 @@ fn coroutine_delay() -> std::io::Result<()> {
     Ok(())
 }
 
+#[test]
+fn sp_in_bounds() -> std::io::Result<()> {
+    let mut coroutine = co!(|suspender, input| {
+        let current = Coroutine::<(), (), ()>::current().unwrap();
+        let stack_info = current.stack_infos().back().cloned().unwrap();
+        assert!(current.stack_ptr_in_bounds(psm::stack_pointer() as u64));
+        assert_eq!(
+            current.stack_ptr_in_bounds(stack_info.stack_top as u64 + 1),
+            false
+        );
+        assert_eq!(
+            current.stack_ptr_in_bounds(stack_info.stack_bottom as u64 - 1),
+            false
+        );
+        assert_eq!(1, input);
+        assert_eq!(3, suspender.suspend_with(2));
+        4
+    })?;
+    assert_eq!(CoroutineState::Suspend(2, 0), coroutine.resume_with(1)?);
+    assert_eq!(CoroutineState::Complete(4), coroutine.resume_with(3)?);
+    println!("{:?}", coroutine);
+    println!("{}", coroutine);
+    Ok(())
+}
+
+#[test]
+fn thread_stack_growth() {
+    fn recurse(i: u32, p: &mut [u8; 10240]) {
+        Coroutine::<(), (), ()>::maybe_grow(|| {
+            // Ensure the stack allocation isn't optimized away.
+            unsafe { std::ptr::read_volatile(&p) };
+            if i > 0 {
+                recurse(i - 1, &mut [0; 10240]);
+            }
+        })
+        .expect("allocate stack failed")
+    }
+    // Use 10MB of stack.
+    recurse(1000, &mut [0; 10240]);
+}
+
 #[test]
 fn coroutine_stack_growth() -> std::io::Result<()> {
     let mut coroutine = co!(|_: &Suspender<(), ()>, ()| {