Merge pull request #2 from adafruit/master

Master

1 files changed, 98 insertions, 12 deletions

diff --git a/py/objint.c b/py/objint.c
index 2ac370d17..b78c9f25b 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -42,9 +42,9 @@
 #endif
 
 // This dispatcher function is expected to be independent of the implementation of long int
-STATIC mp_obj_t mp_obj_int_make_new(const mp_obj_type_t *type_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+STATIC mp_obj_t mp_obj_int_make_new(const mp_obj_type_t *type_in, size_t n_args, const mp_obj_t *args, mp_map_t *kw_args) {
     (void)type_in;
-    mp_arg_check_num(n_args, n_kw, 0, 2, false);
+    mp_arg_check_num(n_args, kw_args, 0, 2, false);
 
     switch (n_args) {
         case 0:
@@ -300,6 +300,76 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
     return b;
 }
 
+#if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
+
+void mp_obj_int_buffer_overflow_check(mp_obj_t self_in, size_t nbytes, bool is_signed)
+{
+    if (is_signed) {
+        // self must be < 2**(bits - 1)
+        mp_obj_t edge = mp_binary_op(MP_BINARY_OP_INPLACE_LSHIFT,
+                                     mp_obj_new_int(1),
+                                     mp_obj_new_int(nbytes * 8 - 1));
+
+        if (mp_binary_op(MP_BINARY_OP_LESS, self_in, edge) == mp_const_true) {
+            // and >= -2**(bits - 1)
+            edge = mp_unary_op(MP_UNARY_OP_NEGATIVE, edge);
+            if (mp_binary_op(MP_BINARY_OP_MORE_EQUAL, self_in, edge) == mp_const_true) {
+                return;
+            }
+        }
+    } else {
+        // self must be >= 0
+        if (mp_obj_int_sign(self_in) >= 0) {
+            // and < 2**(bits)
+            mp_obj_t edge = mp_binary_op(MP_BINARY_OP_INPLACE_LSHIFT,
+                                         mp_obj_new_int(1),
+                                         mp_obj_new_int(nbytes * 8));
+
+            if (mp_binary_op(MP_BINARY_OP_LESS, self_in, edge) == mp_const_true) {
+                return;
+            }
+        }
+    }
+
+    mp_raise_OverflowError_varg(translate("value must fit in %d byte(s)"), nbytes);
+}
+
+#endif // MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
+
+void mp_small_int_buffer_overflow_check(mp_int_t val, size_t nbytes, bool is_signed) {
+    // Fast path for zero.
+    if (val == 0) {
+        return;
+    }
+
+    // Trying to store negative values in unsigned bytes falls through to failure.
+    if (is_signed || val >= 0) {
+
+        if (nbytes >= sizeof(val)) {
+            // All non-negative N bit signed integers fit in an unsigned N bit integer.
+            // This case prevents shifting too far below.
+            return;
+        }
+
+        if (is_signed) {
+            mp_int_t edge = ((mp_int_t)1 << (nbytes * 8 - 1));
+            if (-edge <= val && val < edge) {
+                return;
+            }
+            // Out of range, fall through to failure.
+        } else {
+            // Unsigned. We already know val >= 0.
+            mp_int_t edge = ((mp_int_t)1 << (nbytes * 8));
+            if (val < edge) {
+                return;
+            }
+        }
+        // Fall through to failure.
+    }
+
+    mp_raise_OverflowError_varg(translate("value must fit in %d byte(s)"), nbytes);
+}
+
 #if MICROPY_LONGINT_IMPL == MICROPY_LONGINT_IMPL_NONE
 
 int mp_obj_int_sign(mp_obj_t self_in) {
@@ -420,15 +490,24 @@ STATIC mp_obj_t int_from_bytes(size_t n_args, const mp_obj_t *args) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(int_from_bytes_fun_obj, 3, 4, int_from_bytes);
 STATIC MP_DEFINE_CONST_CLASSMETHOD_OBJ(int_from_bytes_obj, MP_ROM_PTR(&int_from_bytes_fun_obj));
 
-STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
-    // TODO: Support signed param (assumes signed=False)
-    (void)n_args;
-
-    mp_int_t len = mp_obj_get_int(args[1]);
+STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    enum { ARG_length, ARG_byteorder, ARG_signed };
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_length,    MP_ARG_REQUIRED | MP_ARG_INT },
+        { MP_QSTR_byteorder, MP_ARG_REQUIRED | MP_ARG_OBJ },
+        { MP_QSTR_signed,    MP_ARG_KW_ONLY | MP_ARG_BOOL, {.u_bool = false} },
+    };
+    mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
+    mp_arg_parse_all(n_args - 1, pos_args + 1, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
+
+    mp_int_t len = args[ARG_length].u_int;
     if (len < 0) {
         mp_raise_ValueError(NULL);
     }
-    bool big_endian = args[2] != MP_OBJ_NEW_QSTR(MP_QSTR_little);
+
+    mp_obj_t self = pos_args[0];
+    bool big_endian = args[ARG_byteorder].u_obj != MP_OBJ_NEW_QSTR(MP_QSTR_little);
+    bool signed_ = args[ARG_signed].u_bool;
 
     vstr_t vstr;
     vstr_init_len(&vstr, len);
@@ -436,19 +515,26 @@ STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
     memset(data, 0, len);
 
     #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
-    if (!MP_OBJ_IS_SMALL_INT(args[0])) {
-        mp_obj_int_to_bytes_impl(args[0], big_endian, len, data);
+    if (!MP_OBJ_IS_SMALL_INT(self)) {
+        mp_obj_int_buffer_overflow_check(self, len, signed_);
+        mp_obj_int_to_bytes_impl(self, big_endian, len, data);
     } else
     #endif
     {
-        mp_int_t val = MP_OBJ_SMALL_INT_VALUE(args[0]);
+        mp_int_t val = MP_OBJ_SMALL_INT_VALUE(self);
+        // Small int checking is separate, to be fast.
+        mp_small_int_buffer_overflow_check(val, len, signed_);
         size_t l = MIN((size_t)len, sizeof(val));
+        if (val < 0) {
+            // Sign extend negative numbers.
+            memset(data, -1, len);
+        }
         mp_binary_set_int(l, big_endian, data + (big_endian ? (len - l) : 0), val);
     }
 
     return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(int_to_bytes_obj, 3, 4, int_to_bytes);
+STATIC MP_DEFINE_CONST_FUN_OBJ_KW(int_to_bytes_obj, 3, int_to_bytes);
 
 STATIC const mp_rom_map_elem_t int_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_from_bytes), MP_ROM_PTR(&int_from_bytes_obj) },