micropython · arbrauns · Jul 14, 2023 · Jul 13, 2023 · Jul 17, 2023 · Jul 13, 2023
diff --git a/py/mpz.c b/py/mpz.c
@@ -850,25 +850,43 @@ size_t mpz_set_from_str(mpz_t *z, const char *str, size_t len, bool neg, unsigne
     return cur - str;
 }
 
-void mpz_set_from_bytes(mpz_t *z, bool big_endian, size_t len, const byte *buf) {
+void mpz_set_from_bytes(mpz_t *z, bool is_signed, bool big_endian, size_t len, const byte *buf) {
     int delta = 1;
     if (big_endian) {
+        z->neg = is_signed && (buf[0] & 0x80);
         buf += len - 1;
         delta = -1;
+    } else {
+        z->neg = is_signed && (buf[len - 1] & 0x80);
     }
 
     mpz_need_dig(z, (len * 8 + DIG_SIZE - 1) / DIG_SIZE);
 
     mpz_dig_t d = 0;
+    byte carry = 1;
     int num_bits = 0;
-    z->neg = 0;
     z->len = 0;
     while (len) {
-        while (len && num_bits < DIG_SIZE) {
-            d |= *buf << num_bits;
+        while (num_bits < DIG_SIZE) {
+            byte b;
+            if (len) {
+                b = *buf;
+                buf += delta;
+                len--;
+            } else if (z->neg) {
+                // sign-extend missing bytes
+                b = 0xff;
+            } else {
+                b = 0;
+            }
+
+            if (z->neg) {
+                b = (~b) + carry;
+                carry &= b == 0;
+            }
+
+            d |= b << num_bits;
             num_bits += 8;
-            buf += delta;
-            len--;
         }
         z->dig[z->len++] = d & DIG_MASK;
         // Need this #if because it's C undefined behavior to do: uint32_t >> 32

diff --git a/py/mpz.h b/py/mpz.h
@@ -114,7 +114,7 @@ void mpz_set_from_ll(mpz_t *z, long long i, bool is_signed);
 void mpz_set_from_float(mpz_t *z, mp_float_t src);
 #endif
 size_t mpz_set_from_str(mpz_t *z, const char *str, size_t len, bool neg, unsigned int base);
-void mpz_set_from_bytes(mpz_t *z, bool big_endian, size_t len, const byte *buf);
+void mpz_set_from_bytes(mpz_t *z, bool is_signed, bool big_endian, size_t len, const byte *buf);
 
 static inline bool mpz_is_zero(const mpz_t *z) {
     return z->len == 0;

diff --git a/py/objint.c b/py/objint.c
@@ -409,7 +409,7 @@ STATIC mp_obj_t int_from_bytes(size_t n_args, const mp_obj_t *args) {
         #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
         if (value > (MP_SMALL_INT_MAX >> 8)) {
             // Result will overflow a small-int so construct a big-int
-            return mp_obj_int_from_bytes_impl(args[2] != MP_OBJ_NEW_QSTR(MP_QSTR_little), bufinfo.len, bufinfo.buf);
+            return mp_obj_int_from_bytes_impl(false, args[2] != MP_OBJ_NEW_QSTR(MP_QSTR_little), bufinfo.len, bufinfo.buf);
         }
         #endif
         value = (value << 8) | *buf;

diff --git a/py/objint.h b/py/objint.h
@@ -54,7 +54,8 @@ char *mp_obj_int_formatted(char **buf, size_t *buf_size, size_t *fmt_size, mp_co
 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);
 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);
+mp_obj_t mp_obj_int_from_bytes_impl(bool is_signed, bool big_endian, size_t len, const byte *buf);
+size_t mp_obj_int_max_bytes_needed_impl(mp_obj_t self_in);
 void mp_obj_int_to_bytes_impl(mp_obj_t self_in, bool big_endian, size_t len, byte *buf);
 int mp_obj_int_sign(mp_obj_t self_in);
 mp_obj_t mp_obj_int_unary_op(mp_unary_op_t op, mp_obj_t o_in);

diff --git a/py/objint_longlong.c b/py/objint_longlong.c
@@ -43,7 +43,7 @@
 const mp_obj_int_t mp_sys_maxsize_obj = {{&mp_type_int}, MP_SSIZE_MAX};
 #endif
 
-mp_obj_t mp_obj_int_from_bytes_impl(bool big_endian, size_t len, const byte *buf) {
+mp_obj_t mp_obj_int_from_bytes_impl(bool is_signed, bool big_endian, size_t len, const byte *buf) {
     int delta = 1;
     if (!big_endian) {
         buf += len - 1;
@@ -57,6 +57,11 @@ mp_obj_t mp_obj_int_from_bytes_impl(bool big_endian, size_t len, const byte *buf
     return mp_obj_new_int_from_ll(value);
 }
 
+size_t mp_obj_int_max_bytes_needed_impl(mp_obj_t self_in) {
+    assert(mp_obj_is_exact_type(self_in, &mp_type_int));
+    return sizeof(mp_longint_impl_t);
+}
+
 void mp_obj_int_to_bytes_impl(mp_obj_t self_in, bool big_endian, size_t len, byte *buf) {
     assert(mp_obj_is_exact_type(self_in, &mp_type_int));
     mp_obj_int_t *self = self_in;

diff --git a/py/objint_mpz.c b/py/objint_mpz.c
@@ -106,12 +106,19 @@ char *mp_obj_int_formatted_impl(char **buf, size_t *buf_size, size_t *fmt_size,
     return str;
 }
 
-mp_obj_t mp_obj_int_from_bytes_impl(bool big_endian, size_t len, const byte *buf) {
+mp_obj_t mp_obj_int_from_bytes_impl(bool is_signed, bool big_endian, size_t len, const byte *buf) {
     mp_obj_int_t *o = mp_obj_int_new_mpz();
-    mpz_set_from_bytes(&o->mpz, big_endian, len, buf);
+    mpz_set_from_bytes(&o->mpz, is_signed, big_endian, len, buf);
     return MP_OBJ_FROM_PTR(o);
 }
 
+size_t mp_obj_int_max_bytes_needed_impl(mp_obj_t self_in) {
+    assert(mp_obj_is_exact_type(self_in, &mp_type_int));
+    mp_obj_int_t *self = MP_OBJ_TO_PTR(self_in);
+    // two's complement may require one more bit than the magnitude itself
+    return (mpz_max_num_bits(&self->mpz) + 1 + 7) / 8;
+}
+
 void mp_obj_int_to_bytes_impl(mp_obj_t self_in, bool big_endian, size_t len, byte *buf) {
     assert(mp_obj_is_exact_type(self_in, &mp_type_int));
     mp_obj_int_t *self = MP_OBJ_TO_PTR(self_in);