py: More robust int conversion and overflow checking.

1 files changed, 28 insertions, 3 deletions

diff --git a/py/mpz.c b/py/mpz.c
index 79d200d93..4a8941e29 100644
--- a/py/mpz.c
+++ b/py/mpz.c
@@ -1139,6 +1139,7 @@ mpz_t *mpz_mod(const mpz_t *lhs, const mpz_t *rhs) {
 }
 #endif
 
+// TODO check that this correctly handles overflow in all cases
 machine_int_t mpz_as_int(const mpz_t *i) {
     machine_int_t val = 0;
     mpz_dig_t *d = i->dig + i->len;
@@ -1147,11 +1148,13 @@ machine_int_t mpz_as_int(const mpz_t *i) {
         machine_int_t oldval = val;
         val = (val << DIG_SIZE) | *d;
         if (val < oldval) {
-            // TODO need better handling of conversion overflow
+            // overflow, return +/- "infinity"
             if (i->neg == 0) {
-                return 0x7fffffff;
+                // +infinity
+                return ~WORD_MSBIT_HIGH;
             } else {
-                return 0x80000000;
+                // -infinity
+                return WORD_MSBIT_HIGH;
             }
         }
     }
@@ -1163,6 +1166,28 @@ machine_int_t mpz_as_int(const mpz_t *i) {
     return val;
 }
 
+// TODO check that this correctly handles overflow in all cases
+bool mpz_as_int_checked(const mpz_t *i, machine_int_t *value) {
+    machine_int_t val = 0;
+    mpz_dig_t *d = i->dig + i->len;
+
+    while (--d >= i->dig) {
+        machine_int_t oldval = val;
+        val = (val << DIG_SIZE) | *d;
+        if (val < oldval) {
+            // overflow
+            return false;
+        }
+    }
+
+    if (i->neg != 0) {
+        val = -val;
+    }
+
+    *value = val;
+    return true;
+}
+
 #if MICROPY_ENABLE_FLOAT
 mp_float_t mpz_as_float(const mpz_t *i) {
     mp_float_t val = 0;