From 8bb7d958f173eec58892bc00115516c160f93243 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 11 Oct 2016 13:11:32 +1100
Subject: [PATCH] py: Factor duplicated function to calculate size of formatted
 int.

---
 py/mpz.c        | 26 +-------------------------
 py/mpz.h        |  1 +
 py/objint.c     | 10 +++++-----
 py/objint.h     |  2 ++
 py/objint_mpz.c |  2 +-
 5 files changed, 10 insertions(+), 31 deletions(-)

diff --git a/py/mpz.c b/py/mpz.c
index 6415e1df4a..cceb079cd3 100644
--- a/py/mpz.c
+++ b/py/mpz.c
@@ -645,18 +645,6 @@ STATIC void mpn_div(mpz_dig_t *num_dig, mp_uint_t *num_len, const mpz_dig_t *den
 
 #define MIN_ALLOC (2)
 
-STATIC const uint8_t log_base2_floor[] = {
-    0,
-    0, 1, 1, 2,
-    2, 2, 2, 3,
-    3, 3, 3, 3,
-    3, 3, 3, 4,
-    4, 4, 4, 4,
-    4, 4, 4, 4,
-    4, 4, 4, 4,
-    4, 4, 4, 5
-};
-
 void mpz_init_zero(mpz_t *z) {
     z->neg = 0;
     z->fixed_dig = 0;
@@ -1652,18 +1640,6 @@ mp_float_t mpz_as_float(const mpz_t *i) {
 }
 #endif
 
-mp_uint_t mpz_as_str_size(const mpz_t *i, mp_uint_t base, const char *prefix, char comma) {
-    if (base < 2 || base > 32) {
-        return 0;
-    }
-
-    mp_uint_t num_digits = i->len * DIG_SIZE / log_base2_floor[base] + 1;
-    mp_uint_t num_commas = comma ? num_digits / 3: 0;
-    mp_uint_t prefix_len = prefix ? strlen(prefix) : 0;
-
-    return num_digits + num_commas + prefix_len + 2; // +1 for sign, +1 for null byte
-}
-
 #if 0
 this function is unused
 char *mpz_as_str(const mpz_t *i, mp_uint_t base) {
@@ -1673,7 +1649,7 @@ char *mpz_as_str(const mpz_t *i, mp_uint_t base) {
 }
 #endif
 
-// assumes enough space as calculated by mpz_as_str_size
+// assumes enough space as calculated by mp_int_format_size
 // returns length of string, not including null byte
 mp_uint_t mpz_as_str_inpl(const mpz_t *i, mp_uint_t base, const char *prefix, char base_char, char comma, char *str) {
     if (str == NULL || base < 2 || base > 32) {
diff --git a/py/mpz.h b/py/mpz.h
index 63ac772ffd..55ef3e15ff 100644
--- a/py/mpz.h
+++ b/py/mpz.h
@@ -127,6 +127,7 @@ void mpz_or_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs);
 void mpz_xor_inpl(mpz_t *dest, const mpz_t *lhs, const mpz_t *rhs);
 void mpz_divmod_inpl(mpz_t *dest_quo, mpz_t *dest_rem, const mpz_t *lhs, const mpz_t *rhs);
 
+static inline size_t mpz_max_num_bits(const mpz_t *z) { return z->len * MPZ_DIG_SIZE; }
 mp_int_t mpz_hash(const mpz_t *z);
 bool mpz_as_int_checked(const mpz_t *z, mp_int_t *value);
 bool mpz_as_uint_checked(const mpz_t *z, mp_uint_t *value);
diff --git a/py/objint.c b/py/objint.c
index 9f948a1455..31067aaa52 100644
--- a/py/objint.c
+++ b/py/objint.c
@@ -162,14 +162,14 @@ STATIC const uint8_t log_base2_floor[] = {
     4, 4, 4, 5
 };
 
-STATIC uint int_as_str_size_formatted(uint base, const char *prefix, char comma) {
+size_t mp_int_format_size(size_t num_bits, int base, const char *prefix, char comma) {
     if (base < 2 || base > 32) {
         return 0;
     }
 
-    uint num_digits = sizeof(fmt_int_t) * 8 / log_base2_floor[base] + 1;
-    uint num_commas = comma ? num_digits / 3: 0;
-    uint prefix_len = prefix ? strlen(prefix) : 0;
+    size_t num_digits = num_bits / log_base2_floor[base] + 1;
+    size_t num_commas = comma ? num_digits / 3 : 0;
+    size_t prefix_len = prefix ? strlen(prefix) : 0;
     return num_digits + num_commas + prefix_len + 2; // +1 for sign, +1 for null byte
 }
 
@@ -211,7 +211,7 @@ char *mp_obj_int_formatted(char **buf, mp_uint_t *buf_size, mp_uint_t *fmt_size,
         sign = '-';
     }
 
-    uint needed_size = int_as_str_size_formatted(base, prefix, comma);
+    uint needed_size = mp_int_format_size(sizeof(fmt_int_t) * 8, base, prefix, comma);
     if (needed_size > *buf_size) {
         *buf = m_new(char, needed_size);
         *buf_size = needed_size;
diff --git a/py/objint.h b/py/objint.h
index c79eb874a9..dc53e0708a 100644
--- a/py/objint.h
+++ b/py/objint.h
@@ -50,6 +50,8 @@ typedef enum {
 mp_fp_as_int_class_t mp_classify_fp_as_int(mp_float_t val);
 #endif // MICROPY_PY_BUILTINS_FLOAT
 
+size_t mp_int_format_size(size_t num_bits, int base, const char *prefix, char comma);
+
 void mp_obj_int_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind);
 char *mp_obj_int_formatted(char **buf, mp_uint_t *buf_size, mp_uint_t *fmt_size, mp_const_obj_t self_in,
                            int base, const char *prefix, char base_char, char comma);
diff --git a/py/objint_mpz.c b/py/objint_mpz.c
index 3a30eb9d9b..137d514de2 100644
--- a/py/objint_mpz.c
+++ b/py/objint_mpz.c
@@ -95,7 +95,7 @@ char *mp_obj_int_formatted_impl(char **buf, mp_uint_t *buf_size, mp_uint_t *fmt_
     assert(MP_OBJ_IS_TYPE(self_in, &mp_type_int));
     const mp_obj_int_t *self = MP_OBJ_TO_PTR(self_in);
 
-    mp_uint_t needed_size = mpz_as_str_size(&self->mpz, base, prefix, comma);
+    mp_uint_t needed_size = mp_int_format_size(mpz_max_num_bits(&self->mpz), base, prefix, comma);
     if (needed_size > *buf_size) {
         *buf = m_new(char, needed_size);
         *buf_size = needed_size;