lib-modulo 0.2.2

use std::num::NonZero;

use crate::{prime::primality_test, Context64};

/// Factorize integer and writes prime factors to `factor` in any order.
///
/// This function is probabilistic and may fail.
///
/// # Time complexity
///
/// O(`x`^0.25) expected
///
/// # Example
///
/// ```
/// use lib_modulo::factorize::*;
///
/// let mut factor = Vec::new();
/// // panics if factorization fails
/// assert!(factorize(998_244_353 * 1_000_000_007, &mut factor).is_ok());
///
/// factor.sort_unstable();
/// assert_eq!(factor, vec![998_244_353, 1_000_000_007])
/// ```
pub fn factorize(mut x: u64, factor: &mut Vec<u64>) -> Result<(), ()> {
    if x < 2 {
        return Ok(());
    }
    factor.reserve(64);

    // trial division by small primes less than 2^16
    {
        factor.extend(std::iter::repeat_n(2, x.trailing_zeros() as usize));
        x >>= x.trailing_zeros();
    }
    for &(n, inv_n, r2_mod_n) in SMALL_ODD_PRIME_CONTEXT_16.iter() {
        let ctx = Context64 {
            n: n as u64,
            inv_n,
            r2_mod_n: r2_mod_n as u64,
        };

        while ctx.can_divide(x) {
            x /= ctx.n;
            factor.push(ctx.n);
        }

        if x == 1 {
            return Ok(());
        }
    }

    // find large prime factors (up to 3) by Pollard's rho
    while x > 1 {
        if primality_test(x) {
            factor.push(x);
            return Ok(());
        }

        if let Some(d) = pollard_rho(x) {
            let d = d.get();
            while x % d == 0 {
                x /= d;
                factor.push(d);
            }
        } else {
            return Err(());
        }
    }

    Ok(())
}

/// Find prime factor of `x`.
///
/// This function is probabilistic and may fail.
///
/// # Time complexity
///
/// *O*(p^0.25) where p is a prime factor of `x`
fn pollard_rho(x: u64) -> Option<NonZero<u64>> {
    let ctx = Context64::new(x);
    let one = ctx.modulo(1);

    for c in 1..100 {
        // a = b (mod x) => f(a) = f(b) (mod x)
        let f = |x: u64| ctx.mul_add(x, x, c);

        let mut y0 = ctx.modulo(1);
        let mut y1 = y0;

        let mut prod = one;
        let mut step = 0;
        let mut memo = [[0, 0, one.value]; 1 << 5];

        'cycle_detection: while !prod.is_zero() {
            y0.value = f(y0.value);
            y1.value = f(f(y1.value));
            prod *= y1 - y0;
            step += 1;

            if step % (1 << 5) == 0 {
                memo[(step >> 5) % (1 << 5)] = [y0.value, y1.value, prod.value];
            }
            if step % (1 << 10) == 0 || prod.is_zero() {
                let g = binary_gcd(prod.value, x);

                if g == 1 {
                    continue 'cycle_detection;
                } else if primality_test(g) {
                    return NonZero::new(g);
                }

                for i in 0..memo.len() {
                    let g = binary_gcd(memo[i][2], x);

                    if g != 1 {
                        if primality_test(g) {
                            return NonZero::new(g);
                        }

                        y0.value = memo[i][0];
                        y1.value = memo[i][1];
                        for _ in 0..1 << 5 {
                            let g = binary_gcd((y0 - y1).value, x);

                            if g != 1 {
                                if primality_test(g) {
                                    return NonZero::new(g);
                                } else if g != x {
                                    // FIXME: `x` is composed of at most 3 primes, so return `x/g`
                                    return pollard_rho(g);
                                } else {
                                    break 'cycle_detection;
                                }
                            }

                            y0.value = f(y0.value);
                            y1.value = f(f(y1.value));
                        }
                    }
                }
            }
        }
    }

    None
}

#[inline(always)]
fn binary_gcd(mut a: u64, mut b: u64) -> u64 {
    if b == 0 {
        return a;
    }

    let shift = (a | b).trailing_zeros();
    b >>= b.trailing_zeros();

    while a != 0 {
        a >>= a.trailing_zeros();

        if a < b {
            (a, b) = (b, a)
        }
        a -= b
    }

    b << shift
}

#[cfg(test)]
mod tests {
    use rand::{rng, seq::SliceRandom, Rng};

    use super::*;

    #[test]
    fn random() {
        let mut rng = rng();
        for n in std::iter::repeat_with(|| rng.random_range(1 << 55..=u64::MAX)).take(10_000) {
            let mut factor = Vec::new();

            assert!(factorize(n, &mut factor).is_ok());
            assert_eq!(n, factor.iter().product())
        }
    }

    #[test]
    fn random_square() {
        let mut rng = rng();
        for n in std::iter::repeat_with(|| rng.random_range(1 << 20..1 << 32)).take(5000) {
            let mut factor = Vec::new();

            assert!(factorize(n * n, &mut factor).is_ok());
            assert_eq!(n * n, factor.iter().product())
        }
    }

    #[test]
    fn prime_square() {
        for n in (0..1 << 32)
            .rev()
            .step_by(2)
            .filter(|n| primality_test(*n))
            .take(500)
        {
            let mut factor = Vec::new();

            assert!(factorize(n * n, &mut factor).is_ok());
            assert_eq!(n * n, factor.iter().product())
        }
    }

    // fast since p is relatively small
    #[test]
    fn prime_cube() {
        let p = Vec::from_iter(
            (0..1 << 21)
                .rev()
                .step_by(2)
                .filter(|n| primality_test(*n))
                .take(500),
        );

        for p in p {
            let n = p.pow(3);
            let mut factor = Vec::new();

            assert!(factorize(n, &mut factor).is_ok());
            assert_eq!(n, factor.iter().product())
        }
    }

    #[test]
    fn prime_double() {
        let mut p: Vec<u64> = (0..1 << 32)
            .rev()
            .step_by(2)
            .filter(|n| primality_test(*n))
            .take(500)
            .collect();
        p.shuffle(&mut rng());

        for p in p.windows(2) {
            let n = p[0] * p[1];
            let mut factor = Vec::new();

            assert!(factorize(n, &mut factor).is_ok());
            assert_eq!(n, factor.iter().product())
        }
    }

    #[test]
    fn prime_triple() {
        let mut p: Vec<u64> = (0..1 << 21)
            .rev()
            .step_by(2)
            .filter(|n| primality_test(*n))
            .take(500)
            .collect();
        p.shuffle(&mut rng());

        for p in p.windows(3) {
            let n = p[0] * p[1] * p[2];
            let mut factor = Vec::new();

            assert!(factorize(n, &mut factor).is_ok());
            assert_eq!(n, factor.iter().product())
        }
    }
}

// 12 bytes * 171 ~ 2 KiB
static SMALL_ODD_PRIME_CONTEXT_16: [(u16, u64, u16); 171] = [
    (3, 12297829382473034411, 1),
    (5, 14757395258967641293, 1),
    (7, 7905747460161236407, 4),
    (11, 3353953467947191203, 3),
    (13, 5675921253449092805, 9),
    (17, 17361641481138401521, 1),
    (19, 9708812670373448219, 4),
    (23, 15238614669586151335, 13),
    (29, 3816567739388183093, 25),
    (31, 17256631552825064415, 8),
    (37, 1495681951922396077, 33),
    (41, 10348173504763894809, 10),
    (43, 9437869060967677571, 4),
    (47, 5887258746928580303, 14),
    (53, 2436362424829563421, 13),
    (59, 14694863923124558067, 25),
    (61, 5745707170499696405, 12),
    (67, 17345445920055250027, 21),
    (71, 1818693077689674103, 29),
    (73, 9097024474706080249, 4),
    (79, 11208148297950107311, 73),
    (83, 11779246215742243803, 51),
    (89, 17617676924329347049, 39),
    (97, 11790702397628785569, 35),
    (101, 4200743699953660269, 80),
    (103, 15760325033848937303, 38),
    (107, 8619973866219416643, 11),
    (109, 12015769075535579493, 105),
    (113, 10447713457676206225, 109),
    (127, 9150747060186627967, 4),
    (131, 281629680514649643, 33),
    (137, 16292379802327414201, 38),
    (139, 4246732448623781667, 30),
    (149, 16094474695182830269, 123),
    (151, 8062815290495565607, 105),
    (157, 6579730370240349621, 39),
    (163, 2263404180823257867, 152),
    (167, 10162278172342986519, 63),
    (173, 9809829218388894501, 133),
    (179, 17107036403551874683, 161),
    (181, 3770881385233444253, 126),
    (191, 2124755861893246783, 103),
    (193, 8124213711219232577, 108),
    (197, 14513935692512591373, 175),
    (199, 2780916192016515319, 155),
    (211, 13900627050804827995, 119),
    (223, 7527595115280579359, 171),
    (227, 1950316554048586955, 147),
    (229, 2094390156840385773, 104),
    (233, 7204522363551799129, 135),
    (239, 7255204782128442895, 34),
    (241, 17298606475760824337, 15),
    (251, 2939720171109091891, 243),
    (257, 18374966859414961921, 1),
    (263, 15430736487513693367, 33),
    (269, 10354863773718001093, 21),
    (271, 15383631589145234927, 36),
    (277, 17181443938689762877, 155),
    (281, 14245350405676059433, 85),
    (283, 5149444458738708755, 151),
    (293, 2707201348701401773, 161),
    (307, 17305088903023944187, 199),
    (311, 9134400602415662215, 35),
    (313, 6365010734698503433, 132),
    (317, 17050145153302519317, 235),
    (331, 3455281367280943203, 256),
    (337, 9196002980365592497, 4),
    (347, 9941040754419844819, 129),
    (349, 15751088062938241781, 148),
    (353, 8779186981255537313, 22),
    (359, 5600822016808749655, 264),
    (367, 9751139919072624015, 129),
    (373, 3511310534137743069, 68),
    (379, 17181268226964305331, 171),
    (383, 14834457375202459263, 150),
    (389, 12661389891209383757, 164),
    (397, 185861401246443845, 273),
    (401, 3220129888178724721, 360),
    (409, 2074694932495450793, 265),
    (419, 1849076971589024267, 100),
    (421, 14897608040525528621, 255),
    (431, 8046375605237577039, 216),
    (433, 7540585914657253201, 150),
    (439, 15379290047785184263, 36),
    (443, 15615189678648040307, 153),
    (449, 205420312624827969, 18),
    (457, 686202733595322489, 68),
    (461, 3041111821262312197, 444),
    (463, 8127723090792113455, 60),
    (467, 15247201739725667931, 264),
    (479, 8010277176057592351, 28),
    (487, 2386334448960373207, 467),
    (491, 1051952818867347139, 429),
    (499, 12494988971771199291, 462),
    (503, 17969989256695189447, 378),
    (509, 5436172123882971989, 93),
    (521, 1805727346946616377, 130),
    (523, 7195288319381928355, 7),
    (541, 13911777416032342069, 254),
    (547, 13219604528142859659, 318),
    (557, 5133295029488295333, 414),
    (563, 18151858289227516155, 424),
    (569, 6386658317259721737, 107),
    (571, 1873749835858413299, 396),
    (577, 8184343991108570561, 546),
    (587, 8107768264083584867, 456),
    (593, 13407330009728190129, 277),
    (599, 16999336775772408167, 205),
    (601, 10926856722530117097, 8),
    (607, 7810235958720518559, 197),
    (613, 10111102787547160429, 112),
    (617, 5112468778937331161, 334),
    (619, 10400506755613301315, 170),
    (631, 14909412801254946631, 281),
    (641, 18417966001831689601, 1),
    (643, 12450835035754191915, 314),
    (647, 16650538700226241335, 119),
    (653, 18023005695293558853, 386),
    (659, 13744083976011213723, 624),
    (661, 3544230707051608253, 222),
    (673, 7016889276180750689, 417),
    (677, 7711120491668837677, 667),
    (683, 18338710433453565955, 555),
    (691, 16604739238563445883, 100),
    (701, 4578792394900801685, 253),
    (709, 17796294705807240205, 211),
    (719, 7799457855921701935, 631),
    (727, 3501582781529460967, 128),
    (733, 8027982755134170485, 669),
    (739, 12755461734324196043, 707),
    (743, 8416482154761154775, 338),
    (751, 73688724661955599, 167),
    (757, 11233750354002778461, 112),
    (761, 8193166224591101769, 591),
    (769, 8635666926574042369, 360),
    (773, 11955781087876436429, 597),
    (787, 1757948926973591323, 441),
    (797, 3332912354597459765, 57),
    (809, 12062209660064712985, 676),
    (811, 2229076349227541379, 19),
    (821, 4044353146489304861, 487),
    (823, 5536264624794968711, 100),
    (827, 9502192231439261171, 262),
    (829, 9813044796750195733, 345),
    (839, 17941052639030982263, 295),
    (853, 13018686907823153661, 712),
    (857, 12247604875076703465, 315),
    (859, 858986918449804499, 677),
    (863, 13637338028999645343, 730),
    (877, 4964004106494246501, 270),
    (881, 14447506709261737361, 487),
    (883, 17047047751015848379, 179),
    (887, 15763959422628906567, 364),
    (907, 11979197639928253475, 558),
    (911, 16462352285319281519, 142),
    (919, 16098246732442938407, 67),
    (929, 13264181960428396641, 437),
    (937, 11792529028977610905, 174),
    (941, 1725094026021722149, 501),
    (947, 18271431828024877947, 152),
    (953, 18137040080866579081, 443),
    (967, 9614435380713147895, 279),
    (971, 10828675717831559651, 343),
    (977, 12064963626510136625, 496),
    (983, 9326583728009813991, 260),
    (991, 13160290676198438943, 139),
    (997, 11785933776281829869, 299),
    (1009, 6965519813759503633, 142),
    (1013, 7775675932353384541, 92),
    (1019, 10372899268140896051, 821),
    (1021, 7497942008412795221, 646),
];