RustPython · windelbouwman · Aug 28, 2019 · Sep 9, 2019 · Sep 28, 2019 · Sep 28, 2019
diff --git a/bytecode/src/bytecode.rs b/bytecode/src/bytecode.rs
@@ -45,6 +45,9 @@ pub struct CodeObject {
     pub first_line_number: usize,
     pub obj_name: String, // Name of the object that created this code object
     pub is_generator: bool,
+
+    /// Maximum value stack size required for this code object.
+    pub max_stack_size: usize,
 }
 
 bitflags! {
@@ -367,6 +370,7 @@ impl CodeObject {
             first_line_number,
             obj_name,
             is_generator: false,
+            max_stack_size: 0,
         }
     }
 

diff --git a/compiler/src/lib.rs b/compiler/src/lib.rs
@@ -10,4 +10,5 @@ pub mod compile;
 pub mod error;
 pub(crate) mod output_stream;
 pub mod peephole;
+mod stack_effect;
 pub mod symboltable;
diff --git a/compiler/src/output_stream.rs b/compiler/src/output_stream.rs
@@ -1,5 +1,7 @@
 use rustpython_bytecode::bytecode::{CodeObject, Instruction, Label, Location};
 
+use crate::stack_effect::stack_effect;
+
 pub trait OutputStream: From<CodeObject> + Into<CodeObject> {
     /// Output an instruction
     fn emit(&mut self, instruction: Instruction, location: Location);
@@ -11,11 +13,17 @@ pub trait OutputStream: From<CodeObject> + Into<CodeObject> {
 
 pub struct CodeObjectStream {
     code: CodeObject,
+    max_stack: isize,
+    current_stack_size: isize,
 }
 
 impl From<CodeObject> for CodeObjectStream {
     fn from(code: CodeObject) -> Self {
-        CodeObjectStream { code }
+        CodeObjectStream {
+            code,
+            max_stack: 0,
+            current_stack_size: 0,
+        }
     }
 }
 impl From<CodeObjectStream> for CodeObject {
@@ -26,13 +34,21 @@ impl From<CodeObjectStream> for CodeObject {
 
 impl OutputStream for CodeObjectStream {
     fn emit(&mut self, instruction: Instruction, location: Location) {
+        let effect = stack_effect(&instruction);
+        self.current_stack_size += effect;
+        if self.current_stack_size > self.max_stack {
+            self.max_stack = self.current_stack_size;
+        }
+
         self.code.instructions.push(instruction);
         self.code.locations.push(location);
     }
+
     fn set_label(&mut self, label: Label) {
         let position = self.code.instructions.len();
         self.code.label_map.insert(label, position);
     }
+
     fn mark_generator(&mut self) {
         self.code.is_generator = true;
     }

diff --git a/compiler/src/stack_effect.rs b/compiler/src/stack_effect.rs
@@ -0,0 +1,88 @@
+use rustpython_bytecode::bytecode::{CallType, Instruction};
+
+/// Determine the effect on the stack of the given instruction.
+///
+/// This function must match what is executed in frame.rs
+/// The return value is the amount of stack change created by the
+/// given opcode.
+pub fn stack_effect(instruction: &Instruction) -> isize {
+    use Instruction::*;
+
+    match instruction {
+        ImportStar => -1,
+        Import { .. } => 1,
+        ImportFrom { .. } => 1,
+        PopException => 0,
+        Jump { .. } => 0,
+        JumpIfFalse { .. } => -1,
+        JumpIfTrue { .. } => -1,
+        JumpIfFalseOrPop { .. } => -1,
+        JumpIfTrueOrPop { .. } => -1,
+        Break => 0,
+        Continue => 0,
+        PopBlock => 0,
+        GetIter => 0,
+        ForIter { .. } => 1,
+        CallFunction { typ } => match typ {
+            CallType::Positional(amount) => -(*amount as isize),
+            CallType::Keyword(amount) => -1 - (*amount as isize),
+            CallType::Ex(has_kwargs) => {
+                if *has_kwargs {
+                    -2
+                } else {
+                    -1
+                }
+            }
+        },
+        UnaryOperation { .. } => 0,
+        BinaryOperation { .. } => -1,
+        Pop => -1,
+        Duplicate => 1,
+        Rotate { .. } => 0,
+        Reverse { .. } => 0,
+        Subscript { .. } => -1,
+        StoreSubscript { .. } => -3,
+        DeleteSubscript { .. } => -2,
+        LoadAttr { .. } => 0,
+        StoreAttr { .. } => -2,
+        DeleteAttr { .. } => -1,
+        LoadName { .. } => 1,
+        StoreName { .. } => -1,
+        DeleteName { .. } => 0,
+        LoadConst { .. } => 1,
+        ReturnValue { .. } => -1,
+        YieldValue { .. } => 0,
+        YieldFrom { .. } => 0,
+        CompareOperation { .. } => -1,
+        BuildList { size, .. }
+        | BuildSet { size, .. }
+        | BuildTuple { size, .. }
+        | BuildSlice { size, .. }
+        | BuildString { size } => 1 - (*size as isize),
+        BuildMap { size, unpack } => {
+            let size = *size as isize;
+            if *unpack {
+                1 - size
+            } else {
+                1 - size * 2
+            }
+        }
+        ListAppend { .. } => -1,
+        SetAdd { .. } => -1,
+        MapAdd { .. } => -2,
+        UnpackEx { before, after } => -1 + (*before as isize) + (*after as isize) + 1,
+        UnpackSequence { size } => -1 + (*size as isize),
+        SetupLoop { .. } => 0,
+        SetupWith { .. } => 0,
+        CleanupWith { .. } => 0,
+        SetupExcept { .. } => 0,
+        SetupFinally { .. } => 0,
+        EnterFinally { .. } => 0,
+        EndFinally { .. } => 0,
+        LoadBuildClass => 1,
+        MakeFunction { .. } => 0,
+        Raise { argc } => -(*argc as isize),
+        PrintExpr => -1,
+        FormatValue { .. } => 0,
+    }
+}
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -1,4 +1,4 @@
-use std::cell::RefCell;
+use std::cell::{Cell, RefCell};
 use std::fmt;
 
 use crate::builtins;
@@ -85,7 +85,7 @@ pub struct Frame {
     stack: RefCell<Vec<PyObjectRef>>, // The main data frame of the stack machine
     blocks: RefCell<Vec<Block>>,      // Block frames, for controlling loops and exceptions
     pub scope: Scope,                 // Variables
-    pub lasti: RefCell<usize>,        // index of last instruction ran
+    lasti: Cell<usize>,               // index of last instruction ran
 }
 
 impl PyValue for Frame {
@@ -105,25 +105,15 @@ pub type FrameResult = PyResult<Option<ExecutionResult>>;
 
 impl Frame {
     pub fn new(code: PyCodeRef, scope: Scope) -> Frame {
-        //populate the globals and locals
-        //TODO: This is wrong, check https://github.com/nedbat/byterun/blob/31e6c4a8212c35b5157919abff43a7daa0f377c6/byterun/pyvm2.py#L95
-        /*
-        let globals = match globals {
-            Some(g) => g,
-            None => HashMap::new(),
-        };
-        */
-        // let locals = globals;
-        // locals.extend(callargs);
+        let code = code.code.clone();
+        let max_stack_size = code.max_stack_size;
 
         Frame {
-            code: code.code.clone(),
-            stack: RefCell::new(vec![]),
+            code,
+            stack: RefCell::new(Vec::with_capacity(max_stack_size)),
             blocks: RefCell::new(vec![]),
-            // save the callargs as locals
-            // globals: locals.clone(),
             scope,
-            lasti: RefCell::new(0),
+            lasti: Cell::new(0),
         }
     }
 
@@ -189,8 +179,8 @@ impl Frame {
     }
 
     pub fn fetch_instruction(&self) -> &bytecode::Instruction {
-        let ins2 = &self.code.instructions[*self.lasti.borrow()];
-        *self.lasti.borrow_mut() += 1;
+        let ins2 = &self.code.instructions[self.get_lasti()];
+        self.lasti.set(self.lasti.get() + 1);
         ins2
     }
 
@@ -896,31 +886,40 @@ impl Frame {
     }
 
     fn execute_raise(&self, vm: &VirtualMachine, argc: usize) -> FrameResult {
-        let cause = match argc {
-            2 => self.get_exception(vm, true)?,
-            _ => vm.get_none(),
+        if argc > 2 {
+            // TODO: argc of 3 not implemented.
+            panic!("Invalid parameter for RAISE_VARARGS, must be between 0 to 3");
+        }
+
+        let cause = if argc < 2 {
+            vm.get_none()
+        } else {
+            self.get_exception(vm, true)?
         };
-        let exception = match argc {
-            0 => match vm.current_exception() {
+
+        let exception = if argc == 0 {
+            match vm.current_exception() {
                 Some(exc) => exc,
                 None => {
                     return Err(vm.new_exception(
                         vm.ctx.exceptions.runtime_error.clone(),
                         "No active exception to reraise".to_string(),
                     ));
                 }
-            },
-            1 | 2 => self.get_exception(vm, false)?,
-            3 => panic!("Not implemented!"),
-            _ => panic!("Invalid parameter for RAISE_VARARGS, must be between 0 to 3"),
+            }
+        } else {
+            self.get_exception(vm, false)?
         };
-        let context = match argc {
-            0 => vm.get_none(), // We have already got the exception,
-            _ => match vm.current_exception() {
+
+        let context = if argc == 0 {
+            vm.get_none() // We have already got the exception,
+        } else {
+            match vm.current_exception() {
                 Some(exc) => exc,
                 None => vm.get_none(),
-            },
+            }
         };
+
         info!(
             "Exception raised: {:?} with cause: {:?} and context: {:?}",
             exception, cause, context
@@ -940,7 +939,7 @@ impl Frame {
         match next_obj {
             Some(value) => {
                 // Set back program counter:
-                *self.lasti.borrow_mut() -= 1;
+                self.lasti.set(self.get_lasti() - 1);
                 Ok(Some(ExecutionResult::Yield(value)))
             }
             None => Ok(None),
@@ -1011,6 +1010,11 @@ impl Frame {
             }
         }
     }
+
+    pub fn get_lasti(&self) -> usize {
+        self.lasti.get()
+    }
+
     fn execute_make_function(
         &self,
         vm: &VirtualMachine,
@@ -1080,39 +1084,40 @@ impl Frame {
         op: &bytecode::BinaryOperator,
         inplace: bool,
     ) -> FrameResult {
+        use bytecode::BinaryOperator::*;
         let b_ref = self.pop_value();
         let a_ref = self.pop_value();
         let value = if inplace {
             match *op {
-                bytecode::BinaryOperator::Subtract => vm._isub(a_ref, b_ref),
-                bytecode::BinaryOperator::Add => vm._iadd(a_ref, b_ref),
-                bytecode::BinaryOperator::Multiply => vm._imul(a_ref, b_ref),
-                bytecode::BinaryOperator::MatrixMultiply => vm._imatmul(a_ref, b_ref),
-                bytecode::BinaryOperator::Power => vm._ipow(a_ref, b_ref),
-                bytecode::BinaryOperator::Divide => vm._itruediv(a_ref, b_ref),
-                bytecode::BinaryOperator::FloorDivide => vm._ifloordiv(a_ref, b_ref),
-                bytecode::BinaryOperator::Modulo => vm._imod(a_ref, b_ref),
-                bytecode::BinaryOperator::Lshift => vm._ilshift(a_ref, b_ref),
-                bytecode::BinaryOperator::Rshift => vm._irshift(a_ref, b_ref),
-                bytecode::BinaryOperator::Xor => vm._ixor(a_ref, b_ref),
-                bytecode::BinaryOperator::Or => vm._ior(a_ref, b_ref),
-                bytecode::BinaryOperator::And => vm._iand(a_ref, b_ref),
+                Subtract => vm._isub(a_ref, b_ref),
+                Add => vm._iadd(a_ref, b_ref),
+                Multiply => vm._imul(a_ref, b_ref),
+                MatrixMultiply => vm._imatmul(a_ref, b_ref),
+                Power => vm._ipow(a_ref, b_ref),
+                Divide => vm._itruediv(a_ref, b_ref),
+                FloorDivide => vm._ifloordiv(a_ref, b_ref),
+                Modulo => vm._imod(a_ref, b_ref),
+                Lshift => vm._ilshift(a_ref, b_ref),
+                Rshift => vm._irshift(a_ref, b_ref),
+                Xor => vm._ixor(a_ref, b_ref),
+                Or => vm._ior(a_ref, b_ref),
+                And => vm._iand(a_ref, b_ref),
             }?
         } else {
             match *op {
-                bytecode::BinaryOperator::Subtract => vm._sub(a_ref, b_ref),
-                bytecode::BinaryOperator::Add => vm._add(a_ref, b_ref),
-                bytecode::BinaryOperator::Multiply => vm._mul(a_ref, b_ref),
-                bytecode::BinaryOperator::MatrixMultiply => vm._matmul(a_ref, b_ref),
-                bytecode::BinaryOperator::Power => vm._pow(a_ref, b_ref),
-                bytecode::BinaryOperator::Divide => vm._truediv(a_ref, b_ref),
-                bytecode::BinaryOperator::FloorDivide => vm._floordiv(a_ref, b_ref),
-                bytecode::BinaryOperator::Modulo => vm._mod(a_ref, b_ref),
-                bytecode::BinaryOperator::Lshift => vm._lshift(a_ref, b_ref),
-                bytecode::BinaryOperator::Rshift => vm._rshift(a_ref, b_ref),
-                bytecode::BinaryOperator::Xor => vm._xor(a_ref, b_ref),
-                bytecode::BinaryOperator::Or => vm._or(a_ref, b_ref),
-                bytecode::BinaryOperator::And => vm._and(a_ref, b_ref),
+                Subtract => vm._sub(a_ref, b_ref),
+                Add => vm._add(a_ref, b_ref),
+                Multiply => vm._mul(a_ref, b_ref),
+                MatrixMultiply => vm._matmul(a_ref, b_ref),
+                Power => vm._pow(a_ref, b_ref),
+                Divide => vm._truediv(a_ref, b_ref),
+                FloorDivide => vm._floordiv(a_ref, b_ref),
+                Modulo => vm._mod(a_ref, b_ref),
+                Lshift => vm._lshift(a_ref, b_ref),
+                Rshift => vm._rshift(a_ref, b_ref),
+                Xor => vm._xor(a_ref, b_ref),
+                Or => vm._or(a_ref, b_ref),
+                And => vm._and(a_ref, b_ref),
             }?
         };
 
@@ -1208,7 +1213,7 @@ impl Frame {
     }
 
     pub fn get_lineno(&self) -> bytecode::Location {
-        self.code.locations[*self.lasti.borrow()].clone()
+        self.code.locations[self.get_lasti()].clone()
     }
 
     fn push_block(&self, typ: BlockType) {

diff --git a/vm/src/obj/objframe.rs b/vm/src/obj/objframe.rs
@@ -48,6 +48,6 @@ impl FrameRef {
     }
 
     fn f_lasti(self, vm: &VirtualMachine) -> PyObjectRef {
-        vm.ctx.new_int(*self.lasti.borrow())
+        vm.ctx.new_int(self.get_lasti())
     }
 }