Keep a stack of frames in the VM.

cthulahoops · cthulahoops · commit 4dd6592aad4e · 2019-02-27T10:47:33.000Z
diff --git a/tests/snippets/getframe.py b/tests/snippets/getframe.py
@@ -10,6 +10,8 @@ def test_function():
     x = 17
     assert sys._getframe().f_locals is not locals_dict
     assert sys._getframe().f_locals['x'] == 17
+    assert sys._getframe(1).f_locals['foo'] == 'bar'
+    print(sys._getframe(1).f_locals)
 
 test_function()
 
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -1,8 +1,8 @@
 extern crate rustpython_parser;
 
 use self::rustpython_parser::ast;
+use std::cell::RefCell;
 use std::fmt;
-use std::mem;
 use std::path::PathBuf;
 
 use crate::builtins;
@@ -49,10 +49,10 @@ enum BlockType {
 pub struct Frame {
     pub code: bytecode::CodeObject,
     // We need 1 stack per frame
-    stack: Vec<PyObjectRef>, // The main data frame of the stack machine
-    blocks: Vec<Block>,      // Block frames, for controlling loops and exceptions
-    pub locals: PyObjectRef, // Variables
-    pub lasti: usize,        // index of last instruction ran
+    stack: RefCell<Vec<PyObjectRef>>, // The main data frame of the stack machine
+    blocks: Vec<Block>,               // Block frames, for controlling loops and exceptions
+    pub locals: PyObjectRef,          // Variables
+    pub lasti: usize,                 // index of last instruction ran
 }
 
 // Running a frame can result in one of the below:
@@ -79,7 +79,7 @@ impl Frame {
 
         Frame {
             code: objcode::get_value(&code),
-            stack: vec![],
+            stack: RefCell::new(vec![]),
             blocks: vec![],
             // save the callargs as locals
             // globals: locals.clone(),
@@ -88,18 +88,9 @@ impl Frame {
         }
     }
 
-    pub fn run_frame_full(&mut self, vm: &mut VirtualMachine) -> PyResult {
-        match self.run_frame(vm)? {
-            ExecutionResult::Return(value) => Ok(value),
-            _ => panic!("Got unexpected result from function"),
-        }
-    }
-
-    pub fn run_frame(&mut self, vm: &mut VirtualMachine) -> Result<ExecutionResult, PyObjectRef> {
+    pub fn run(&mut self, vm: &mut VirtualMachine) -> Result<ExecutionResult, PyObjectRef> {
         let filename = &self.code.source_path.to_string();
 
-        let prev_frame = mem::replace(&mut vm.current_frame, Some(vm.ctx.new_frame(self.clone())));
-
         // This is the name of the object being run:
         let run_obj_name = &self.code.obj_name.to_string();
 
@@ -148,7 +139,6 @@ impl Frame {
             }
         };
 
-        vm.current_frame = prev_frame;
         value
     }
 
@@ -1082,50 +1072,51 @@ impl Frame {
     fn push_block(&mut self, typ: BlockType) {
         self.blocks.push(Block {
             typ,
-            level: self.stack.len(),
+            level: self.stack.borrow().len(),
         });
     }
 
     fn pop_block(&mut self) -> Option<Block> {
         let block = self.blocks.pop()?;
-        self.stack.truncate(block.level);
+        self.stack.borrow_mut().truncate(block.level);
         Some(block)
     }
 
     fn current_block(&self) -> Option<&Block> {
         self.blocks.last()
     }
 
-    pub fn push_value(&mut self, obj: PyObjectRef) {
-        self.stack.push(obj);
+    pub fn push_value(&self, obj: PyObjectRef) {
+        self.stack.borrow_mut().push(obj);
     }
 
-    fn pop_value(&mut self) -> PyObjectRef {
-        self.stack.pop().unwrap()
+    fn pop_value(&self) -> PyObjectRef {
+        self.stack.borrow_mut().pop().unwrap()
     }
 
     fn pop_multiple(&mut self, count: usize) -> Vec<PyObjectRef> {
         let mut objs: Vec<PyObjectRef> = Vec::new();
         for _x in 0..count {
-            objs.push(self.stack.pop().unwrap());
+            objs.push(self.pop_value());
         }
         objs.reverse();
         objs
     }
 
     fn last_value(&self) -> PyObjectRef {
-        self.stack.last().unwrap().clone()
+        self.stack.borrow().last().unwrap().clone()
     }
 
     fn nth_value(&self, depth: usize) -> PyObjectRef {
-        self.stack[self.stack.len() - depth - 1].clone()
+        self.stack.borrow_mut()[self.stack.borrow().len() - depth - 1].clone()
     }
 }
 
 impl fmt::Debug for Frame {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let stack_str = self
             .stack
+            .borrow()
             .iter()
             .map(|elem| format!("\n  > {:?}", elem))
             .collect::<Vec<_>>()
diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs
@@ -2,9 +2,7 @@
  * The mythical generator.
  */
 
-use std::cell::RefCell;
-
-use crate::frame::{ExecutionResult, Frame};
+use crate::frame::ExecutionResult;
 use crate::pyobject::{
     PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
 };
@@ -29,11 +27,9 @@ pub fn init(context: &PyContext) {
     );
 }
 
-pub fn new_generator(vm: &mut VirtualMachine, frame: Frame) -> PyResult {
+pub fn new_generator(vm: &mut VirtualMachine, frame: PyObjectRef) -> PyResult {
     Ok(PyObject::new(
-        PyObjectPayload::Generator {
-            frame: RefCell::new(frame),
-        },
+        PyObjectPayload::Generator { frame: frame },
         vm.ctx.generator_type.clone(),
     ))
 }
@@ -60,9 +56,13 @@ fn generator_send(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 
 fn send(vm: &mut VirtualMachine, gen: &PyObjectRef, value: &PyObjectRef) -> PyResult {
     if let PyObjectPayload::Generator { ref frame } = gen.payload {
-        let mut frame_mut = frame.borrow_mut();
-        frame_mut.push_value(value.clone());
-        match frame_mut.run_frame(vm)? {
+        if let PyObjectPayload::Frame { ref frame } = frame.payload {
+            frame.push_value(value.clone());
+        } else {
+            panic!("Generator frame isn't a frame.");
+        }
+
+        match vm.run_frame(frame.clone())? {
             ExecutionResult::Yield(value) => Ok(value),
             ExecutionResult::Return(_value) => {
                 // Stop iteration!
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -591,8 +591,13 @@ impl PyContext {
         )
     }
 
-    pub fn new_frame(&self, frame: Frame) -> PyObjectRef {
-        PyObject::new(PyObjectPayload::Frame { frame }, self.frame_type())
+    pub fn new_frame(&self, code: PyObjectRef, scope: PyObjectRef) -> PyObjectRef {
+        PyObject::new(
+            PyObjectPayload::Frame {
+                frame: Frame::new(code, scope),
+            },
+            self.frame_type(),
+        )
     }
 
     pub fn new_property<F: 'static + Fn(&mut VirtualMachine, PyFuncArgs) -> PyResult>(
@@ -1246,7 +1251,7 @@ pub enum PyObjectPayload {
         defaults: PyObjectRef,
     },
     Generator {
-        frame: RefCell<Frame>,
+        frame: PyObjectRef,
     },
     BoundMethod {
         function: PyObjectRef,
diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -1,5 +1,7 @@
+use crate::obj::objint;
 use crate::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
 use crate::vm::VirtualMachine;
+use num_traits::ToPrimitive;
 use std::rc::Rc;
 use std::{env, mem};
 
@@ -13,12 +15,30 @@ fn argv(ctx: &PyContext) -> PyObjectRef {
     ctx.new_list(argv)
 }
 
-fn getframe(vm: &mut VirtualMachine, _args: PyFuncArgs) -> PyResult {
-    if let Some(frame) = &vm.current_frame {
-        Ok(frame.clone())
-    } else {
-        panic!("Current frame is undefined!")
+fn frame_idx(vm: &mut VirtualMachine, offset: Option<&PyObjectRef>) -> Result<usize, PyObjectRef> {
+    if let Some(int) = offset {
+        if let Some(offset) = objint::get_value(&int).to_usize() {
+            if offset > vm.frames.len() - 1 {
+                return Err(vm.new_value_error("call stack is not deep enough".to_string()));
+            }
+            return Ok(offset);
+        }
     }
+    return Ok(0);
+}
+
+fn getframe(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [],
+        optional = [(offset, Some(vm.ctx.int_type()))]
+    );
+
+    let idx = frame_idx(vm, offset)?;
+    let idx = vm.frames.len() - idx - 1;
+    let frame = &vm.frames[idx];
+    Ok(frame.clone())
 }
 
 fn sys_getrefcount(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -12,9 +12,10 @@ use std::sync::{Mutex, MutexGuard};
 
 use crate::builtins;
 use crate::bytecode;
-use crate::frame::Frame;
+use crate::frame::ExecutionResult;
 use crate::obj::objbool;
 use crate::obj::objcode;
+use crate::obj::objframe;
 use crate::obj::objgenerator;
 use crate::obj::objiter;
 use crate::obj::objsequence;
@@ -39,7 +40,7 @@ pub struct VirtualMachine {
     pub sys_module: PyObjectRef,
     pub stdlib_inits: HashMap<String, stdlib::StdlibInitFunc>,
     pub ctx: PyContext,
-    pub current_frame: Option<PyObjectRef>,
+    pub frames: Vec<PyObjectRef>,
     pub wasm_id: Option<String>,
 }
 
@@ -62,14 +63,29 @@ impl VirtualMachine {
             sys_module: sysmod,
             stdlib_inits,
             ctx,
-            current_frame: None,
+            frames: vec![],
             wasm_id: None,
         }
     }
 
     pub fn run_code_obj(&mut self, code: PyObjectRef, scope: PyObjectRef) -> PyResult {
-        let mut frame = Frame::new(code, scope);
-        frame.run_frame_full(self)
+        let frame = self.ctx.new_frame(code, scope);
+        self.run_frame_full(frame)
+    }
+
+    pub fn run_frame_full(&mut self, frame: PyObjectRef) -> PyResult {
+        match self.run_frame(frame)? {
+            ExecutionResult::Return(value) => Ok(value),
+            _ => panic!("Got unexpected result from function"),
+        }
+    }
+
+    pub fn run_frame(&mut self, frame: PyObjectRef) -> Result<ExecutionResult, PyObjectRef> {
+        self.frames.push(frame.clone());
+        let mut frame = objframe::get_value(&frame);
+        let result = frame.run(self);
+        self.frames.pop();
+        result
     }
 
     /// Create a new python string object.
@@ -312,13 +328,13 @@ impl VirtualMachine {
         self.fill_scope_from_args(&code_object, &scope, args, defaults)?;
 
         // Construct frame:
-        let mut frame = Frame::new(code.clone(), scope);
+        let frame = self.ctx.new_frame(code.clone(), scope);
 
         // If we have a generator, create a new generator
         if code_object.is_generator {
             objgenerator::new_generator(self, frame)
         } else {
-            frame.run_frame_full(self)
+            self.run_frame_full(frame)
         }
     }
 
