Add faster base58 encode/decode

Add optional gmpy2.mpz for even faster encode/decode
Add longer random benchmark

Author: malvidin

base58/__init__.py

@@ -12,6 +12,12 @@
 from functools import lru_cache
 from hashlib import sha256
 from typing import Mapping, Union
+from math import log
+
+try:
+    from gmpy2 import mpz
+except ImportError:
+    mpz = None
 
 __version__ = '2.1.1'
 
@@ -20,6 +26,10 @@
     b'123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'
 RIPPLE_ALPHABET = b'rpshnaf39wBUDNEGHJKLM4PQRST7VWXYZ2bcdeCg65jkm8oFqi1tuvAxyz'
 XRP_ALPHABET = RIPPLE_ALPHABET
+POWERS = {
+    45: {2 ** i: 45 ** (2 ** i) for i in range(4, 20)},
+    58: {2 ** i: 58 ** (2 ** i) for i in range(4, 20)}
+}
 
 # Retro compatibility
 alphabet = BITCOIN_ALPHABET
@@ -32,19 +42,64 @@ def scrub_input(v: Union[str, bytes]) -> bytes:
     return v
 
 
+def _encode_int(i: int, base: int = 58, alphabet: bytes = BITCOIN_ALPHABET) -> bytes:
+    """
+    Encode integer to bytes with base 58 alphabet by powers of 58
+    """
+    min_val = POWERS[base][2**8]
+    if i <= min_val:
+        string = bytearray()
+        while i:
+            i, idx = divmod(i, base)
+            string.append(idx)
+        return string[::-1]
+    else:
+        origlen0 = int(log(i, 58))//2
+        try:
+            split_num = POWERS[base][2**origlen0]
+        except KeyError:
+            POWERS[base][2**origlen0] = split_num = base ** origlen0
+        i1, i0 = divmod(i, split_num)
+
+        v1 = _encode_int(i1, base, alphabet)
+        v0 = _encode_int(i0, base, alphabet)
+        newlen0 = len(v0)
+        if newlen0 < origlen0:
+            v0[:0] = b'\0' * (origlen0 - newlen0)
+
+        return v1 + v0
+
+
+def _mpz_encode(i: int, alphabet: bytes) -> bytes:
+    """
+    Encode an integer to arbitrary base using gmpy2 mpz
+    """
+    base = len(alphabet)
+
+    raw: bytes = mpz(i).digits(base).encode()
+    tr_bytes = bytes.maketrans(''.join([mpz(x).digits(base) for x in range(base)]).encode(), alphabet)
+    encoded: bytes = raw.translate(tr_bytes)
+
+    return encoded
+
+
 def b58encode_int(
     i: int, default_one: bool = True, alphabet: bytes = BITCOIN_ALPHABET
 ) -> bytes:
     """
     Encode an integer using Base58
     """
-    if not i and default_one:
-        return alphabet[0:1]
-    string = b""
+    if not i:
+        if default_one:
+            return alphabet[0:1]
+        return b''
+    if mpz:
+        return _mpz_encode(i, alphabet)
+
     base = len(alphabet)
-    while i:
-        i, idx = divmod(i, base)
-        string = alphabet[idx:idx+1] + string
+    raw_string = _encode_int(i, base, alphabet)
+    string = raw_string.translate(bytes.maketrans(bytearray(range(len(alphabet))), alphabet))
+
     return string
 
 
@@ -82,6 +137,24 @@ def _get_base58_decode_map(alphabet: bytes,
     return invmap
 
 
+def _decode(data: bytes, min_split: int = 256, base: int = 58) -> int:
+    """
+    Decode larger data blocks recursively
+    """
+    if len(data) <= min_split:
+        ret_int = 0
+        for val in data:
+            ret_int = base * ret_int + val
+        return ret_int
+    else:
+        split_len = 2**(len(data).bit_length()-2)
+        try:
+            base_pow = POWERS[base][split_len]
+        except KeyError:
+            POWERS[base] = base_pow = base ** split_len
+        return (base_pow * _decode(data[:-split_len])) + _decode(data[-split_len:])
+
+
 def b58decode_int(
     v: Union[str, bytes], alphabet: bytes = BITCOIN_ALPHABET, *,
     autofix: bool = False
@@ -93,18 +166,29 @@ def b58decode_int(
         v = v.rstrip()
     v = scrub_input(v)
 
+    base = len(alphabet)
     map = _get_base58_decode_map(alphabet, autofix=autofix)
+    if mpz:
+        tr_bytes = bytes.maketrans(bytearray(map.keys()), ''.join([mpz(x).digits(base) for x in map.values()]).encode())
+    else:
+        tr_bytes = bytes.maketrans(bytearray(map.keys()), bytearray(map.values()))
+    del_chars = bytes(bytearray(x for x in range(256) if x not in map))
 
-    decimal = 0
-    base = len(alphabet)
-    try:
-        for char in v:
-            decimal = decimal * base + map[char]
-    except KeyError as e:
-        raise ValueError(
-            "Invalid character {!r}".format(chr(e.args[0]))
-        ) from None
-    return decimal
+    cv = v.translate(tr_bytes, delete=del_chars)
+    if len(v) != len(cv):
+        err_char = chr(next(c for c in v if c not in map))
+        raise ValueError("Invalid character {!r}".format(err_char))
+
+    if cv == b'':
+        return 0
+
+    if mpz:
+        try:
+            return int(mpz(cv, base=base))
+        except ValueError:
+            raise ValueError(cv, base)
+
+    return _decode(cv, base=base)
 
 
 def b58decode(
@@ -123,12 +207,7 @@ def b58decode(
 
     acc = b58decode_int(v, alphabet=alphabet, autofix=autofix)
 
-    result = []
-    while acc > 0:
-        acc, mod = divmod(acc, 256)
-        result.append(mod)
-
-    return b'\0' * (origlen - newlen) + bytes(reversed(result))
+    return acc.to_bytes(origlen - newlen + (acc.bit_length() + 7) // 8, "big")
 
 
 def b58encode_check(

test_base58.py

@@ -134,14 +134,14 @@ def test_invalid_input():
         raises(ValueError, "Invalid character '\\\\x08'"))
 
 
-@pytest.mark.parametrize('length', [8, 32, 256, 1024])
+@pytest.mark.parametrize('length', [8, 32, 256, 1024, 8192])
 def test_encode_random(benchmark, length) -> None:
     data = getrandbits(length * 8).to_bytes(length, byteorder='big')
     encoded = benchmark(lambda: b58encode(data))
     assert_that(b58decode(encoded), equal_to(data))
 
 
-@pytest.mark.parametrize('length', [8, 32, 256, 1024])
+@pytest.mark.parametrize('length', [8, 32, 256, 1024, 8192])
 def test_decode_random(benchmark, length) -> None:
     origdata = getrandbits(length * 8).to_bytes(length, byteorder='big')
     encoded = b58encode(origdata)