py: Fix float/complex binop returning NULL; implement complex power.

4 files changed, 45 insertions, 5 deletions

diff --git a/py/obj.h b/py/obj.h
index b944d2fcb5..af6ef83189 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -429,11 +429,11 @@ typedef struct _mp_obj_float_t {
     mp_float_t value;
 } mp_obj_float_t;
 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);
+mp_obj_t mp_obj_float_binary_op(int op, mp_float_t lhs_val, mp_obj_t rhs); // can return MP_OBJ_NULL
 
 // complex
 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);
+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); // can return MP_OBJ_NULL
 #endif
 
 // tuple
diff --git a/py/objcomplex.c b/py/objcomplex.c
index 66f971da0e..23c67eb6e1 100644
--- a/py/objcomplex.c
+++ b/py/objcomplex.c
@@ -11,6 +11,8 @@
 
 #if MICROPY_ENABLE_FLOAT
 
+#include <math.h>
+
 #if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
 #include "formatfloat.h"
 #endif
@@ -176,6 +178,33 @@ mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_im
             }
             break;
 
+        case MP_BINARY_OP_POWER:
+        case MP_BINARY_OP_INPLACE_POWER: {
+            // z1**z2 = exp(z2*ln(z1))
+            //        = exp(z2*(ln(|z1|)+i*arg(z1)))
+            //        = exp( (x2*ln1 - y2*arg1) + i*(y2*ln1 + x2*arg1) )
+            //        = exp(x3 + i*y3)
+            //        = exp(x3)*(cos(y3) + i*sin(y3))
+            mp_float_t abs1 = MICROPY_FLOAT_C_FUN(sqrt)(lhs_real*lhs_real + lhs_imag*lhs_imag);
+            if (abs1 == 0) {
+                if (rhs_imag == 0) {
+                    lhs_real = 1;
+                    rhs_real = 0;
+                } else {
+                    nlr_raise(mp_obj_new_exception_msg(&mp_type_ZeroDivisionError, "0.0 to a complex power"));
+                }
+            } else {
+                mp_float_t ln1 = MICROPY_FLOAT_C_FUN(log)(abs1);
+                mp_float_t arg1 = MICROPY_FLOAT_C_FUN(atan2)(lhs_imag, lhs_real);
+                mp_float_t x3 = rhs_real * ln1 - rhs_imag * arg1;
+                mp_float_t y3 = rhs_imag * ln1 + rhs_real * arg1;
+                mp_float_t exp_x3 = MICROPY_FLOAT_C_FUN(exp)(x3);
+                lhs_real = exp_x3 * MICROPY_FLOAT_C_FUN(cos)(y3);
+                lhs_imag = exp_x3 * MICROPY_FLOAT_C_FUN(sin)(y3);
+            }
+            break;
+        }
+
         default:
             return MP_OBJ_NULL; // op not supported
     }
diff --git a/py/objfloat.c b/py/objfloat.c
index 8ccba1024c..1f96ac18e5 100644
--- a/py/objfloat.c
+++ b/py/objfloat.c
@@ -136,7 +136,7 @@ check_zero_division:
         case MP_BINARY_OP_MORE_EQUAL: return MP_BOOL(lhs_val >= rhs_val);
 
         default:
-            return NULL; // op not supported
+            return MP_OBJ_NULL; // op not supported
     }
     return mp_obj_new_float(lhs_val);
 }
diff --git a/py/runtime.c b/py/runtime.c
index ef07e39bff..499905a0fa 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -386,9 +386,19 @@ mp_obj_t mp_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
             }
 #if MICROPY_ENABLE_FLOAT
         } else if (MP_OBJ_IS_TYPE(rhs, &mp_type_float)) {
-            return mp_obj_float_binary_op(op, lhs_val, rhs);
+            mp_obj_t res = mp_obj_float_binary_op(op, lhs_val, rhs);
+            if (res == MP_OBJ_NULL) {
+                goto unsupported_op;
+            } else {
+                return res;
+            }
         } else if (MP_OBJ_IS_TYPE(rhs, &mp_type_complex)) {
-            return mp_obj_complex_binary_op(op, lhs_val, 0, rhs);
+            mp_obj_t res = mp_obj_complex_binary_op(op, lhs_val, 0, rhs);
+            if (res == MP_OBJ_NULL) {
+                goto unsupported_op;
+            } else {
+                return res;
+            }
 #endif
         }
     }
@@ -438,6 +448,7 @@ generic_binary_op:
     // TODO implement dispatch for reverse binary ops
 
     // TODO specify in error message what the operator is
+unsupported_op:
     nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
         "unsupported operand types for binary operator: '%s', '%s'",
         mp_obj_get_type_str(lhs), mp_obj_get_type_str(rhs)));