Add complex basic type

Author: windelbouwman

Cargo.lock

@@ -300,4 +300,5 @@ pub enum Comparison {
 pub enum Number {
     Integer { value: i32 },
     Float { value: f64 },
+    Complex { real: f64, imag: f64 },
 }

parser/src/ast.rs

@@ -3,6 +3,7 @@
 pub enum Tok {
     Name { name: String },
     Number { value: String },
+    Complex { real: f64, imag: f64 },
     String { value: String },
     Bytes { value: Vec<u8> },
     Newline,

parser/src/token.rs

@@ -37,3 +37,7 @@
     pass
 
 assert int() == 0
+
+a = complex(2, 4)
+assert type(a) is complex
+assert type(a + a) is complex

tests/snippets/basic_types.py

@@ -5,6 +5,7 @@ authors = ["Shing Lyu <[email protected]>"]
 
 [dependencies]
 bitflags = "1.0.4"
+num-complex = "0.2"
 log = "0.3"
 rustpython_parser = {path = "../parser"}
 serde = "1.0.66"

vm/Cargo.toml

@@ -126,7 +126,6 @@ fn builtin_compile(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     compile::compile(vm, &source, mode, None)
 }
 
-// builtin_complex
 // builtin_delattr
 
 fn builtin_dir(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -543,6 +542,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("bytes"), ctx.bytes_type());
     dict.insert(String::from("chr"), ctx.new_rustfunc(builtin_chr));
     dict.insert(String::from("compile"), ctx.new_rustfunc(builtin_compile));
+    dict.insert(String::from("complex"), ctx.complex_type());
     dict.insert(String::from("dict"), ctx.dict_type());
     dict.insert(String::from("divmod"), ctx.new_rustfunc(builtin_divmod));
     dict.insert(String::from("dir"), ctx.new_rustfunc(builtin_dir));

vm/src/builtins.rs

@@ -10,9 +10,9 @@ let call_function = 0x64;
 /*
  * Primitive instruction type, which can be encoded and decoded.
  */
-extern crate rustpython_parser;
 
-use self::rustpython_parser::ast;
+use num_complex::Complex64;
+use rustpython_parser::ast;
 use std::collections::HashMap;
 use std::fmt;
 
@@ -197,6 +197,7 @@ pub enum CallType {
 pub enum Constant {
     Integer { value: i32 }, // TODO: replace by arbitrary big int math.
     Float { value: f64 },
+    Complex { value: Complex64 },
     Boolean { value: bool },
     String { value: String },
     Bytes { value: Vec<u8> },

vm/src/bytecode.rs

@@ -6,12 +6,11 @@
  *   https://github.com/micropython/micropython/blob/master/py/compile.c
  */
 
-extern crate rustpython_parser;
-
-use self::rustpython_parser::{ast, parser};
 use super::bytecode::{self, CallType, CodeObject, Instruction};
 use super::pyobject::{PyObject, PyObjectKind, PyResult};
 use super::vm::VirtualMachine;
+use num_complex::Complex64;
+use rustpython_parser::{ast, parser};
 
 struct Compiler {
     code_object_stack: Vec<CodeObject>,
@@ -846,6 +845,9 @@ impl Compiler {
                 let const_value = match value {
                     ast::Number::Integer { value } => bytecode::Constant::Integer { value: *value },
                     ast::Number::Float { value } => bytecode::Constant::Float { value: *value },
+                    ast::Number::Complex { real, imag } => bytecode::Constant::Complex {
+                        value: Complex64::new(*real, *imag),
+                    },
                 };
                 self.emit(Instruction::LoadConst { value: const_value });
             }
@@ -1301,8 +1303,8 @@ mod tests {
     use super::bytecode::CodeObject;
     use super::bytecode::Constant::*;
     use super::bytecode::Instruction::*;
-    use super::rustpython_parser::parser;
     use super::Compiler;
+    use rustpython_parser::parser;
     fn compile_exec(source: &str) -> CodeObject {
         let mut compiler = Compiler::new();
         compiler.push_new_code_object(Option::None, "<module>".to_string());

vm/src/compile.rs

@@ -1074,6 +1074,7 @@ impl Frame {
         match *value {
             bytecode::Constant::Integer { ref value } => vm.ctx.new_int(*value),
             bytecode::Constant::Float { ref value } => vm.ctx.new_float(*value),
+            bytecode::Constant::Complex { ref value } => vm.ctx.new_complex(*value),
             bytecode::Constant::String { ref value } => vm.new_str(value.clone()),
             bytecode::Constant::Bytes { ref value } => vm.ctx.new_bytes(value.clone()),
             bytecode::Constant::Boolean { ref value } => vm.new_bool(value.clone()),

vm/src/frame.rs

@@ -3,9 +3,12 @@ extern crate bitflags;
 #[macro_use]
 extern crate log;
 // extern crate env_logger;
+extern crate num_complex;
 extern crate serde;
 extern crate serde_json;
 
+extern crate rustpython_parser;
+
 //extern crate eval; use eval::eval::*;
 // use py_code_object::{Function, NativeType, PyCodeObject};
 

vm/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod objbool;
 pub mod objbytes;
+pub mod objcomplex;
 pub mod objdict;
 pub mod objfloat;
 pub mod objfunction;

vm/src/obj/mod.rs

@@ -0,0 +1,71 @@
+use super::super::pyobject::{
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
+};
+use super::super::vm::VirtualMachine;
+use super::objfloat;
+use super::objtype;
+use num_complex::Complex64;
+
+pub fn init(context: &PyContext) {
+    let ref complex_type = context.complex_type;
+    complex_type.set_attr("__add__", context.new_rustfunc(complex_add));
+    complex_type.set_attr("__new__", context.new_rustfunc(complex_new));
+    complex_type.set_attr("__repr__", context.new_rustfunc(complex_repr));
+}
+
+pub fn get_value(obj: &PyObjectRef) -> Complex64 {
+    if let PyObjectKind::Complex { value } = &obj.borrow().kind {
+        *value
+    } else {
+        panic!("Inner error getting complex");
+    }
+}
+
+fn complex_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(cls, None)],
+        optional = [(real, None), (imag, None)]
+    );
+
+    if !objtype::issubclass(cls, &vm.ctx.complex_type()) {
+        return Err(vm.new_type_error(format!("{:?} is not a subtype of complex", cls)));
+    }
+
+    let real = match real {
+        None => 0.0,
+        Some(value) => objfloat::make_float(vm, value)?,
+    };
+
+    let imag = match imag {
+        None => 0.0,
+        Some(value) => objfloat::make_float(vm, value)?,
+    };
+
+    let value = Complex64::new(real, imag);
+
+    Ok(PyObject::new(PyObjectKind::Complex { value }, cls.clone()))
+}
+
+fn complex_add(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(i, Some(vm.ctx.complex_type())), (i2, None)]
+    );
+
+    let v1 = get_value(i);
+    if objtype::isinstance(i2, &vm.ctx.complex_type()) {
+        Ok(vm.ctx.new_complex(v1 + get_value(i2)))
+    } else {
+        Err(vm.new_type_error(format!("Cannot add {:?} and {:?}", i, i2)))
+    }
+}
+
+fn complex_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, Some(vm.ctx.complex_type()))]);
+    let v = get_value(obj);
+    Ok(vm.new_str(v.to_string()))
+}

vm/src/obj/objcomplex.rs

@@ -5,7 +5,7 @@ use super::super::vm::VirtualMachine;
 use super::objint;
 use super::objtype;
 
-fn float_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> Result<PyObjectRef, PyObjectRef> {
+fn float_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(float, Some(vm.ctx.float_type()))]);
     let v = get_value(float);
     Ok(vm.new_str(v.to_string()))
@@ -38,6 +38,16 @@ pub fn get_value(obj: &PyObjectRef) -> f64 {
     }
 }
 
+pub fn make_float(vm: &mut VirtualMachine, obj: &PyObjectRef) -> Result<f64, PyObjectRef> {
+    if objtype::isinstance(obj, &vm.ctx.float_type()) {
+        Ok(get_value(obj))
+    } else if objtype::isinstance(obj, &vm.ctx.int_type()) {
+        Ok(objint::get_value(obj) as f64)
+    } else {
+        Err(vm.new_type_error(format!("Cannot cast {:?} to float", obj)))
+    }
+}
+
 fn set_value(obj: &PyObjectRef, value: f64) {
     obj.borrow_mut().kind = PyObjectKind::Float { value };
 }

vm/src/obj/objfloat.rs

@@ -1,5 +1,6 @@
 use super::super::pyobject::{
-    AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objbool;
@@ -56,7 +57,10 @@ fn list_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         }
     };
 
-    Ok(vm.ctx.new_list(elements))
+    Ok(PyObject::new(
+        PyObjectKind::List { elements: elements },
+        cls.clone(),
+    ))
 }
 
 fn list_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {

vm/src/obj/objlist.rs

@@ -3,6 +3,7 @@ use super::exceptions;
 use super::frame::Frame;
 use super::obj::objbool;
 use super::obj::objbytes;
+use super::obj::objcomplex;
 use super::obj::objdict;
 use super::obj::objfloat;
 use super::obj::objfunction;
@@ -16,6 +17,7 @@ use super::obj::objstr;
 use super::obj::objtuple;
 use super::obj::objtype;
 use super::vm::VirtualMachine;
+use num_complex::Complex64;
 use std::cell::RefCell;
 use std::collections::HashMap;
 use std::fmt;
@@ -57,6 +59,7 @@ pub struct PyContext {
     pub dict_type: PyObjectRef,
     pub int_type: PyObjectRef,
     pub float_type: PyObjectRef,
+    pub complex_type: PyObjectRef,
     pub bytes_type: PyObjectRef,
     pub bool_type: PyObjectRef,
     pub true_value: PyObjectRef,
@@ -130,6 +133,7 @@ impl PyContext {
         let set_type = create_type("set", &type_type, &object_type, &dict_type);
         let int_type = create_type("int", &type_type, &object_type, &dict_type);
         let float_type = create_type("float", &type_type, &object_type, &dict_type);
+        let complex_type = create_type("complex", &type_type, &object_type, &dict_type);
         let bytes_type = create_type("bytes", &type_type, &object_type, &dict_type);
         let tuple_type = create_type("tuple", &type_type, &object_type, &dict_type);
         let iter_type = create_type("iter", &type_type, &object_type, &dict_type);
@@ -146,6 +150,7 @@ impl PyContext {
         let context = PyContext {
             int_type: int_type,
             float_type: float_type,
+            complex_type: complex_type,
             bytes_type: bytes_type,
             list_type: list_type,
             set_type: set_type,
@@ -176,6 +181,7 @@ impl PyContext {
         objgenerator::init(&context);
         objint::init(&context);
         objfloat::init(&context);
+        objcomplex::init(&context);
         objbytes::init(&context);
         objstr::init(&context);
         objtuple::init(&context);
@@ -193,6 +199,10 @@ impl PyContext {
         self.float_type.clone()
     }
 
+    pub fn complex_type(&self) -> PyObjectRef {
+        self.complex_type.clone()
+    }
+
     pub fn bytes_type(&self) -> PyObjectRef {
         self.bytes_type.clone()
     }
@@ -262,6 +272,10 @@ impl PyContext {
         PyObject::new(PyObjectKind::Float { value: i }, self.float_type())
     }
 
+    pub fn new_complex(&self, i: Complex64) -> PyObjectRef {
+        PyObject::new(PyObjectKind::Complex { value: i }, self.complex_type())
+    }
+
     pub fn new_str(&self, s: String) -> PyObjectRef {
         PyObject::new(PyObjectKind::String { value: s }, self.str_type())
     }
@@ -624,6 +638,9 @@ pub enum PyObjectKind {
     Float {
         value: f64,
     },
+    Complex {
+        value: Complex64,
+    },
     Bytes {
         value: Vec<u8>,
     },
@@ -690,6 +707,7 @@ impl fmt::Debug for PyObjectKind {
             &PyObjectKind::String { ref value } => write!(f, "str \"{}\"", value),
             &PyObjectKind::Integer { ref value } => write!(f, "int {}", value),
             &PyObjectKind::Float { ref value } => write!(f, "float {}", value),
+            &PyObjectKind::Complex { ref value } => write!(f, "complex {}", value),
             &PyObjectKind::Bytes { ref value } => write!(f, "bytes {:?}", value),
             &PyObjectKind::List { elements: _ } => write!(f, "list"),
             &PyObjectKind::Tuple { elements: _ } => write!(f, "tuple"),
@@ -740,6 +758,7 @@ impl PyObject {
             PyObjectKind::String { ref value } => value.clone(),
             PyObjectKind::Integer { ref value } => format!("{:?}", value),
             PyObjectKind::Float { ref value } => format!("{:?}", value),
+            PyObjectKind::Complex { ref value } => format!("{:?}", value),
             PyObjectKind::Bytes { ref value } => format!("b'{:?}'", value),
             PyObjectKind::List { ref elements } => format!(
                 "[{}]",

vm/src/pyobject.rs

@@ -13,6 +13,7 @@ use super::super::pyobject::{
     AttributeProtocol, DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::VirtualMachine;
+use num_complex::Complex64;
 use std::ops::Deref;
 
 /*
@@ -391,6 +392,9 @@ fn expression_to_ast(ctx: &PyContext, expression: &ast::Expression) -> PyObjectR
             let py_n = match value {
                 ast::Number::Integer { value } => ctx.new_int(*value),
                 ast::Number::Float { value } => ctx.new_float(*value),
+                ast::Number::Complex { real, imag } => {
+                    ctx.new_complex(Complex64::new(*real, *imag))
+                }
             };
             node.set_attr("n", py_n);