perf(vm): Only check for stack overflow when entering a function

Author: Marwes

vm/src/channel.rs

@@ -229,7 +229,7 @@ fn spawn_<'vm>(value: WithVM<'vm, Function<&'vm Thread, fn(())>>) -> VmResult<Ro
         };
         value_variant.clone().push(&mut context)?;
         context.push(ValueRepr::Int(0));
-        context.context().stack.enter_scope(1, &*callable);
+        context.context().stack.enter_scope(1, &*callable)?;
     }
     Ok(thread)
 }

vm/src/stack.rs

@@ -10,8 +10,9 @@ use crate::{
     gc::{self, CloneUnrooted, CopyUnrooted, GcPtr, Trace},
     types::VmIndex,
     value::{ClosureData, DataStruct, ExternFunction, Value, ValueRepr},
-    Variants,
+    Error, Result, Variants,
 };
+use gc::Borrow;
 
 pub trait StackPrimitive {
     fn push_to(&self, stack: &mut Stack);
@@ -199,6 +200,7 @@ pub trait StackState: CopyUnrooted + Sized {
     fn from_state(state: &State) -> &Self;
     fn from_state_mut(state: &mut State) -> &mut Self;
     fn to_state(&self) -> gc::Borrow<State>;
+    fn max_stack_size(&self) -> VmIndex;
 }
 
 impl StackState for State {
@@ -211,6 +213,14 @@ impl StackState for State {
     fn to_state(&self) -> gc::Borrow<State> {
         gc::Borrow::new(self)
     }
+
+    fn max_stack_size(&self) -> VmIndex {
+        match self {
+            State::Unknown => 0,
+            State::Closure(closure) => closure.max_stack_size(),
+            State::Extern(ext) => ext.max_stack_size(),
+        }
+    }
 }
 
 impl StackState for ClosureState {
@@ -229,6 +239,9 @@ impl StackState for ClosureState {
     fn to_state(&self) -> gc::Borrow<State> {
         construct_gc!(State::Closure(@ gc::Borrow::new(self)))
     }
+    fn max_stack_size(&self) -> VmIndex {
+        self.closure.function.max_stack_size
+    }
 }
 
 impl StackState for ExternState {
@@ -247,6 +260,9 @@ impl StackState for ExternState {
     fn to_state(&self) -> gc::Borrow<State> {
         construct_gc!(State::Extern(@ gc::Borrow::new(self)))
     }
+    fn max_stack_size(&self) -> VmIndex {
+        0
+    }
 }
 
 #[derive(Debug, PartialEq)]
@@ -408,6 +424,7 @@ pub struct Stack {
     values: Vec<Value>,
     #[cfg_attr(feature = "serde_derive", serde(state))]
     frames: Vec<Frame<State>>,
+    max_stack_size: VmIndex,
 }
 
 unsafe impl Trace for Stack {
@@ -419,9 +436,14 @@ impl Stack {
         Stack {
             values: Vec::new(),
             frames: Vec::new(),
+            max_stack_size: VmIndex::MAX,
         }
     }
 
+    pub fn set_max_stack_size(&mut self, max_stack_size: VmIndex) {
+        self.max_stack_size = max_stack_size;
+    }
+
     fn assert_pop(&self, count: VmIndex) {
         let frame = self.frames.last().unwrap();
         let args = if let State::Extern(ExternState {
@@ -656,9 +678,9 @@ impl<'a: 'b, 'b> StackFrame<'b, State> {
         }
     }
 
-    pub fn new_frame(stack: &'b mut Stack, args: VmIndex, state: State) -> StackFrame<'b> {
-        let frame = unsafe { Self::add_new_frame(stack, args, &state, false).unrooted() };
-        StackFrame { stack, frame }
+    pub fn new_frame(stack: &'b mut Stack, args: VmIndex, state: State) -> Result<StackFrame<'b>> {
+        let frame = unsafe { Self::add_new_frame(stack, args, &state, false)?.unrooted() };
+        Ok(StackFrame { stack, frame })
     }
 }
 
@@ -825,7 +847,7 @@ where
         }
     }
 
-    pub(crate) fn enter_scope<T>(self, args: VmIndex, state: &T) -> StackFrame<'b, T>
+    pub(crate) fn enter_scope<T>(self, args: VmIndex, state: &T) -> Result<StackFrame<'b, T>>
     where
         T: StackState,
     {
@@ -837,16 +859,16 @@ where
         args: VmIndex,
         state: &T,
         excess: bool,
-    ) -> StackFrame<'b, T>
+    ) -> Result<StackFrame<'b, T>>
     where
         T: StackState,
     {
         let stack = self.stack;
-        let frame = unsafe { Self::add_new_frame(stack, args, state, excess).unrooted() };
-        StackFrame { stack, frame }
+        let frame = unsafe { Self::add_new_frame(stack, args, state, excess)?.unrooted() };
+        Ok(StackFrame { stack, frame })
     }
 
-    pub(crate) fn exit_scope(self) -> Result<StackFrame<'b, State>, &'b mut Stack>
+    pub(crate) fn exit_scope(self) -> std::result::Result<StackFrame<'b, State>, &'b mut Stack>
     where
         S: StackState,
     {
@@ -886,7 +908,7 @@ where
         args: VmIndex,
         state: &'gc T,
         excess: bool,
-    ) -> gc::Borrow<'gc, Frame<T>>
+    ) -> Result<Borrow<'gc, Frame<T>>>
     where
         T: StackState,
     {
@@ -911,13 +933,18 @@ where
                 );
             }
         }
+        // Before entering a function check that the stack cannot exceed `max_stack_size`
+        if stack.len() + frame.state.max_stack_size() > stack.max_stack_size {
+            return Err(Error::StackOverflow(stack.max_stack_size));
+        }
+
         // SAFETY The frame's gc pointers are scanned the `Stack::trace` since they are on
         // the stack
         unsafe {
             stack.frames.push(frame.to_state().clone_unrooted());
         }
         debug!("----> Store {} {:?}", stack.frames.len(), frame.to_state());
-        frame
+        Ok(frame)
     }
 }
 
@@ -1158,15 +1185,15 @@ mod tests {
         let _ = ::env_logger::try_init();
 
         let mut stack = Stack::new();
-        let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown);
+        let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
         frame.push(Int(0));
         frame.push(Int(1));
 
-        frame = frame.enter_scope(2, &State::Unknown);
+        frame = frame.enter_scope(2, &State::Unknown).unwrap();
         frame.push(Int(2));
         frame.push(Int(3));
 
-        frame = frame.enter_scope(1, &State::Unknown);
+        frame = frame.enter_scope(1, &State::Unknown).unwrap();
         frame.push(Int(4));
         frame.push(Int(5));
         frame.push(Int(6));
@@ -1193,14 +1220,14 @@ mod tests {
 
         let mut stack = Stack::new();
         {
-            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown);
+            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
             frame.push(Int(0));
             frame.push(Int(1));
-            let frame = frame.enter_scope(2, &*ExternState::new(&ext));
+            let frame = frame.enter_scope(2, &*ExternState::new(&ext)).unwrap();
             let _lock = frame.into_lock();
         }
         // Panic as it attempts to access past the lock
-        StackFrame::new_frame(&mut stack, 1, State::Unknown);
+        StackFrame::new_frame(&mut stack, 1, State::Unknown).unwrap();
 
         unsafe { gc.clear() }
     }
@@ -1215,9 +1242,9 @@ mod tests {
 
         let mut stack = Stack::new();
         {
-            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown);
+            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
             frame.push(Int(0));
-            let frame = frame.enter_scope(1, &*ExternState::new(&ext));
+            let frame = frame.enter_scope(1, &*ExternState::new(&ext)).unwrap();
             let _lock = frame.into_lock();
         }
         // Panic as it attempts to pop a locked value
@@ -1235,14 +1262,14 @@ mod tests {
 
         let mut stack = Stack::new();
         let lock = {
-            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown);
+            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
             frame.push(Int(0));
             frame.push(Int(1));
-            let frame = frame.enter_scope(2, &*ExternState::new(&ext));
+            let frame = frame.enter_scope(2, &*ExternState::new(&ext)).unwrap();
             frame.into_lock()
         };
         {
-            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown);
+            let mut frame = StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
             frame.push(Int(2));
             frame = frame.exit_scope().unwrap();
             frame.exit_scope().unwrap_err();
@@ -1259,12 +1286,12 @@ mod tests {
         let _ = ::env_logger::try_init();
 
         let mut stack = Stack::new();
-        StackFrame::new_frame(&mut stack, 0, State::Unknown);
+        StackFrame::new_frame(&mut stack, 0, State::Unknown).unwrap();
         let mut stack = StackFrame::<State>::current(&mut stack);
         stack.push(Int(0));
         stack.insert_slice(0, &[Int(2).into(), Int(1).into()]);
         assert_eq!(&stack[..], [Int(2).into(), Int(1).into(), Int(0).into()]);
-        stack = stack.enter_scope(2, &State::Unknown);
+        stack = stack.enter_scope(2, &State::Unknown).unwrap();
         stack.insert_slice(1, &[Int(10).into()]);
         assert_eq!(&stack[..], [Int(1).into(), Int(10).into(), Int(0).into()]);
         stack.insert_slice(1, &[]);