From f820aeb1ea46c46a0f302edfc14fbece2cfb540c Mon Sep 17 00:00:00 2001
From: Joey Hain <jmhain@protonmail.com>
Date: Mon, 4 Mar 2019 08:46:40 -0800
Subject: [PATCH] Convert complex payload

---
 vm/src/obj/objcomplex.rs | 30 +++++++++++++++++++++++-------
 vm/src/pyobject.rs       | 28 +++++++++++++++-------------
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/vm/src/obj/objcomplex.rs b/vm/src/obj/objcomplex.rs
index c6be5e5f80..66205559c9 100644
--- a/vm/src/obj/objcomplex.rs
+++ b/vm/src/obj/objcomplex.rs
@@ -2,12 +2,30 @@ use super::objfloat;
 use super::objint;
 use super::objtype;
 use crate::pyobject::{
-    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectPayload2, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use crate::vm::VirtualMachine;
 use num_complex::Complex64;
 use num_traits::ToPrimitive;
 
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct PyComplex {
+    value: Complex64,
+}
+
+impl PyObjectPayload2 for PyComplex {
+    fn required_type(ctx: &PyContext) -> PyObjectRef {
+        ctx.complex_type()
+    }
+}
+
+impl From<Complex64> for PyComplex {
+    fn from(value: Complex64) -> Self {
+        PyComplex { value }
+    }
+}
+
 pub fn init(context: &PyContext) {
     let complex_type = &context.complex_type;
 
@@ -45,11 +63,7 @@ pub fn init(context: &PyContext) {
 }
 
 pub fn get_value(obj: &PyObjectRef) -> Complex64 {
-    if let PyObjectPayload::Complex { value } = &obj.payload {
-        *value
-    } else {
-        panic!("Inner error getting complex");
-    }
+    obj.payload::<PyComplex>().unwrap().value
 }
 
 fn complex_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -77,7 +91,9 @@ fn complex_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let value = Complex64::new(real, imag);
 
     Ok(PyObject::new(
-        PyObjectPayload::Complex { value },
+        PyObjectPayload::AnyRustValue {
+            value: Box::new(PyComplex { value }),
+        },
         cls.clone(),
     ))
 }
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
index 2e07e30f6e..e581b1fae0 100644
--- a/vm/src/pyobject.rs
+++ b/vm/src/pyobject.rs
@@ -5,6 +5,11 @@ use std::iter;
 use std::ops::RangeInclusive;
 use std::rc::{Rc, Weak};
 
+use num_bigint::BigInt;
+use num_bigint::ToBigInt;
+use num_complex::Complex64;
+use num_traits::{One, Zero};
+
 use crate::bytecode;
 use crate::exceptions;
 use crate::frame::{Frame, Scope, ScopeRef};
@@ -12,11 +17,11 @@ use crate::obj::objbool;
 use crate::obj::objbytearray;
 use crate::obj::objbytes;
 use crate::obj::objcode;
-use crate::obj::objcomplex;
+use crate::obj::objcomplex::{self, PyComplex};
 use crate::obj::objdict;
 use crate::obj::objenumerate;
 use crate::obj::objfilter;
-use crate::obj::objfloat;
+use crate::obj::objfloat::{self, PyFloat};
 use crate::obj::objframe;
 use crate::obj::objfunction;
 use crate::obj::objgenerator;
@@ -37,10 +42,6 @@ use crate::obj::objtuple;
 use crate::obj::objtype;
 use crate::obj::objzip;
 use crate::vm::VirtualMachine;
-use num_bigint::BigInt;
-use num_bigint::ToBigInt;
-use num_complex::Complex64;
-use num_traits::{One, Zero};
 
 /* Python objects and references.
 
@@ -463,14 +464,19 @@ impl PyContext {
     pub fn new_float(&self, value: f64) -> PyObjectRef {
         PyObject::new(
             PyObjectPayload::AnyRustValue {
-                value: Box::new(objfloat::PyFloat::from(value)),
+                value: Box::new(PyFloat::from(value)),
             },
             self.float_type(),
         )
     }
 
-    pub fn new_complex(&self, i: Complex64) -> PyObjectRef {
-        PyObject::new(PyObjectPayload::Complex { value: i }, self.complex_type())
+    pub fn new_complex(&self, value: Complex64) -> PyObjectRef {
+        PyObject::new(
+            PyObjectPayload::AnyRustValue {
+                value: Box::new(PyComplex::from(value)),
+            },
+            self.complex_type(),
+        )
     }
 
     pub fn new_str(&self, s: String) -> PyObjectRef {
@@ -1408,9 +1414,6 @@ pub enum PyObjectPayload {
     Integer {
         value: BigInt,
     },
-    Complex {
-        value: Complex64,
-    },
     Sequence {
         elements: RefCell<Vec<PyObjectRef>>,
     },
@@ -1494,7 +1497,6 @@ impl fmt::Debug for PyObjectPayload {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             PyObjectPayload::Integer { ref value } => write!(f, "int {}", value),
-            PyObjectPayload::Complex { ref value } => write!(f, "complex {}", value),
             PyObjectPayload::MemoryView { ref obj } => write!(f, "bytes/bytearray {:?}", obj),
             PyObjectPayload::Sequence { .. } => write!(f, "list or tuple"),
             PyObjectPayload::Dict { .. } => write!(f, "dict"),