Keep a stack of frames on the VM.

cthulahoops · cthulahoops · commit ee86229ff6c6 · 2018-11-14T07:49:41.000Z
diff --git a/tests/snippets/getframe.py b/tests/snippets/getframe.py
@@ -9,7 +9,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(0).f_locals['x'] == 17
+    assert sys._getframe(1).f_locals['foo'] == 'bar'
 
 test_function()
 
@@ -18,3 +19,8 @@ def __init__(self):
         assert sys._getframe().f_locals['self'] == self
 
 TestClass()
+
+try:
+    sys._getframe(100)
+except ValueError:
+    pass
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -3,7 +3,6 @@ extern crate rustpython_parser;
 use self::rustpython_parser::ast;
 use std::collections::hash_map::HashMap;
 use std::fmt;
-use std::mem;
 use std::path::PathBuf;
 
 use super::builtins;
@@ -82,22 +81,13 @@ 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 = if let Some(source_path) = &self.code.source_path {
             source_path.to_string()
         } else {
             "<unknown>".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();
 
@@ -146,7 +136,6 @@ impl Frame {
             }
         };
 
-        vm.current_frame = prev_frame;
         value
     }
 
diff --git a/vm/src/obj/objgenerator.rs b/vm/src/obj/objgenerator.rs
@@ -48,7 +48,7 @@ fn generator_send(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 fn send(vm: &mut VirtualMachine, gen: &PyObjectRef, value: &PyObjectRef) -> PyResult {
     if let PyObjectKind::Generator { ref mut frame } = gen.borrow_mut().kind {
         frame.push_value(value.clone());
-        match frame.run_frame(vm)? {
+        match vm.run_frame(frame.clone())? {
             ExecutionResult::Yield(value) => Ok(value),
             ExecutionResult::Return(_value) => {
                 // Stop iteration!
diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -1,9 +1,11 @@
+use super::obj::objint;
+use super::obj::objtype;
+use super::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
+use super::vm::VirtualMachine;
 use num_bigint::ToBigInt;
-use obj::objtype;
-use pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol};
+use num_traits::ToPrimitive;
+use std::env;
 use std::rc::Rc;
-use std::{env, mem};
-use vm::VirtualMachine;
 
 /*
  * The magic sys module.
@@ -15,12 +17,31 @@ 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 getframe(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [],
+        optional = [(offset, Some(vm.ctx.int_type()))]
+    );
+    let idx = match offset {
+        Some(int) => {
+            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()));
+                }
+                offset
+            } else {
+                0
+            }
+        }
+        None => 0,
+    };
+
+    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
@@ -10,8 +10,9 @@ use std::collections::hash_map::HashMap;
 
 use super::builtins;
 use super::bytecode;
-use super::frame::Frame;
+use super::frame::{ExecutionResult, Frame};
 use super::obj::objcode::copy_code;
+use super::obj::objframe;
 use super::obj::objgenerator;
 use super::obj::objiter;
 use super::obj::objsequence;
@@ -35,7 +36,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>,
 }
 
 impl VirtualMachine {
@@ -52,13 +53,28 @@ impl VirtualMachine {
             sys_module: sysmod,
             stdlib_inits,
             ctx: ctx,
-            current_frame: None,
+            frames: vec![],
         }
     }
 
     pub fn run_code_obj(&mut self, code: PyObjectRef, scope: PyObjectRef) -> PyResult {
-        let mut frame = Frame::new(code, scope);
-        frame.run_frame_full(self)
+        self.run_frame_full(Frame::new(code, scope))
+    }
+
+    pub fn run_frame_full(&mut self, frame: Frame) -> PyResult {
+        match self.run_frame(frame)? {
+            ExecutionResult::Return(value) => Ok(value),
+            _ => panic!("Got unexpected result from function"),
+        }
+    }
+
+    pub fn run_frame(&mut self, frame: Frame) -> Result<ExecutionResult, PyObjectRef> {
+        let frame = self.ctx.new_frame(frame);
+        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.
@@ -260,13 +276,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 = Frame::new(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)
         }
     }
 
