A faster implementation of `BigUint::to_string` inspired by `_pylong.py`

bigint_to_string.rs

// A faster implementation of `BigUint::to_string` ported from `_pylong.py`
// See: https://github.com/python/cpython/blob/main/Lib/_pylong.py

use num_bigint::BigUint;
use num_integer::Integer;
use std::collections::HashMap;

fn div2n1n(a: &BigUint, b: &BigUint, n: usize) -> (BigUint, BigUint) {
    if a.bits() as usize <= 4000 + n {
        return a.div_rem(b);
    }
    let aa;
    let bb;
    let pad = n & 1;
    let (a, b, n) = if pad != 0 {
        aa = a << 1;
        bb = b << 1;
        (&aa, &bb, n + 1)
    } else {
        (a, b, n)
    };
    let half_n = n >> 1;
    let mask = (BigUint::from(1u32) << half_n) - 1u32;
    let b1 = b >> half_n;
    let b2 = b & &mask;
    let (q1, r) = div3n2n(&(a >> n), &((a >> half_n) & &mask), b, &b1, &b2, half_n);
    let (q2, mut r) = div3n2n(&r, &(a & &mask), b, &b1, &b2, half_n);
    if pad != 0 {
        r >>= 1;
    }
    ((q1 << half_n) | q2, r)
}

fn div3n2n(
    a12: &BigUint,
    a3: &BigUint,
    b: &BigUint,
    b1: &BigUint,
    b2: &BigUint,
    n: usize,
) -> (BigUint, BigUint) {
    let (mut q, r) = if a12 >> n == *b1 {
        ((BigUint::from(1u32) << n) - 1u32, a12 - (b1 << n) + b1)
    } else {
        div2n1n(a12, b1, n)
    };
    let mut r1 = r << n | a3;
    let r2 = &q * b2;
    while r1 < r2 {
        q -= 1u32;
        r1 += b;
    }
    (q, r1 - r2)
}

fn int2digits(a: &BigUint, n: usize) -> Vec<BigUint> {
    let length = (a.bits() as usize).div_ceil(n);
    let mut a_digits = vec![BigUint::ZERO; length];

    fn inner(x: &BigUint, left: usize, right: usize, n: usize, a_digits: &mut Vec<BigUint>) {
        if left + 1 == right {
            a_digits[left] = x.clone();
            return;
        }
        let mid = (left + right) >> 1;
        let shift = (mid - left) * n;
        let upper = x >> shift;
        let lower = x - (&upper << shift);
        inner(&lower, left, mid, n, a_digits);
        inner(&upper, mid, right, n, a_digits);
    }

    if a != &BigUint::ZERO {
        inner(a, 0, a_digits.len(), n, &mut a_digits);
    }
    a_digits
}

fn digits2int(digits: &[BigUint], n: usize) -> BigUint {
    fn inner(left: usize, right: usize, n: usize, digits: &[BigUint]) -> BigUint {
        if left + 1 == right {
            return digits[left].clone();
        }
        let mid = (left + right) >> 1;
        let shift = (mid - left) * n;
        (inner(mid, right, n, digits) << shift) + inner(left, mid, n, digits)
    }

    if digits.is_empty() {
        BigUint::ZERO
    } else {
        inner(0, digits.len(), n, digits)
    }
}

fn divmod_pos(a: &BigUint, b: &BigUint) -> (BigUint, BigUint) {
    let n = b.bits() as usize;
    let a_digits = int2digits(a, n);

    let mut r = BigUint::ZERO;
    let mut q_digits = Vec::new();
    for a_digit in a_digits.iter().rev() {
        let a = &(r << n) + a_digit;
        let (q_digit, new_r) = div2n1n(&a, b, n);
        q_digits.push(q_digit);
        r = new_r;
    }
    q_digits.reverse();
    let q = digits2int(&q_digits, n);
    (q, r)
}

fn int_divmod(a: &BigUint, b: &BigUint) -> (BigUint, BigUint) {
    if b == &BigUint::ZERO {
        panic!("division by zero");
    } else if a == &BigUint::ZERO {
        (BigUint::ZERO, BigUint::ZERO)
    } else {
        divmod_pos(a, b)
    }
}

fn bigint_to_string(n: &BigUint) -> String {
    let mut cache: HashMap<u32, BigUint> = HashMap::new();

    fn inner(n: &BigUint, w: u32, cache: &mut HashMap<u32, BigUint>) -> String {
        if w <= 1000 {
            return n.to_string();
        }
        let q = w / 2;
        let w1 = q + w % 2;
        let w2 = q;
        let power10 = cache
            .entry(w2)
            .or_insert_with_key(|k| BigUint::from(10u32).pow(*k));
        let (hi, lo) = int_divmod(n, power10);
        let hi_str = inner(&hi, w1, cache);
        let lo_str = inner(&lo, w2, cache);
        format!(
            "{}{}{}",
            hi_str,
            "0".repeat(w2 as usize - lo_str.len()),
            lo_str
        )
    }

    let w = (n.bits() as f64 * 2f64.log10() + 1f64) as u32;
    let s = inner(n, w, &mut cache);
    if s.starts_with('0') && n != &BigUint::ZERO {
        s.trim_start_matches('0').to_string()
    } else {
        s
    }
}

fn main() {
    let x = BigUint::from(1u32) << (1u32 << 23);
    bigint_to_string(&x);
}

#[test]
fn test() {
    let x = BigUint::from(1u32) << (1u32 << 23);
    let s = bigint_to_string(&x);
    assert_eq!(s, x.to_string());
}