foundation for complex numbers

src/analysis/type_analysis.cpp

@@ -413,6 +413,8 @@ private:
                 return FLOAT;
             case AST_Num::LONG:
                 return LONG;
+            case AST_Num::COMPLEX:
+                return BOXED_COMPLEX;
         }
         abort();
     }

src/codegen/compvars.cpp

@@ -1133,6 +1133,11 @@ ConcreteCompilerVariable* makeLong(IREmitter& emitter, std::string& n_long) {
     return new ConcreteCompilerVariable(LONG, v, true);
 }
 
+ConcreteCompilerVariable* makePureImaginary(IREmitter& emitter, double imag) {
+    llvm::Value* v = emitter.getBuilder()->CreateCall(g.funcs.createPureImaginary, getConstantDouble(imag));
+    return new ConcreteCompilerVariable(BOXED_COMPLEX, v, true);
+}
+
 class KnownClassobjType : public ValuedCompilerType<BoxedClass*> {
 private:
     BoxedClass* cls;
@@ -1169,7 +1174,6 @@ CompilerType* typeOfClassobj(BoxedClass* cls) {
     return KnownClassobjType::fromClass(cls);
 }
 
-
 class NormalObjectType : public ConcreteCompilerType {
 private:
     BoxedClass* cls;
@@ -1907,6 +1911,6 @@ ConcreteCompilerVariable* undefVariable() {
 }
 
 
-ConcreteCompilerType* LIST, *SLICE, *MODULE, *DICT, *SET, *FROZENSET, *LONG;
+ConcreteCompilerType* LIST, *SLICE, *MODULE, *DICT, *SET, *FROZENSET, *LONG, *BOXED_COMPLEX;
 
 } // namespace pyston

src/codegen/compvars.h

@@ -30,7 +30,7 @@ class CompilerType;
 class IREmitter;
 
 extern ConcreteCompilerType* INT, *BOXED_INT, *LONG, *FLOAT, *BOXED_FLOAT, *VOID, *UNKNOWN, *BOOL, *STR, *NONE, *LIST,
-    *SLICE, *MODULE, *DICT, *BOOL, *BOXED_BOOL, *BOXED_TUPLE, *SET, *FROZENSET, *CLOSURE, *GENERATOR;
+    *SLICE, *MODULE, *DICT, *BOOL, *BOXED_BOOL, *BOXED_TUPLE, *SET, *FROZENSET, *CLOSURE, *GENERATOR, *BOXED_COMPLEX;
 extern CompilerType* UNDEF;
 
 class CompilerType {
@@ -342,6 +342,7 @@ ConcreteCompilerVariable* makeInt(int64_t);
 ConcreteCompilerVariable* makeFloat(double);
 ConcreteCompilerVariable* makeBool(bool);
 ConcreteCompilerVariable* makeLong(IREmitter& emitter, std::string&);
+ConcreteCompilerVariable* makePureImaginary(IREmitter& emitter, double imag);
 CompilerVariable* makeStr(const std::string*);
 CompilerVariable* makeFunction(IREmitter& emitter, CLFunction*, CompilerVariable* closure, bool isGenerator,
                                const std::vector<ConcreteCompilerVariable*>& defaults);

src/codegen/irgen/irgenerator.cpp

@@ -758,6 +758,8 @@ private:
             return makeInt(node->n_int);
         else if (node->num_type == AST_Num::FLOAT)
             return makeFloat(node->n_float);
+        else if (node->num_type == AST_Num::COMPLEX)
+            return makePureImaginary(emitter, node->n_float);
         else
             return makeLong(emitter, node->n_long);
     }

src/codegen/irgen/util.cpp

@@ -107,6 +107,10 @@ llvm::Constant* getConstantInt(int n) {
     return getConstantInt(n, g.i64);
 }
 
+llvm::Constant* getConstantDouble(double val) {
+    return llvm::ConstantFP::get(g.double_, val);
+}
+
 class PrettifyingMaterializer : public llvm::ValueMaterializer {
 private:
     llvm::Module* module;

src/codegen/irgen/util.h

@@ -29,6 +29,7 @@ llvm::Constant* getStringConstantPtr(const std::string& str);
 llvm::Constant* getStringConstantPtr(const char* str);
 llvm::Constant* embedConstantPtr(const void* addr, llvm::Type*);
 llvm::Constant* getConstantInt(int val);
+llvm::Constant* getConstantDouble(double val);
 llvm::Constant* getConstantInt(int val, llvm::Type*);
 
 void dumpPrettyIR(llvm::Function* f);

src/codegen/parse_ast.py

@@ -117,6 +117,8 @@ def convert(n, f):
             f.write('\x30')
         elif isinstance(n.n, float):
             f.write('\x20')
+        elif isinstance(n.n, complex):
+            f.write('\x40')
         else:
             raise Exception(type(n.n))
 
@@ -160,6 +162,8 @@ def convert(n, f):
             _print_str(str(v), f)
         elif isinstance(v, float):
             f.write(struct.pack(">d", v))
+        elif isinstance(v, complex):
+            f.write(struct.pack(">d", v.imag))
         elif v is None or isinstance(v, _ast.AST):
             convert(v, f)
         else:

src/codegen/parser.cpp

@@ -532,6 +532,8 @@ AST_Num* read_num(BufferedReader* reader) {
         rtn->n_long = readString(reader);
     } else if (rtn->num_type == AST_Num::FLOAT) {
         rtn->n_float = reader->readDouble();
+    } else if (rtn->num_type == AST_Num::COMPLEX) {
+        rtn->n_float = reader->readDouble();
     } else {
         RELEASE_ASSERT(0, "%d", rtn->num_type);
     }

src/codegen/runtime_hooks.cpp

@@ -32,6 +32,7 @@
 #include "codegen/irgen/util.h"
 #include "core/threading.h"
 #include "core/types.h"
+#include "runtime/complex.h"
 #include "runtime/float.h"
 #include "runtime/generator.h"
 #include "runtime/inline/boxing.h"
@@ -171,6 +172,7 @@ void initGlobalFuncs(GlobalState& g) {
     GET(createClosure);
     GET(createGenerator);
     GET(createLong);
+    GET(createPureImaginary);
     GET(createSet);
 
     GET(getattr);

src/codegen/runtime_hooks.h

@@ -34,7 +34,7 @@ struct GlobalFuncs {
 
     llvm::Value* boxInt, *unboxInt, *boxFloat, *unboxFloat, *boxStringPtr, *boxCLFunction, *unboxCLFunction,
         *boxInstanceMethod, *boxBool, *unboxBool, *createTuple, *createDict, *createList, *createSlice,
-        *createUserClass, *createClosure, *createGenerator, *createLong, *createSet;
+        *createUserClass, *createClosure, *createGenerator, *createLong, *createSet, *createPureImaginary;
     llvm::Value* getattr, *setattr, *delattr, *delitem, *delGlobal, *print, *nonzero, *binop, *compare, *augbinop,
         *unboxedLen, *getitem, *getclsattr, *getGlobal, *setitem, *unaryop, *import, *importFrom, *importStar, *repr,
         *str, *isinstance, *yield, *getiter;

src/core/ast.cpp

@@ -1431,6 +1431,8 @@ bool PrintVisitor::visit_num(AST_Num* node) {
         printf("%sL", node->n_long.c_str());
     } else if (node->num_type == AST_Num::FLOAT) {
         printf("%f", node->n_float);
+    } else if (node->num_type == AST_Num::COMPLEX) {
+        printf("%fj", node->n_float);
     } else {
         RELEASE_ASSERT(0, "");
     }

src/core/ast.h

@@ -621,6 +621,9 @@ public:
         INT = 0x10,
         FLOAT = 0x20,
         LONG = 0x30,
+
+        // for COMPLEX, n_float is the imaginary part, real part is 0
+        COMPLEX = 0x40,
     } num_type;
 
     union {

src/runtime/complex.cpp

+// Copyright (c) 2014 Dropbox, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "runtime/complex.h"
+
+#include "codegen/compvars.h"
+#include "core/types.h"
+#include "runtime/float.h"
+#include "runtime/inline/boxing.h"
+#include "runtime/objmodel.h"
+#include "runtime/types.h"
+
+namespace pyston {
+
+static inline void raiseDivZeroExc() {
+    raiseExcHelper(ZeroDivisionError, "complex divide by zero");
+}
+
+extern "C" Box* createPureImaginary(double i) {
+    return new BoxedComplex(0.0, i);
+}
+
+// addition
+
+extern "C" Box* complexAddComplex(BoxedComplex* lhs, BoxedComplex* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == complex_cls);
+    return boxComplex(lhs->real + rhs->real, lhs->imag + rhs->imag);
+}
+
+extern "C" Box* complexAddFloat(BoxedComplex* lhs, BoxedFloat* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == float_cls);
+    return boxComplex(lhs->real + rhs->d, lhs->imag);
+}
+
+extern "C" Box* complexAddInt(BoxedComplex* lhs, BoxedInt* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == int_cls);
+    return boxComplex(lhs->real + (double)rhs->n, lhs->imag);
+}
+
+extern "C" Box* complexAdd(BoxedComplex* lhs, Box* rhs) {
+    assert(lhs->cls == complex_cls);
+    if (rhs->cls == int_cls) {
+        return complexAddInt(lhs, static_cast<BoxedInt*>(rhs));
+    } else if (rhs->cls == float_cls) {
+        return complexAddFloat(lhs, static_cast<BoxedFloat*>(rhs));
+    } else if (rhs->cls == complex_cls) {
+        return complexAddComplex(lhs, static_cast<BoxedComplex*>(rhs));
+    } else {
+        return NotImplemented;
+    }
+}
+
+// subtraction
+
+extern "C" Box* complexSubComplex(BoxedComplex* lhs, BoxedComplex* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == complex_cls);
+    return boxComplex(lhs->real - rhs->real, lhs->imag - rhs->imag);
+}
+
+extern "C" Box* complexSubFloat(BoxedComplex* lhs, BoxedFloat* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == float_cls);
+    return boxComplex(lhs->real - rhs->d, lhs->imag);
+}
+
+extern "C" Box* complexSubInt(BoxedComplex* lhs, BoxedInt* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == int_cls);
+    return boxComplex(lhs->real - (double)rhs->n, lhs->imag);
+}
+
+extern "C" Box* complexSub(BoxedComplex* lhs, Box* rhs) {
+    assert(lhs->cls == complex_cls);
+    if (rhs->cls == int_cls) {
+        return complexSubInt(lhs, static_cast<BoxedInt*>(rhs));
+    } else if (rhs->cls == float_cls) {
+        return complexSubFloat(lhs, static_cast<BoxedFloat*>(rhs));
+    } else if (rhs->cls == complex_cls) {
+        return complexSubComplex(lhs, static_cast<BoxedComplex*>(rhs));
+    } else {
+        return NotImplemented;
+    }
+}
+
+// multiplication
+
+extern "C" Box* complexMulComplex(BoxedComplex* lhs, BoxedComplex* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == complex_cls);
+    return boxComplex(lhs->real * rhs->real - lhs->imag * rhs->imag, lhs->real * rhs->imag + lhs->imag * rhs->real);
+}
+
+extern "C" Box* complexMulFloat(BoxedComplex* lhs, BoxedFloat* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == float_cls);
+    return boxComplex(lhs->real * rhs->d, lhs->imag * rhs->d);
+}
+
+extern "C" Box* complexMulInt(BoxedComplex* lhs, BoxedInt* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == int_cls);
+    return boxComplex(lhs->real * (double)rhs->n, lhs->imag * (double)rhs->n);
+}
+
+extern "C" Box* complexMul(BoxedComplex* lhs, Box* rhs) {
+    assert(lhs->cls == complex_cls);
+    if (rhs->cls == int_cls) {
+        return complexMulInt(lhs, static_cast<BoxedInt*>(rhs));
+    } else if (rhs->cls == float_cls) {
+        return complexMulFloat(lhs, static_cast<BoxedFloat*>(rhs));
+    } else if (rhs->cls == complex_cls) {
+        return complexMulComplex(lhs, static_cast<BoxedComplex*>(rhs));
+    } else {
+        return NotImplemented;
+    }
+}
+
+// division
+
+extern "C" Box* complexDivComplex(BoxedComplex* lhs, BoxedComplex* rhs) {
+    // TODO implement this
+    // NOTE: the "naive" implementation of complex division has numerical issues
+    // see notes in CPython, Objects/complexobject.c, c_quot
+    return NotImplemented;
+}
+
+extern "C" Box* complexDivFloat(BoxedComplex* lhs, BoxedFloat* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == float_cls);
+    if (rhs->d == 0.0) {
+        raiseDivZeroExc();
+    }
+    return boxComplex(lhs->real / rhs->d, lhs->imag / rhs->d);
+}
+
+extern "C" Box* complexDivInt(BoxedComplex* lhs, BoxedInt* rhs) {
+    assert(lhs->cls == complex_cls);
+    assert(rhs->cls == int_cls);
+    if (rhs->n == 0) {
+        raiseDivZeroExc();
+    }
+    return boxComplex(lhs->real / (float)rhs->n, lhs->imag / (float)rhs->n);
+}
+
+extern "C" Box* complexDiv(BoxedComplex* lhs, Box* rhs) {
+    assert(lhs->cls == complex_cls);
+    if (rhs->cls == int_cls) {
+        return complexDivInt(lhs, static_cast<BoxedInt*>(rhs));
+    } else if (rhs->cls == float_cls) {
+        return complexDivFloat(lhs, static_cast<BoxedFloat*>(rhs));
+    } else if (rhs->cls == complex_cls) {
+        return complexDivComplex(lhs, static_cast<BoxedComplex*>(rhs));
+    } else {
+        return NotImplemented;
+    }
+}
+
+// str and repr
+// For now, just print the same way as ordinary doubles.
+// TODO this is wrong, e.g. if real or imaginary part is an integer, there should
+// be no decimal point, maybe some other differences. Need to dig deeper into
+// how CPython formats floats and complex numbers.
+// (complex_format in Objects/complexobject.c)
+std::string complexFmt(double r, double i, int precision, char code) {
+    if (r == 0. && copysign(1.0, r) == 1.0) {
+        return floatFmt(i, precision, code) + "j";
+    } else {
+        return "(" + floatFmt(r, precision, code) + (isnan(i) || i >= 0.0 ? "+" : "") + floatFmt(i, precision, code)
+               + "j)";
+    }
+}
+
+static void _addFunc(const char* name, ConcreteCompilerType* rtn_type, void* complex_func, void* float_func,
+                     void* int_func, void* boxed_func) {
+    CLFunction* cl = createRTFunction(2, 0, false, false);
+    addRTFunction(cl, complex_func, rtn_type, { BOXED_COMPLEX, BOXED_COMPLEX });
+    addRTFunction(cl, float_func, rtn_type, { BOXED_COMPLEX, BOXED_FLOAT });
+    addRTFunction(cl, int_func, rtn_type, { BOXED_COMPLEX, BOXED_INT });
+    addRTFunction(cl, boxed_func, UNKNOWN, { BOXED_COMPLEX, UNKNOWN });
+    complex_cls->giveAttr(name, new BoxedFunction(cl));
+}
+
+Box* complexStr(BoxedComplex* self) {
+    assert(self->cls == complex_cls);
+    return boxString(complexFmt(self->real, self->imag, 12, 'g'));
+}
+
+Box* complexRepr(BoxedComplex* self) {
+    assert(self->cls == complex_cls);
+    return boxString(complexFmt(self->real, self->imag, 16, 'g'));
+}
+
+void setupComplex() {
+    complex_cls->giveAttr("__name__", boxStrConstant("complex"));
+
+    _addFunc("__add__", BOXED_COMPLEX, (void*)complexAddComplex, (void*)complexAddFloat, (void*)complexAddInt,
+             (void*)complexAdd);
+
+    _addFunc("__sub__", BOXED_COMPLEX, (void*)complexSubComplex, (void*)complexSubFloat, (void*)complexSubInt,
+             (void*)complexSub);
+
+    _addFunc("__mul__", BOXED_COMPLEX, (void*)complexMulComplex, (void*)complexMulFloat, (void*)complexMulInt,
+             (void*)complexMul);
+
+    _addFunc("__div__", BOXED_COMPLEX, (void*)complexDivComplex, (void*)complexDivFloat, (void*)complexDivInt,
+             (void*)complexDiv);
+
+    complex_cls->giveAttr("__str__", new BoxedFunction(boxRTFunction((void*)complexStr, STR, 1)));
+    complex_cls->giveAttr("__repr__", new BoxedFunction(boxRTFunction((void*)complexRepr, STR, 1)));
+    complex_cls->freeze();
+}
+
+void teardownComplex() {
+}
+}

src/runtime/complex.h

+// Copyright (c) 2014 Dropbox, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef PYSTON_RUNTIME_COMPLEX_H
+#define PYSTON_RUNTIME_COMPLEX_H
+
+#include "core/types.h"
+
+namespace pyston {
+
+extern "C" Box* createPureImaginary(double i);
+
+extern "C" double mod_complex_complex(double lhs, double rhs);
+extern "C" double div_complex_complex(double lhs, double rhs);
+extern "C" double floordiv_complex_complex(double lhs, double rhs);
+extern "C" double pow_complex_complex(double lhs, double rhs);
+}
+
+#endif

src/runtime/float.h

@@ -24,6 +24,8 @@ extern "C" double pow_float_float(double lhs, double rhs);
 
 class BoxedFloat;
 bool floatNonzeroUnboxed(BoxedFloat* self);
+
+std::string floatFmt(double x, int precision, char code);
 }
 
 #endif

src/runtime/inline/boxing.h

@@ -26,6 +26,11 @@ extern "C" inline Box* boxFloat(double d) {
     return new BoxedFloat(d);
 }
 
+extern "C" inline Box* boxComplex(double r, double i) __attribute__((visibility("default")));
+extern "C" inline Box* boxComplex(double r, double i) {
+    return new BoxedComplex(r, i);
+}
+
 extern "C" inline Box* boxBool(bool b) __attribute__((visibility("default")));
 extern "C" inline Box* boxBool(bool b) {
     Box* rtn = b ? True : False;

src/runtime/inline/link_forcer.cpp

@@ -17,6 +17,7 @@
 
 #include "core/types.h"
 #include "gc/heap.h"
+#include "runtime/complex.h"
 #include "runtime/float.h"
 #include "runtime/generator.h"
 #include "runtime/inline/boxing.h"
@@ -63,6 +64,7 @@ void force() {
     FORCE(createClosure);
     FORCE(createGenerator);
     FORCE(createLong);
+    FORCE(createPureImaginary);
     FORCE(createSet);
 
     FORCE(getattr);

src/runtime/types.cpp

@@ -294,7 +294,7 @@ extern "C" void closureGCHandler(GCVisitor* v, Box* b) {
 extern "C" {
 BoxedClass* object_cls, *type_cls, *none_cls, *bool_cls, *int_cls, *float_cls, *str_cls, *function_cls,
     *instancemethod_cls, *list_cls, *slice_cls, *module_cls, *dict_cls, *tuple_cls, *file_cls, *member_cls,
-    *closure_cls, *generator_cls;
+    *closure_cls, *generator_cls, *complex_cls;
 
 
 BoxedTuple* EmptyTuple;
@@ -665,6 +665,7 @@ void setupRuntime() {
     // We could have a multi-stage setup process, but that seems overkill for now.
     bool_cls = new BoxedClass(object_cls, NULL, 0, sizeof(BoxedBool), false);
     int_cls = new BoxedClass(object_cls, NULL, 0, sizeof(BoxedInt), false);
+    complex_cls = new BoxedClass(object_cls, NULL, 0, sizeof(BoxedComplex), false);
     // TODO we're leaking long memory!
     long_cls = new BoxedClass(object_cls, NULL, 0, sizeof(BoxedLong), false);
     float_cls = new BoxedClass(object_cls, NULL, 0, sizeof(BoxedFloat), false);
@@ -695,6 +696,7 @@ void setupRuntime() {
     FROZENSET = typeFromClass(frozenset_cls);
     BOXED_TUPLE = typeFromClass(tuple_cls);
     LONG = typeFromClass(long_cls);
+    BOXED_COMPLEX = typeFromClass(complex_cls);
 
     object_cls->giveAttr("__name__", boxStrConstant("object"));
     object_cls->giveAttr("__new__", new BoxedFunction(boxRTFunction((void*)objectNew, UNKNOWN, 1, 0, true, false)));
@@ -736,6 +738,7 @@ void setupRuntime() {
     setupInt();
     setupLong();
     setupFloat();
+    setupComplex();
     setupStr();
     setupList();
     setupDict();
@@ -832,6 +835,7 @@ void teardownRuntime() {
     teardownList();
     teardownInt();
     teardownFloat();
+    teardownComplex();
     teardownStr();
     teardownBool();
     teardownDict();

src/runtime/types.h

@@ -41,6 +41,8 @@ void setupInt();
 void teardownInt();
 void setupFloat();
 void teardownFloat();
+void setupComplex();
+void teardownComplex();
 void setupStr();
 void teardownStr();
 void setupList();
@@ -76,7 +78,7 @@ Box* getSysStdout();
 extern "C" {
 extern BoxedClass* object_cls, *type_cls, *bool_cls, *int_cls, *long_cls, *float_cls, *str_cls, *function_cls,
     *none_cls, *instancemethod_cls, *list_cls, *slice_cls, *module_cls, *dict_cls, *tuple_cls, *file_cls, *xrange_cls,
-    *member_cls, *method_cls, *closure_cls, *generator_cls;
+    *member_cls, *method_cls, *closure_cls, *generator_cls, *complex_cls;
 }
 extern "C" { extern Box* None, *NotImplemented, *True, *False; }
 extern "C" {
@@ -194,6 +196,14 @@ public:
     BoxedFloat(double d) __attribute__((visibility("default"))) : Box(float_cls), d(d) {}
 };
 
+class BoxedComplex : public Box {
+public:
+    double real;
+    double imag;
+
+    BoxedComplex(double r, double i) __attribute__((visibility("default"))) : Box(complex_cls), real(r), imag(i) {}
+};
+
 class BoxedBool : public Box {
 public:
     bool b;

test/tests/complex.py

+# TODO repr is wrong, so for now, only printing complex numbers whose real
+# and imaginary parts are non-integers
+
+print 0.5j + 1.5
+print 0.5j + 1.5
+print 0.5j + 1.5
+
+print 0.5j - 1.5
+print 0.5j - 1.5
+print 0.5j - 1.5
+
+print 0.5j * 1.5
+print 0.5j * 1.5
+print 0.5j * 1.5