Merge pull request #719 from Daetalus/pow

Pow

src/runtime/float.cpp

@@ -1502,6 +1502,8 @@ void setupFloat() {
     float_cls->freeze();
 
     floatFormatInit();
+
+    float_cls->tp_as_number->nb_power = float_pow;
 }
 
 void teardownFloat() {

src/runtime/int.cpp

@@ -334,33 +334,17 @@ extern "C" i64 mod_i64_i64(i64 lhs, i64 rhs) {
     return lhs % rhs;
 }
 
-extern "C" Box* pow_i64_i64(i64 lhs, i64 rhs) {
+extern "C" Box* pow_i64_i64(i64 lhs, i64 rhs, Box* mod) {
     i64 orig_rhs = rhs;
     i64 rtn = 1, curpow = lhs;
 
     if (rhs < 0)
+        // already checked, rhs is a integer,
+        // and mod will be None in this case.
         return boxFloat(pow_float_float(lhs, rhs));
 
-    if (rhs == 0)
-        return boxInt(1);
-
-    assert(rhs > 0);
-    while (true) {
-        if (rhs & 1) {
-            // TODO: could potentially avoid restarting the entire computation on overflow?
-            if (__builtin_smull_overflow(rtn, curpow, &rtn))
-                return longPow(boxLong(lhs), boxLong(orig_rhs));
-        }
-
-        rhs >>= 1;
-
-        if (!rhs)
-            break;
-
-        if (__builtin_smull_overflow(curpow, curpow, &curpow))
-            return longPow(boxLong(lhs), boxLong(orig_rhs));
-    }
-    return boxInt(rtn);
+    // let longPow do the checks.
+    return longPow(boxLong(lhs), boxLong(rhs), mod);
 }
 
 extern "C" Box* mul_i64_i64(i64 lhs, i64 rhs) {
@@ -663,36 +647,62 @@ extern "C" Box* intMul(BoxedInt* lhs, Box* rhs) {
     }
 }
 
-extern "C" Box* intPowInt(BoxedInt* lhs, BoxedInt* rhs) {
+static void _addFuncPow(const char* name, ConcreteCompilerType* rtn_type, void* float_func, void* int_func) {
+    std::vector<ConcreteCompilerType*> v_ifu{ BOXED_INT, BOXED_FLOAT, UNKNOWN };
+    std::vector<ConcreteCompilerType*> v_uuu{ UNKNOWN, UNKNOWN, UNKNOWN };
+
+    CLFunction* cl = createRTFunction(3, 1, false, false);
+    addRTFunction(cl, float_func, UNKNOWN, v_ifu);
+    addRTFunction(cl, int_func, UNKNOWN, v_uuu);
+    int_cls->giveAttr(name, new BoxedFunction(cl, { None }));
+}
+
+extern "C" Box* intPowLong(BoxedInt* lhs, BoxedLong* rhs, Box* mod) {
     assert(isSubclass(lhs->cls, int_cls));
-    assert(isSubclass(rhs->cls, int_cls));
-    BoxedInt* rhs_int = static_cast<BoxedInt*>(rhs);
-    return pow_i64_i64(lhs->n, rhs_int->n);
+    assert(isSubclass(rhs->cls, long_cls));
+    BoxedLong* lhs_long = boxLong(lhs->n);
+    return longPow(lhs_long, rhs, mod);
 }
 
-extern "C" Box* intPowFloat(BoxedInt* lhs, BoxedFloat* rhs) {
+extern "C" Box* intPowFloat(BoxedInt* lhs, BoxedFloat* rhs, Box* mod) {
     assert(isSubclass(lhs->cls, int_cls));
     assert(rhs->cls == float_cls);
-    return boxFloat(pow(lhs->n, rhs->d));
+
+    if (mod != None) {
+        raiseExcHelper(TypeError, "pow() 3rd argument not allowed unless all arguments are integers");
+    }
+    return boxFloat(pow_float_float(lhs->n, rhs->d));
 }
 
 extern "C" Box* intPow(BoxedInt* lhs, Box* rhs, Box* mod) {
     if (!isSubclass(lhs->cls, int_cls))
         raiseExcHelper(TypeError, "descriptor '__pow__' requires a 'int' object but received a '%s'", getTypeName(lhs));
 
-    if (isSubclass(rhs->cls, int_cls)) {
+    if (isSubclass(rhs->cls, long_cls))
+        return intPowLong(lhs, static_cast<BoxedLong*>(rhs), mod);
+    else if (isSubclass(rhs->cls, float_cls))
+        return intPowFloat(lhs, static_cast<BoxedFloat*>(rhs), mod);
+    else if (!isSubclass(rhs->cls, int_cls))
+        return NotImplemented;
+
     BoxedInt* rhs_int = static_cast<BoxedInt*>(rhs);
-        Box* rtn = intPowInt(lhs, rhs_int);
-        if (mod == None)
-            return rtn;
-        return binop(rtn, mod, AST_TYPE::Mod);
-    } else if (rhs->cls == float_cls) {
-        RELEASE_ASSERT(mod == None, "");
-        BoxedFloat* rhs_float = static_cast<BoxedFloat*>(rhs);
-        return intPowFloat(lhs, rhs_float);
-    } else {
+    BoxedInt* mod_int = static_cast<BoxedInt*>(mod);
+
+    if (mod != None) {
+        if (rhs_int->n < 0)
+            raiseExcHelper(TypeError, "pow() 2nd argument "
+                                      "cannot be negative when 3rd argument specified");
+        if (!isSubclass(mod->cls, int_cls)) {
             return NotImplemented;
+        } else if (mod_int->n == 0) {
+            raiseExcHelper(ValueError, "pow() 3rd argument cannot be 0");
+        }
     }
+
+    Box* rtn = pow_i64_i64(lhs->n, rhs_int->n, mod);
+    if (isSubclass(rtn->cls, long_cls))
+        return longInt(rtn);
+    return rtn;
 }
 
 extern "C" Box* intRShiftInt(BoxedInt* lhs, BoxedInt* rhs) {
@@ -1096,9 +1106,7 @@ void setupInt() {
     _addFuncIntFloatUnknown("__truediv__", (void*)intTruedivInt, (void*)intTruedivFloat, (void*)intTruediv);
     _addFuncIntFloatUnknown("__mul__", (void*)intMulInt, (void*)intMulFloat, (void*)intMul);
     _addFuncIntUnknown("__mod__", BOXED_INT, (void*)intModInt, (void*)intMod);
-    int_cls->giveAttr("__pow__",
-                      new BoxedFunction(boxRTFunction((void*)intPow, UNKNOWN, 3, 1, false, false), { None }));
-
+    _addFuncPow("__pow__", BOXED_INT, (void*)intPowFloat, (void*)intPow);
     // Note: CPython implements int comparisons using tp_compare
     int_cls->tp_richcompare = int_richcompare;
 

src/runtime/int.h

@@ -32,7 +32,7 @@ extern "C" i64 mod_i64_i64(i64 lhs, i64 rhs);
 
 extern "C" Box* add_i64_i64(i64 lhs, i64 rhs);
 extern "C" Box* sub_i64_i64(i64 lhs, i64 rhs);
-extern "C" Box* pow_i64_i64(i64 lhs, i64 rhs);
+extern "C" Box* pow_i64_i64(i64 lhs, i64 rhs, Box* mod = None);
 extern "C" Box* mul_i64_i64(i64 lhs, i64 rhs);
 extern "C" i1 eq_i64_i64(i64 lhs, i64 rhs);
 extern "C" i1 ne_i64_i64(i64 lhs, i64 rhs);

src/runtime/long.cpp

@@ -27,6 +27,7 @@
 #include "core/stats.h"
 #include "core/types.h"
 #include "gc/collector.h"
+#include "runtime/float.h"
 #include "runtime/inline/boxing.h"
 #include "runtime/objmodel.h"
 #include "runtime/types.h"
@@ -1191,49 +1192,89 @@ Box* longRTrueDiv(BoxedLong* v1, Box* _v2) {
     return boxFloat(lhs / (double)rhs);
 }
 
-Box* longPow(BoxedLong* v1, Box* _v2, Box* _v3) {
-    if (!isSubclass(v1->cls, long_cls))
-        raiseExcHelper(TypeError, "descriptor '__pow__' requires a 'long' object but received a '%s'", getTypeName(v1));
-
-    BoxedLong* mod = nullptr;
-    if (_v3 != None) {
-        if (isSubclass(_v3->cls, int_cls))
-            mod = boxLong(((BoxedInt*)_v3)->n);
-        else {
-            RELEASE_ASSERT(_v3->cls == long_cls, "");
-            mod = (BoxedLong*)_v3;
-        }
-        RELEASE_ASSERT(mpz_sgn(mod->n) >= 0, "");
-    }
+static void _addFuncPow(const char* name, ConcreteCompilerType* rtn_type, void* float_func, void* long_func) {
+    std::vector<ConcreteCompilerType*> v_lfu{ UNKNOWN, BOXED_FLOAT, UNKNOWN };
+    std::vector<ConcreteCompilerType*> v_uuu{ UNKNOWN, UNKNOWN, UNKNOWN };
 
-    if (isSubclass(_v2->cls, long_cls)) {
-        BoxedLong* v2 = static_cast<BoxedLong*>(_v2);
-        BoxedLong* r = new BoxedLong();
-        mpz_init(r->n);
+    CLFunction* cl = createRTFunction(3, 1, false, false);
+    addRTFunction(cl, float_func, UNKNOWN, v_lfu);
+    addRTFunction(cl, long_func, UNKNOWN, v_uuu);
+    long_cls->giveAttr(name, new BoxedFunction(cl, { None }));
+}
+
+extern "C" Box* longPowFloat(BoxedLong* lhs, BoxedFloat* rhs) {
+    assert(isSubclass(lhs->cls, long_cls));
+    assert(isSubclass(rhs->cls, float_cls));
+    double lhs_float = static_cast<BoxedFloat*>(longFloat(lhs))->d;
+    return boxFloat(pow_float_float(lhs_float, rhs->d));
+}
 
-        RELEASE_ASSERT(mpz_sgn(v2->n) >= 0, "");
+Box* longPow(BoxedLong* lhs, Box* rhs, Box* mod) {
+    if (!isSubclass(lhs->cls, long_cls))
+        raiseExcHelper(TypeError, "descriptor '__pow__' requires a 'long' object but received a '%s'",
+                       getTypeName(lhs));
 
-        if (mod) {
-            mpz_powm(r->n, v1->n, v2->n, mod->n);
+    BoxedLong* mod_long = nullptr;
+    if (mod != None) {
+        if (isSubclass(mod->cls, long_cls)) {
+            mod_long = static_cast<BoxedLong*>(mod);
+        } else if (isSubclass(mod->cls, int_cls)) {
+            mod_long = boxLong(static_cast<BoxedInt*>(mod)->n);
         } else {
-            RELEASE_ASSERT(mpz_fits_ulong_p(v2->n), "");
-            uint64_t n2 = mpz_get_ui(v2->n);
-            mpz_pow_ui(r->n, v1->n, n2);
+            return NotImplemented;
         }
-        return r;
-    } else if (isSubclass(_v2->cls, int_cls)) {
-        BoxedInt* v2 = static_cast<BoxedInt*>(_v2);
-        RELEASE_ASSERT(v2->n >= 0, "");
+    }
+
+    BoxedLong* rhs_long = nullptr;
+    if (isSubclass(rhs->cls, long_cls)) {
+        rhs_long = static_cast<BoxedLong*>(rhs);
+    } else if (isSubclass(rhs->cls, int_cls)) {
+        rhs_long = boxLong(static_cast<BoxedInt*>(rhs)->n);
+    } else {
+        return NotImplemented;
+    }
+
+    if (mod != None) {
+        if (mpz_sgn(rhs_long->n) < 0)
+            raiseExcHelper(TypeError, "pow() 2nd argument "
+                                      "cannot be negative when 3rd argument specified");
+        else if (mpz_sgn(mod_long->n) == 0)
+            raiseExcHelper(ValueError, "pow() 3rd argument cannot be 0");
+    }
+
     BoxedLong* r = new BoxedLong();
     mpz_init(r->n);
-        if (mod)
-            mpz_powm_ui(r->n, v1->n, v2->n, mod->n);
-        else
-            mpz_pow_ui(r->n, v1->n, v2->n);
+
+    if (mpz_sgn(rhs_long->n) == -1) {
+        BoxedFloat* rhs_float = static_cast<BoxedFloat*>(longFloat(rhs_long));
+        BoxedFloat* lhs_float = static_cast<BoxedFloat*>(longFloat(lhs));
+        return boxFloat(pow_float_float(lhs_float->d, rhs_float->d));
+    }
+
+    if (mod != None) {
+        mpz_powm(r->n, lhs->n, rhs_long->n, mod_long->n);
+        if (mpz_sgn(r->n) == 0)
             return r;
+        if (mpz_sgn(mod_long->n) < 0)
+            return longAdd(r, mod_long);
     } else {
-        return NotImplemented;
+        if (mpz_fits_ulong_p(rhs_long->n)) {
+            uint64_t n2 = mpz_get_ui(rhs_long->n);
+            mpz_pow_ui(r->n, lhs->n, n2);
+        } else {
+            if (mpz_cmp_si(lhs->n, 1l) == 0) {
+                mpz_set_ui(r->n, 1l);
+            } else if (mpz_sgn(lhs->n) == 0) {
+                mpz_set_ui(r->n, 0l);
+            } else if (mpz_cmp_si(lhs->n, -1l) == 0) {
+                long rl = mpz_even_p(rhs_long->n) ? 1l : -1l;
+                mpz_set_si(r->n, rl);
+            } else {
+                raiseExcHelper(OverflowError, "the result is too large to convert to long");
+            }
         }
+    }
+    return r;
 }
 
 extern "C" Box* longInvert(BoxedLong* v) {
@@ -1343,7 +1384,8 @@ static PyObject* long_pow(PyObject* v, PyObject* w, PyObject* x) noexcept {
         CONVERT_BINOP(v, w, &a, &b);
         return longPow((BoxedLong*)a, (BoxedLong*)b, x);
     } catch (ExcInfo e) {
-        abort();
+        setCAPIException(e);
+        return NULL;
     }
 }
 
@@ -1369,6 +1411,7 @@ static Box* long1(Box* b, void*) {
 void setupLong() {
     mp_set_memory_functions(customised_allocation, customised_realloc, customised_free);
 
+    _addFuncPow("__pow__", UNKNOWN, (void*)longPowFloat, (void*)longPow);
     long_cls->giveAttr(
         "__new__", new BoxedFunction(boxRTFunction((void*)longNew, UNKNOWN, 3, 2, false, false), { boxInt(0), NULL }));
 
@@ -1391,10 +1434,6 @@ void setupLong() {
 
     long_cls->giveAttr("__add__", new BoxedFunction(boxRTFunction((void*)longAdd, UNKNOWN, 2)));
     long_cls->giveAttr("__radd__", long_cls->getattr(internStringMortal("__add__")));
-
-    long_cls->giveAttr("__pow__",
-                       new BoxedFunction(boxRTFunction((void*)longPow, UNKNOWN, 3, 1, false, false), { None }));
-
     long_cls->giveAttr("__and__", new BoxedFunction(boxRTFunction((void*)longAnd, UNKNOWN, 2)));
     long_cls->giveAttr("__rand__", long_cls->getattr(internStringMortal("__and__")));
     long_cls->giveAttr("__or__", new BoxedFunction(boxRTFunction((void*)longOr, UNKNOWN, 2)));

src/runtime/long.h

@@ -55,6 +55,7 @@ Box* longRshift(BoxedLong* lhs, Box* rhs);
 Box* longHex(BoxedLong* v);
 Box* longOct(BoxedLong* v);
 Box* longStr(BoxedLong* v);
+Box* longInt(Box* v);
 
 bool longNonzeroUnboxed(BoxedLong* n);
 }

test/tests/float.py

@@ -95,3 +95,13 @@ print -(0.0)
 print -(-0.0)
 
 print repr((1e100).__trunc__())
+
+all_args = [None, 1, 1L, -1, -1L, 2.0, 0.5, 0, "",
+            0.0, -0.0, -0.5, float('nan'), float('inf')]
+for lhs in all_args:
+    for rhs in all_args:
+        for mod in all_args:
+            try:
+                print pow(lhs, rhs, mod)
+            except Exception as e:
+                print type(e), e

test/tests/long.py

@@ -57,6 +57,17 @@ print ~(-10L)
 print -(1L)
 print 1L**2L
 print 1L**2
+print 0 ** (1 << 100)
+print pow(1 << 30, 1 << 30, 127)
+print pow(1L << 30, 1L << 30, 127)
+print pow(1 << 100, 1 << 100, 1000)
+print pow(-1, (1<<100))
+print pow(-1, (1<<100) + 1)
+print pow(0, (1<<100))
+print pow(1, (1<<100))
+print pow(5, 3, -7L)
+print pow(-5,  3, 7L)
+print pow(-5,  3, -7L)
 print (11L).__pow__(32, 50L)
 print (11L).__index__()