py/objint_*: Convert back to small int after binary op.

Before this change, long/mpz ints propagated into all future calculations.

For example, a relatively common operation like `x = a * b // c` where
a,b,c all small ints would always result in a long/mpz int, even if it
didn't need to, and then this would impact all future calculations with
x.

This adds +24 bytes on PYBV11 but avoids heap allocations and potential
surprises (e.g. `big-big` is now a small `0`, and can safely be accessed
with MP_OBJ_SMALL_INT_VALUE).

This work was funded through GitHub Sponsors.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

py/objint_longlong.c

@@ -227,25 +227,21 @@ zero_division:
 }
 
 mp_obj_t mp_obj_new_int(mp_int_t value) {
-    if (MP_SMALL_INT_FITS(value)) {
-        return MP_OBJ_NEW_SMALL_INT(value);
-    }
     return mp_obj_new_int_from_ll(value);
 }
 
 mp_obj_t mp_obj_new_int_from_uint(mp_uint_t value) {
-    // SMALL_INT accepts only signed numbers, so make sure the input
-    // value fits completely in the small-int positive range.
-    if ((value & ~MP_SMALL_INT_POSITIVE_MASK) == 0) {
-        return MP_OBJ_NEW_SMALL_INT(value);
-    }
     return mp_obj_new_int_from_ll(value);
 }
 
 mp_obj_t mp_obj_new_int_from_ll(long long val) {
+    if ((long long)(mp_int_t)val == val && MP_SMALL_INT_FITS(val)) {
+        return MP_OBJ_NEW_SMALL_INT(val);
+    }
+
     mp_obj_int_t *o = mp_obj_malloc(mp_obj_int_t, &mp_type_int);
     o->val = val;
-    return o;
+    return MP_OBJ_FROM_PTR(o);
 }
 
 mp_obj_t mp_obj_new_int_from_ull(unsigned long long val) {
@@ -253,19 +249,16 @@ mp_obj_t mp_obj_new_int_from_ull(unsigned long long val) {
     if (val >> (sizeof(unsigned long long) * 8 - 1) != 0) {
         mp_raise_msg(&mp_type_OverflowError, MP_ERROR_TEXT("ulonglong too large"));
     }
-    mp_obj_int_t *o = mp_obj_malloc(mp_obj_int_t, &mp_type_int);
-    o->val = val;
-    return o;
+    return mp_obj_new_int_from_ll(val);
 }
 
 mp_obj_t mp_obj_new_int_from_str_len(const char **str, size_t len, bool neg, unsigned int base) {
     // TODO this does not honor the given length of the string, but it all cases it should anyway be null terminated
     // TODO check overflow
-    mp_obj_int_t *o = mp_obj_malloc(mp_obj_int_t, &mp_type_int);
     char *endptr;
-    o->val = strtoll(*str, &endptr, base);
+    mp_obj_t result = mp_obj_new_int_from_ll(strtoll(*str, &endptr, base));
     *str = endptr;
-    return o;
+    return result;
 }
 
 mp_int_t mp_obj_int_get_truncated(mp_const_obj_t self_in) {

py/objint_mpz.c

@@ -308,6 +308,13 @@ mp_obj_t mp_obj_int_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_i
             }
         }
 
+        mp_int_t res_small;
+        if (mpz_as_int_checked(&res->mpz, &res_small)) {
+            if (MP_SMALL_INT_FITS(res_small)) {
+                return MP_OBJ_NEW_SMALL_INT(res_small);
+            }
+        }
+
         return MP_OBJ_FROM_PTR(res);
 
     } else {

tests/basics/int_big_to_small.py

@@ -0,0 +1,22 @@
+try:
+    import micropython
+    micropython.heap_lock
+except:
+    print("SKIP")
+    raise SystemExit
+
+# All less than small int max.
+for d in (0, 27, 1<<29, -1861, -(1<<29)):
+    i = 1<<70
+    print(i)
+    j = (1<<70) + d
+    print(j)
+    # k should now be a small int.
+    k = j - i
+    print(k)
+
+    # Now verify that working with k doesn't allocate (i.e. it's a small int).
+    micropython.heap_lock()
+    print(k + 20)
+    print(k // 20)
+    micropython.heap_unlock()

tests/basics/int_big_to_small.py.exp

@@ -0,0 +1,25 @@
+1180591620717411303424
+1180591620717411303424
+0
+20
+0
+1180591620717411303424
+1180591620717411303451
+27
+47
+1
+1180591620717411303424
+1180591620717948174336
+536870912
+536870932
+26843545
+1180591620717411303424
+1180591620717411301563
+-1861
+-1841
+-94
+1180591620717411303424
+1180591620716874432512
+-536870912
+-536870892
+-26843546