py: Fix up number operations and coercion.

6 files changed, 131 insertions, 100 deletions

diff --git a/examples/mandel.py b/examples/mandel.py
new file mode 100644
index 0000000000..b13b7d87f8
--- /dev/null
+++ b/examples/mandel.py
@@ -0,0 +1,14 @@
+@micropython.native
+def in_set(c):
+    z = 0
+    for i in range(40):
+        z = z*z + c
+        if abs(z) > 60:
+            return False
+    return True
+
+for v in range(31):
+    line = []
+    for u in range(91):
+        line.append('*' if in_set((u / 30 - 2) + (v / 15 - 1) * 1j) else ' ')
+    print(''.join(line))
diff --git a/py/asmx64.c b/py/asmx64.c
index ed9ca80f5c..226d4efee8 100644
--- a/py/asmx64.c
+++ b/py/asmx64.c
@@ -155,7 +155,7 @@ void asm_x64_end_pass(asm_x64_t *as) {
         //as->code_base = m_new(byte, as->code_size); need to allocale executable memory
         uint actual_alloc;
         as->code_base = alloc_mem(as->code_size, &actual_alloc, true);
-        printf("code_size: %u\n", as->code_size);
+        //printf("code_size: %u\n", as->code_size);
     }
 
     /*
diff --git a/py/obj.h b/py/obj.h
index 88a611ba7b..1ba7427cbb 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -200,10 +200,12 @@ qstr mp_obj_str_get(mp_obj_t self_in);
 // float
 extern const mp_obj_type_t float_type;
 mp_float_t mp_obj_float_get(mp_obj_t self_in);
+mp_obj_t mp_obj_float_binary_op(int op, mp_float_t lhs_val, mp_obj_t rhs);
 
 // complex
 extern const mp_obj_type_t complex_type;
 void mp_obj_complex_get(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag);
+mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in);
 #endif
 
 // tuple
diff --git a/py/objcomplex.c b/py/objcomplex.c
index fc32f96674..46f43b54b5 100644
--- a/py/objcomplex.c
+++ b/py/objcomplex.c
@@ -48,14 +48,14 @@ static mp_obj_t complex_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *a
         {
             mp_float_t real, imag;
             if (MP_OBJ_IS_TYPE(args[1], &complex_type)) {
-                mp_obj_get_complex(args[1], &real, &imag);
+                mp_obj_complex_get(args[1], &real, &imag);
             } else {
                 real = mp_obj_get_float(args[1]);
                 imag = 0;
             }
             if (MP_OBJ_IS_TYPE(args[0], &complex_type)) {
                 mp_float_t real2, imag2;
-                mp_obj_get_complex(args[0], &real2, &imag2);
+                mp_obj_complex_get(args[0], &real2, &imag2);
                 real -= imag2;
                 imag += real2;
             } else {
@@ -80,9 +80,41 @@ static mp_obj_t complex_unary_op(int op, mp_obj_t o_in) {
 }
 
 static mp_obj_t complex_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
-    mp_float_t lhs_real, lhs_imag, rhs_real, rhs_imag;
-    mp_obj_complex_get(lhs_in, &lhs_real, &lhs_imag);
-    mp_obj_complex_get(rhs_in, &rhs_real, &rhs_imag);
+    mp_obj_complex_t *lhs = lhs_in;
+    return mp_obj_complex_binary_op(op, lhs->real, lhs->imag, rhs_in);
+}
+
+const mp_obj_type_t complex_type = {
+    { &mp_const_type },
+    "complex",
+    complex_print, // print
+    complex_make_new, // make_new
+    NULL, // call_n
+    complex_unary_op, // unary_op
+    complex_binary_op, // binary_op
+    NULL, // getiter
+    NULL, // iternext
+    .methods = { { NULL, NULL }, },
+};
+
+mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag) {
+    mp_obj_complex_t *o = m_new_obj(mp_obj_complex_t);
+    o->base.type = &complex_type;
+    o->real = real;
+    o->imag = imag;
+    return o;
+}
+
+void mp_obj_complex_get(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag) {
+    assert(MP_OBJ_IS_TYPE(self_in, &complex_type));
+    mp_obj_complex_t *self = self_in;
+    *real = self->real;
+    *imag = self->imag;
+}
+
+mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in) {
+    mp_float_t rhs_real, rhs_imag;
+    mp_obj_get_complex(rhs_in, &rhs_real, &rhs_imag); // can be any type, this function will convert to float (if possible)
     switch (op) {
         case RT_BINARY_OP_ADD:
         case RT_BINARY_OP_INPLACE_ADD:
@@ -115,32 +147,4 @@ static mp_obj_t complex_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
     return mp_obj_new_complex(lhs_real, lhs_imag);
 }
 
-const mp_obj_type_t complex_type = {
-    { &mp_const_type },
-    "complex",
-    complex_print, // print
-    complex_make_new, // make_new
-    NULL, // call_n
-    complex_unary_op, // unary_op
-    complex_binary_op, // binary_op
-    NULL, // getiter
-    NULL, // iternext
-    .methods = { { NULL, NULL }, },
-};
-
-mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag) {
-    mp_obj_complex_t *o = m_new_obj(mp_obj_complex_t);
-    o->base.type = &complex_type;
-    o->real = real;
-    o->imag = imag;
-    return o;
-}
-
-void mp_obj_complex_get(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag) {
-    assert(MP_OBJ_IS_TYPE(self_in, &complex_type));
-    mp_obj_complex_t *self = self_in;
-    *real = self->real;
-    *imag = self->imag;
-}
-
 #endif
diff --git a/py/objfloat.c b/py/objfloat.c
index 0250172ad3..336ae597fc 100644
--- a/py/objfloat.c
+++ b/py/objfloat.c
@@ -53,27 +53,12 @@ static mp_obj_t float_unary_op(int op, mp_obj_t o_in) {
 }
 
 static mp_obj_t float_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
+    mp_obj_float_t *lhs = lhs_in;
     if (MP_OBJ_IS_TYPE(rhs_in, &complex_type)) {
-        return complex_type.binary_op(op, lhs_in, rhs_in);
+        return mp_obj_complex_binary_op(op, lhs->value, 0, rhs_in);
+    } else {
+        return mp_obj_float_binary_op(op, lhs->value, rhs_in);
     }
-    mp_float_t lhs_val = mp_obj_get_float(lhs_in);
-    mp_float_t rhs_val = mp_obj_get_float(rhs_in);
-    switch (op) {
-        case RT_BINARY_OP_ADD:
-        case RT_BINARY_OP_INPLACE_ADD: lhs_val += rhs_val; break;
-        case RT_BINARY_OP_SUBTRACT:
-        case RT_BINARY_OP_INPLACE_SUBTRACT: lhs_val -= rhs_val; break;
-        case RT_BINARY_OP_MULTIPLY:
-        case RT_BINARY_OP_INPLACE_MULTIPLY: lhs_val *= rhs_val; break;
-        /* TODO floor(?) the value
-        case RT_BINARY_OP_FLOOR_DIVIDE:
-        case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: val = lhs_val / rhs_val; break;
-        */
-        case RT_BINARY_OP_TRUE_DIVIDE:
-        case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: lhs_val /= rhs_val; break;
-        return NULL; // op not supported
-    }
-    return mp_obj_new_float(lhs_val);
 }
 
 const mp_obj_type_t float_type = {
@@ -99,4 +84,24 @@ mp_float_t mp_obj_float_get(mp_obj_t self_in) {
     return self->value;
 }
 
+mp_obj_t mp_obj_float_binary_op(int op, mp_float_t lhs_val, mp_obj_t rhs_in) {
+    mp_float_t rhs_val = mp_obj_get_float(rhs_in); // can be any type, this function will convert to float (if possible)
+    switch (op) {
+        case RT_BINARY_OP_ADD:
+        case RT_BINARY_OP_INPLACE_ADD: lhs_val += rhs_val; break;
+        case RT_BINARY_OP_SUBTRACT:
+        case RT_BINARY_OP_INPLACE_SUBTRACT: lhs_val -= rhs_val; break;
+        case RT_BINARY_OP_MULTIPLY:
+        case RT_BINARY_OP_INPLACE_MULTIPLY: lhs_val *= rhs_val; break;
+        /* TODO floor(?) the value
+        case RT_BINARY_OP_FLOOR_DIVIDE:
+        case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: val = lhs_val / rhs_val; break;
+        */
+        case RT_BINARY_OP_TRUE_DIVIDE:
+        case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: lhs_val /= rhs_val; break;
+        return NULL; // op not supported
+    }
+    return mp_obj_new_float(lhs_val);
+}
+
 #endif
diff --git a/py/runtime.c b/py/runtime.c
index 197c08b55a..6bc71abff7 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -452,57 +452,63 @@ mp_obj_t rt_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
     //   then fail
     // note that list does not implement + or +=, so that inplace_concat is reached first for +=
 
-    if (MP_OBJ_IS_SMALL_INT(lhs) && MP_OBJ_IS_SMALL_INT(rhs)) {
+    if (MP_OBJ_IS_SMALL_INT(lhs)) {
         mp_small_int_t lhs_val = MP_OBJ_SMALL_INT_VALUE(lhs);
-        mp_small_int_t rhs_val = MP_OBJ_SMALL_INT_VALUE(rhs);
-        switch (op) {
-            case RT_BINARY_OP_OR:
-            case RT_BINARY_OP_INPLACE_OR: lhs_val |= rhs_val; break;
-            case RT_BINARY_OP_XOR:
-            case RT_BINARY_OP_INPLACE_XOR: lhs_val ^= rhs_val; break;
-            case RT_BINARY_OP_AND:
-            case RT_BINARY_OP_INPLACE_AND: lhs_val &= rhs_val; break;
-            case RT_BINARY_OP_LSHIFT:
-            case RT_BINARY_OP_INPLACE_LSHIFT: lhs_val <<= rhs_val; break;
-            case RT_BINARY_OP_RSHIFT:
-            case RT_BINARY_OP_INPLACE_RSHIFT: lhs_val >>= rhs_val; break;
-            case RT_BINARY_OP_ADD:
-            case RT_BINARY_OP_INPLACE_ADD: lhs_val += rhs_val; break;
-            case RT_BINARY_OP_SUBTRACT:
-            case RT_BINARY_OP_INPLACE_SUBTRACT: lhs_val -= rhs_val; break;
-            case RT_BINARY_OP_MULTIPLY:
-            case RT_BINARY_OP_INPLACE_MULTIPLY: lhs_val *= rhs_val; break;
-            case RT_BINARY_OP_FLOOR_DIVIDE:
-            case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: lhs_val /= rhs_val; break;
-#if MICROPY_ENABLE_FLOAT
-            case RT_BINARY_OP_TRUE_DIVIDE:
-            case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: return mp_obj_new_float((mp_float_t)lhs_val / (mp_float_t)rhs_val);
-#endif
-
-            // TODO implement modulo as specified by Python
-            case RT_BINARY_OP_MODULO:
-            case RT_BINARY_OP_INPLACE_MODULO: lhs_val %= rhs_val; break;
-
-            // TODO check for negative power, and overflow
-            case RT_BINARY_OP_POWER:
-            case RT_BINARY_OP_INPLACE_POWER:
-            {
-                int ans = 1;
-                while (rhs_val > 0) {
-                    if (rhs_val & 1) {
-                        ans *= lhs_val;
+        if (MP_OBJ_IS_SMALL_INT(rhs)) {
+            mp_small_int_t rhs_val = MP_OBJ_SMALL_INT_VALUE(rhs);
+            switch (op) {
+                case RT_BINARY_OP_OR:
+                case RT_BINARY_OP_INPLACE_OR: lhs_val |= rhs_val; break;
+                case RT_BINARY_OP_XOR:
+                case RT_BINARY_OP_INPLACE_XOR: lhs_val ^= rhs_val; break;
+                case RT_BINARY_OP_AND:
+                case RT_BINARY_OP_INPLACE_AND: lhs_val &= rhs_val; break;
+                case RT_BINARY_OP_LSHIFT:
+                case RT_BINARY_OP_INPLACE_LSHIFT: lhs_val <<= rhs_val; break;
+                case RT_BINARY_OP_RSHIFT:
+                case RT_BINARY_OP_INPLACE_RSHIFT: lhs_val >>= rhs_val; break;
+                case RT_BINARY_OP_ADD:
+                case RT_BINARY_OP_INPLACE_ADD: lhs_val += rhs_val; break;
+                case RT_BINARY_OP_SUBTRACT:
+                case RT_BINARY_OP_INPLACE_SUBTRACT: lhs_val -= rhs_val; break;
+                case RT_BINARY_OP_MULTIPLY:
+                case RT_BINARY_OP_INPLACE_MULTIPLY: lhs_val *= rhs_val; break;
+                case RT_BINARY_OP_FLOOR_DIVIDE:
+                case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: lhs_val /= rhs_val; break;
+    #if MICROPY_ENABLE_FLOAT
+                case RT_BINARY_OP_TRUE_DIVIDE:
+                case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: return mp_obj_new_float((mp_float_t)lhs_val / (mp_float_t)rhs_val);
+    #endif
+
+                // TODO implement modulo as specified by Python
+                case RT_BINARY_OP_MODULO:
+                case RT_BINARY_OP_INPLACE_MODULO: lhs_val %= rhs_val; break;
+
+                // TODO check for negative power, and overflow
+                case RT_BINARY_OP_POWER:
+                case RT_BINARY_OP_INPLACE_POWER:
+                {
+                    int ans = 1;
+                    while (rhs_val > 0) {
+                        if (rhs_val & 1) {
+                            ans *= lhs_val;
+                        }
+                        lhs_val *= lhs_val;
+                        rhs_val /= 2;
                     }
-                    lhs_val *= lhs_val;
-                    rhs_val /= 2;
+                    lhs_val = ans;
+                    break;
                 }
-                lhs_val = ans;
-                break;
-            }
 
-            default: assert(0);
-        }
-        if (fit_small_int(lhs_val)) {
-            return MP_OBJ_NEW_SMALL_INT(lhs_val);
+                default: assert(0);
+            }
+            if (fit_small_int(lhs_val)) {
+                return MP_OBJ_NEW_SMALL_INT(lhs_val);
+            }
+        } else if (MP_OBJ_IS_TYPE(rhs, &float_type)) {
+            return mp_obj_float_binary_op(op, lhs_val, rhs);
+        } else if (MP_OBJ_IS_TYPE(rhs, &complex_type)) {
+            return mp_obj_complex_binary_op(op, lhs_val, 0, rhs);
         }
     } else if (MP_OBJ_IS_OBJ(lhs)) {
         mp_obj_base_t *o = lhs;