Merge pull request #447 from dhylands/str-format-mpz

Add string formatting support for longlong and mpz.

1 files changed, 37 insertions, 9 deletions

diff --git a/py/mpz.c b/py/mpz.c
index 21b390996a..a6b024ca87 100644
--- a/py/mpz.c
+++ b/py/mpz.c
@@ -1216,15 +1216,27 @@ uint mpz_as_str_size(const mpz_t *i, uint base) {
     return i->len * DIG_SIZE / log_base2_floor[base] + 2 + 1; // +1 for null byte termination
 }
 
+uint mpz_as_str_size_formatted(const mpz_t *i, uint base, const char *prefix, char comma) {
+    if (base < 2 || base > 32) {
+        return 0;
+    }
+
+    uint num_digits = i->len * DIG_SIZE / log_base2_floor[base] + 1;
+    uint num_commas = comma ? num_digits / 3: 0;
+    uint prefix_len = prefix ? strlen(prefix) : 0;
+
+    return num_digits + num_commas + prefix_len + 2; // +1 for sign, +1 for null byte
+}
+
 char *mpz_as_str(const mpz_t *i, uint base) {
     char *s = m_new(char, mpz_as_str_size(i, base));
-    mpz_as_str_inpl(i, base, s);
+    mpz_as_str_inpl(i, base, "", 'a', 0, s);
     return s;
 }
 
 // assumes enough space as calculated by mpz_as_str_size
 // returns length of string, not including null byte
-uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
+uint mpz_as_str_inpl(const mpz_t *i, uint base, const char *prefix, char base_char, char comma, char *str) {
     if (str == NULL || base < 2 || base > 32) {
         str[0] = 0;
         return 0;
@@ -1232,10 +1244,15 @@ uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
 
     uint ilen = i->len;
 
+    char *s = str;
     if (ilen == 0) {
-        str[0] = '0';
-        str[1] = 0;
-        return 1;
+        if (prefix) {
+            while (*prefix)
+                *s++ = *prefix++;
+        }
+        *s++ = '0';
+        *s = '\0';
+        return s - str;
     }
 
     // make a copy of mpz digits
@@ -1243,7 +1260,7 @@ uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
     memcpy(dig, i->dig, ilen * sizeof(mpz_dig_t));
 
     // convert
-    char *s = str;
+    char *last_comma = str;
     bool done;
     do {
         mpz_dig_t *d = dig + ilen;
@@ -1259,7 +1276,7 @@ uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
         // convert to character
         a += '0';
         if (a > '9') {
-            a += 'a' - '9' - 1;
+            a += base_char - '9' - 1;
         }
         *s++ = a;
 
@@ -1271,8 +1288,19 @@ uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
                 break;
             }
         }
-    } while (!done);
+        if (comma && (s - last_comma) == 3) {
+            *s++ = comma;
+            last_comma = s;
+        }
+    }
+    while (!done);
 
+    if (prefix) {
+        const char *p = &prefix[strlen(prefix)];
+        while (p > prefix) {
+            *s++ = *--p;
+        }
+    }
     if (i->neg != 0) {
         *s++ = '-';
     }
@@ -1284,7 +1312,7 @@ uint mpz_as_str_inpl(const mpz_t *i, uint base, char *str) {
         *v = temp;
     }
 
-    s[0] = 0; // null termination
+    *s = '\0'; // null termination
 
     return s - str;
 }