Proper construction of scope for exec/eval.

cthulahoops · cthulahoops · commit c5789a03a32d · 2019-03-04T14:32:52.000Z
diff --git a/tests/snippets/test_exec.py b/tests/snippets/test_exec.py
@@ -0,0 +1,12 @@
+exec("def square(x):\n return x * x\n")
+assert 16 == square(4)
+
+d = {}
+exec("def square(x):\n return x * x\n", {}, d)
+assert 16 == d['square'](4)
+
+exec("assert 2 == x", {}, {'x': 2})
+exec("assert 2 == x", {'x': 2}, {})
+exec("assert 4 == x", {'x': 2}, {'x': 4})
+
+exec("assert max(1, 2) == 2", {}, {})
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -18,7 +18,6 @@ use crate::frame::{Scope, ScopeRef};
 use crate::pyobject::{
     AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol,
 };
-use std::rc::Rc;
 
 #[cfg(not(target_arch = "wasm32"))]
 use crate::stdlib::io::io_open;
@@ -192,7 +191,7 @@ fn builtin_eval(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         args,
         required = [(source, None)],
         optional = [
-            (_globals, Some(vm.ctx.dict_type())),
+            (globals, Some(vm.ctx.dict_type())),
             (locals, Some(vm.ctx.dict_type()))
         ]
     );
@@ -215,7 +214,7 @@ fn builtin_eval(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         return Err(vm.new_type_error("code argument must be str or code object".to_string()));
     };
 
-    let scope = make_scope(vm, locals);
+    let scope = make_scope(vm, globals, locals);
 
     // Run the source:
     vm.run_code_obj(code_obj.clone(), scope)
@@ -229,7 +228,7 @@ fn builtin_exec(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         args,
         required = [(source, None)],
         optional = [
-            (_globals, Some(vm.ctx.dict_type())),
+            (globals, Some(vm.ctx.dict_type())),
             (locals, Some(vm.ctx.dict_type()))
         ]
     );
@@ -252,26 +251,28 @@ fn builtin_exec(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         return Err(vm.new_type_error("source argument must be str or code object".to_string()));
     };
 
-    let scope = make_scope(vm, locals);
+    let scope = make_scope(vm, globals, locals);
 
     // Run the code:
     vm.run_code_obj(code_obj, scope)
 }
 
-fn make_scope(vm: &mut VirtualMachine, locals: Option<&PyObjectRef>) -> ScopeRef {
-    // handle optional global and locals
-    let locals = if let Some(locals) = locals {
-        locals.clone()
-    } else {
-        vm.new_dict()
+fn make_scope(
+    vm: &mut VirtualMachine,
+    globals: Option<&PyObjectRef>,
+    locals: Option<&PyObjectRef>,
+) -> ScopeRef {
+    let current_scope = vm.current_scope();
+    let parent = match globals {
+        Some(dict) => Some(Scope::new(dict.clone(), Some(vm.get_builtin_scope()))),
+        None => current_scope.parent.clone(),
+    };
+    let locals = match locals {
+        Some(dict) => dict.clone(),
+        None => current_scope.locals.clone(),
     };
 
-    // TODO: handle optional globals
-    // Construct new scope:
-    Rc::new(Scope {
-        locals,
-        parent: None,
-    })
+    Scope::new(locals, parent)
 }
 
 fn builtin_format(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -35,6 +35,12 @@ pub struct Scope {
 }
 pub type ScopeRef = Rc<Scope>;
 
+impl Scope {
+    pub fn new(locals: PyObjectRef, parent: Option<ScopeRef>) -> ScopeRef {
+        Rc::new(Scope { locals, parent })
+    }
+}
+
 #[derive(Clone, Debug)]
 struct Block {
     /// The type of block.
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -90,6 +90,12 @@ impl VirtualMachine {
         result
     }
 
+    pub fn current_scope(&self) -> &ScopeRef {
+        let current_frame = &self.frames[self.frames.len() - 1];
+        let frame = objframe::get_value(current_frame);
+        &frame.scope
+    }
+
     /// Create a new python string object.
     pub fn new_str(&self, s: String) -> PyObjectRef {
         self.ctx.new_str(s)
@@ -218,7 +224,7 @@ impl VirtualMachine {
         &self.ctx
     }
 
-    pub fn get_builtin_scope(&mut self) -> ScopeRef {
+    pub fn get_builtin_scope(&self) -> ScopeRef {
         let a2 = &*self.builtins;
         match a2.payload {
             PyObjectPayload::Module { ref scope, .. } => scope.clone(),
