Handle truth values; speed up smallint checks

3 files changed, 53 insertions, 5 deletions

diff --git a/py/binary.c b/py/binary.c
index 0e4a1d5f6..6b46425cf 100644
--- a/py/binary.c
+++ b/py/binary.c
@@ -308,6 +308,7 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte **
             bool signed_type = is_signed(val_type);
             #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
             if (MP_OBJ_IS_TYPE(val_in, &mp_type_int)) {
+                // It's a longint.
                 mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
                 mp_obj_int_to_bytes_impl(val_in, struct_type == '>', size, p);
                 return;
@@ -315,7 +316,8 @@ void mp_binary_set_val(char struct_type, char val_type, mp_obj_t val_in, byte **
             #endif
             {
                 val = mp_obj_get_int(val_in);
-                mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
+                // Small int checking is separate, to be fast.
+                mp_small_int_buffer_overflow_check(val, size, signed_type);
                 // zero/sign extend if needed
                 if (BYTES_PER_WORD < 8 && size > sizeof(val)) {
                     int c = (is_signed(val_type) && (mp_int_t)val < 0) ? 0xff : 0x00;
@@ -353,6 +355,7 @@ void mp_binary_set_val_array(char typecode, void *p, mp_uint_t index, mp_obj_t v
 
             #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
             if (MP_OBJ_IS_TYPE(val_in, &mp_type_int)) {
+                // It's a long int.
                 mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
                 mp_obj_int_to_bytes_impl(val_in, MP_ENDIANNESS_BIG,
                     size, (uint8_t*)p + index * size);
@@ -360,7 +363,8 @@ void mp_binary_set_val_array(char typecode, void *p, mp_uint_t index, mp_obj_t v
             }
             #endif
             mp_int_t val = mp_obj_get_int(val_in);
-            mp_obj_int_buffer_overflow_check(val_in, size, signed_type);
+            // Small int checking is separate, to be fast.
+            mp_small_int_buffer_overflow_check(val, size, signed_type);
             mp_binary_set_val_array_from_int(typecode, p, index, val);
         }
     }
diff --git a/py/objint.c b/py/objint.c
index 80f6dfe55..3b3a3d9c0 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -300,6 +300,8 @@ 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) {
@@ -329,7 +331,43 @@ void mp_obj_int_buffer_overflow_check(mp_obj_t self_in, size_t nbytes, bool is_s
         }
     }
 
-    mp_raise_OverflowError_varg(translate("value would overflow a %d byte buffer"), nbytes);
+    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;
+    if (!is_signed) {
+        if (val >= 0) {
+            // Using signed constants here, not UINT8_MAX, etc. to avoid any unintended conversions.
+            if (val <= 0xff) return;        // Small values fit in any number of nbytes.
+            if (nbytes == 2 && val <= 0xffff) return;
+#if !defined(__LP64__)
+            // 32-bit ints and pointers
+            if (nbytes >= 4) return;   // Any mp_int_t will fit.
+#else
+            // 64-bit ints and pointers
+            if (nbytes == 4 && val <= 0xffffffff) return;
+            if (nbytes >= 8) return;   // Any mp_int_t will fit.
+#endif
+        } // Negative, fall through to failure.
+    } else {
+        // signed
+        if (val >= INT8_MIN && val <= INT8_MAX) return; // Small values fit in any number of nbytes.
+        if (nbytes == 2 && val >= INT16_MIN && val <= INT16_MAX) return;
+#if !defined(__LP64__)
+        // 32-bit ints and pointers
+        if (nbytes >= 4) return;   // Any mp_int_t will fit.
+#else
+        // 64-bit ints and pointers
+        if (nbytes == 4 && val >= INT32_MIN && val <= INT32_MAX) return;
+        if (nbytes >= 8) return;   // Any mp_int_t will fit.
+#endif
+    } // 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
@@ -467,15 +505,16 @@ STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
     byte *data = (byte*)vstr.buf;
     memset(data, 0, len);
 
-    mp_obj_int_buffer_overflow_check(args[0], len, false);
-
     #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
     if (!MP_OBJ_IS_SMALL_INT(args[0])) {
+        mp_obj_int_buffer_overflow_check(args[0], len, false);
         mp_obj_int_to_bytes_impl(args[0], big_endian, len, data);
     } else
     #endif
     {
         mp_int_t val = MP_OBJ_SMALL_INT_VALUE(args[0]);
+        // Small int checking is separate, to be fast.
+        mp_small_int_buffer_overflow_check(val, len, false);
         size_t l = MIN((size_t)len, sizeof(val));
         mp_binary_set_int(l, big_endian, data + (big_endian ? (len - l) : 0), val);
     }
diff --git a/py/objint.h b/py/objint.h
index 654404ac9..e8c9bc3e0 100644
--- a/py/objint.h
+++ b/py/objint.h
@@ -53,7 +53,12 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
                            int base, const char *prefix, char base_char, char comma);
 char *mp_obj_int_formatted_impl(char **buf, size_t *buf_size, size_t *fmt_size, mp_const_obj_t self_in,
                                 int base, const char *prefix, char base_char, char comma);
+#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);
+#endif
+
+void mp_small_int_buffer_overflow_check(mp_int_t val, size_t nbytes, bool is_signed);
+
 mp_int_t mp_obj_int_hash(mp_obj_t self_in);
 mp_obj_t mp_obj_int_from_bytes_impl(bool big_endian, size_t len, const byte *buf);
 void mp_obj_int_to_bytes_impl(mp_obj_t self_in, bool big_endian, size_t len, byte *buf);