From 8270e3853dc167d2d7946bb0de7a0f0bb2adde48 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Thu, 3 Apr 2014 11:00:54 +0000
Subject: [PATCH] py: More robust int conversion and overflow checking.

---
 py/mpz.c        | 31 ++++++++++++++++++++++++++++---
 py/mpz.h        |  1 +
 py/obj.c        | 18 ++++++++++++++++++
 py/obj.h        |  1 +
 py/objint_mpz.c | 14 ++++++++++++--
 py/objlist.c    |  7 +++----
 6 files changed, 63 insertions(+), 9 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;
diff --git a/py/mpz.h b/py/mpz.h
index 777889322..cbe60eb8d 100644
--- a/py/mpz.h
+++ b/py/mpz.h
@@ -71,6 +71,7 @@ mpz_t *mpz_div(const mpz_t *lhs, const mpz_t *rhs);
 mpz_t *mpz_mod(const mpz_t *lhs, const mpz_t *rhs);
 
 machine_int_t mpz_as_int(const mpz_t *z);
+bool mpz_as_int_checked(const mpz_t *z, machine_int_t *value);
 #if MICROPY_ENABLE_FLOAT
 mp_float_t mpz_as_float(const mpz_t *z);
 #endif
diff --git a/py/obj.c b/py/obj.c
index 95052d16d..bf3e1d50e 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -203,6 +203,24 @@ machine_int_t mp_obj_get_int(mp_obj_t arg) {
     }
 }
 
+// returns false if arg is not of integral type
+// returns true and sets *value if it is of integral type
+// can throw OverflowError if arg is of integral type, but doesn't fit in a machine_int_t
+bool mp_obj_get_int_maybe(mp_obj_t arg, machine_int_t *value) {
+    if (arg == mp_const_false) {
+        *value = 0;
+    } else if (arg == mp_const_true) {
+        *value = 1;
+    } else if (MP_OBJ_IS_SMALL_INT(arg)) {
+        *value = MP_OBJ_SMALL_INT_VALUE(arg);
+    } else if (MP_OBJ_IS_TYPE(arg, &mp_type_int)) {
+        *value = mp_obj_int_get_checked(arg);
+    } else {
+        return false;
+    }
+    return true;
+}
+
 #if MICROPY_ENABLE_FLOAT
 mp_float_t mp_obj_get_float(mp_obj_t arg) {
     if (arg == mp_const_false) {
diff --git a/py/obj.h b/py/obj.h
index 952187e46..f5ce873d6 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -359,6 +359,7 @@ bool mp_obj_equal(mp_obj_t o1, mp_obj_t o2);
 bool mp_obj_less(mp_obj_t o1, mp_obj_t o2);
 
 machine_int_t mp_obj_get_int(mp_obj_t arg);
+bool mp_obj_get_int_maybe(mp_obj_t arg, machine_int_t *value);
 #if MICROPY_ENABLE_FLOAT
 mp_float_t mp_obj_get_float(mp_obj_t self_in);
 void mp_obj_get_complex(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag);
diff --git a/py/objint_mpz.c b/py/objint_mpz.c
index 27f77d14b..25a069cf9 100644
--- a/py/objint_mpz.c
+++ b/py/objint_mpz.c
@@ -243,8 +243,18 @@ machine_int_t mp_obj_int_get(mp_obj_t self_in) {
 }
 
 machine_int_t mp_obj_int_get_checked(mp_obj_t self_in) {
-    // TODO: Check overflow
-    return mp_obj_int_get(self_in);
+    if (MP_OBJ_IS_SMALL_INT(self_in)) {
+        return MP_OBJ_SMALL_INT_VALUE(self_in);
+    } else {
+        mp_obj_int_t *self = self_in;
+        machine_int_t value;
+        if (mpz_as_int_checked(&self->mpz, &value)) {
+            return value;
+        } else {
+            // overflow
+            nlr_jump(mp_obj_new_exception_msg(&mp_type_OverflowError, "overflow converting long int to machine word"));
+        }
+    }
 }
 
 #if MICROPY_ENABLE_FLOAT
diff --git a/py/objlist.c b/py/objlist.c
index 244d4a596..2a5f2c883 100644
--- a/py/objlist.c
+++ b/py/objlist.c
@@ -122,12 +122,11 @@ STATIC mp_obj_t list_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
             list_extend(lhs, rhs);
             return o;
         }
-        case MP_BINARY_OP_MULTIPLY:
-        {
-            if (!MP_OBJ_IS_SMALL_INT(rhs)) {
+        case MP_BINARY_OP_MULTIPLY: {
+            machine_int_t n;
+            if (!mp_obj_get_int_maybe(rhs, &n)) {
                 return NULL;
             }
-            int n = MP_OBJ_SMALL_INT_VALUE(rhs);
             mp_obj_list_t *s = list_new(o->len * n);
             mp_seq_multiply(o->items, sizeof(*o->items), o->len, n, s->items);
             return s;