function

OddCoincidence · OddCoincidence · commit dd56d1d5a215 · 2019-03-09T07:09:53.000-08:00
diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs
@@ -1,8 +1,34 @@
+use crate::frame::ScopeRef;
 use crate::pyobject::{
-    AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObjectPayload, PyResult, TypeProtocol,
+    AttributeProtocol, IdProtocol, PyContext, PyFuncArgs, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 
+#[derive(Debug)]
+pub struct PyFunction {
+    // TODO: these shouldn't be public
+    pub code: PyObjectRef,
+    pub scope: ScopeRef,
+    pub defaults: PyObjectRef,
+}
+
+impl PyFunction {
+    pub fn new(code: PyObjectRef, scope: ScopeRef, defaults: PyObjectRef) -> Self {
+        PyFunction {
+            code,
+            scope,
+            defaults,
+        }
+    }
+}
+
+impl PyObjectPayload2 for PyFunction {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.function_type()
+    }
+}
+
 pub fn init(context: &PyContext) {
     let function_type = &context.function_type;
     context.set_attr(&function_type, "__get__", context.new_rustfunc(bind_method));
@@ -79,9 +105,9 @@ fn bind_method(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 fn function_code(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    match args.args[0].payload {
-        PyObjectPayload::Function { ref code, .. } => Ok(code.clone()),
-        _ => Err(vm.new_type_error("no code".to_string())),
+    match args.args[0].payload() {
+        Some(PyFunction { ref code, .. }) => Ok(code.clone()),
+        None => Err(vm.new_type_error("no code".to_string())),
     }
 }
 
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -24,7 +24,7 @@ use crate::obj::objenumerate;
 use crate::obj::objfilter;
 use crate::obj::objfloat::{self, PyFloat};
 use crate::obj::objframe;
-use crate::obj::objfunction;
+use crate::obj::objfunction::{self, PyFunction};
 use crate::obj::objgenerator;
 use crate::obj::objint::{self, PyInt};
 use crate::obj::objiter;
@@ -657,10 +657,8 @@ impl PyContext {
         defaults: PyObjectRef,
     ) -> PyObjectRef {
         PyObject::new(
-            PyObjectPayload::Function {
-                code: code_obj,
-                scope,
-                defaults,
+            PyObjectPayload::AnyRustValue {
+                value: Box::new(PyFunction::new(code_obj, scope, defaults)),
             },
             self.function_type(),
         )
@@ -1510,11 +1508,6 @@ pub enum PyObjectPayload {
     Frame {
         frame: Frame,
     },
-    Function {
-        code: PyObjectRef,
-        scope: ScopeRef,
-        defaults: PyObjectRef,
-    },
     Generator {
         frame: PyObjectRef,
     },
@@ -1548,7 +1541,6 @@ impl fmt::Debug for PyObjectPayload {
             PyObjectPayload::WeakRef { .. } => write!(f, "weakref"),
             PyObjectPayload::Iterator { .. } => write!(f, "iterator"),
             PyObjectPayload::Slice { .. } => write!(f, "slice"),
-            PyObjectPayload::Function { .. } => write!(f, "function"),
             PyObjectPayload::Generator { .. } => write!(f, "generator"),
             PyObjectPayload::BoundMethod {
                 ref function,
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -18,6 +18,7 @@ use crate::frame::{Scope, ScopeRef};
 use crate::obj::objbool;
 use crate::obj::objcode;
 use crate::obj::objframe;
+use crate::obj::objfunction::PyFunction;
 use crate::obj::objgenerator;
 use crate::obj::objiter;
 use crate::obj::objlist::PyList;
@@ -289,13 +290,16 @@ impl VirtualMachine {
     {
         let args = args.into();
         trace!("Invoke: {:?} {:?}", func_ref, args);
+        if let Some(PyFunction {
+            ref code,
+            ref scope,
+            ref defaults,
+        }) = func_ref.payload()
+        {
+            return self.invoke_python_function(code, scope, defaults, args);
+        }
         match func_ref.payload {
             PyObjectPayload::RustFunction { ref function } => function(self, args),
-            PyObjectPayload::Function {
-                ref code,
-                ref scope,
-                ref defaults,
-            } => self.invoke_python_function(code, scope, defaults, args),
             PyObjectPayload::BoundMethod {
                 ref function,
                 ref object,
@@ -331,11 +335,11 @@ impl VirtualMachine {
     }
 
     pub fn invoke_with_locals(&mut self, function: PyObjectRef, locals: PyObjectRef) -> PyResult {
-        if let PyObjectPayload::Function {
+        if let Some(PyFunction {
             code,
             scope,
-            defaults: _defaults,
-        } = &function.payload
+            defaults: _,
+        }) = &function.payload()
         {
             let scope = Rc::new(Scope {
                 locals,
