randperm_crt/
lib.rs

1#![warn(clippy::must_use_candidate)]
2#![deny(clippy::use_self)]
3#![deny(clippy::double_must_use)]
4#![deny(clippy::if_not_else)]
5#![deny(clippy::inconsistent_struct_constructor)]
6#![deny(clippy::iter_not_returning_iterator)]
7#![deny(clippy::map_unwrap_or)]
8#![deny(clippy::mod_module_files)]
9#![deny(clippy::semicolon_if_nothing_returned)]
10
11mod crt;
12
13use rand::Rng;
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16struct FactoredInteger {
17    factors: Vec<(u8, u8)>,
18}
19
20impl FactoredInteger {
21    fn new(mut n: u64) -> Option<Self> {
22        let mut factors = Vec::new();
23
24        let pow2 = n.trailing_zeros() as u8;
25        if pow2 != 0 {
26            n >>= pow2;
27            factors.push((2, pow2));
28        }
29
30        for p in (3..u8::MAX).step_by(2) {
31            let q = p as u64;
32
33            let mut counter = 0;
34            while n % q == 0 {
35                counter += 1;
36                n /= q;
37            }
38
39            if counter > 0 {
40                factors.push((p, counter));
41            }
42
43            if n == 1 {
44                break;
45            }
46        }
47
48        if n == 1 {
49            Some(Self { factors })
50        } else {
51            None
52        }
53    }
54}
55
56pub trait Permutation: Sized {
57    fn num_points(&self) -> u64;
58    fn nth(&self, n: u64) -> Option<u64>;
59
60    fn iter(&self) -> PermutationIter<'_, Self> {
61        PermutationIter { perm: self, idx: 0 }
62    }
63}
64
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct RandomPermutation {
67    num_points: u64,
68    sub_perms: Vec<Vec<u64>>,
69}
70
71impl RandomPermutation {
72    #[must_use]
73    pub fn new(n: u64) -> Option<Self> {
74        Self::with_rng(n, &mut rand::rng())
75    }
76
77    pub fn with_rng<R: Rng>(n: u64, rng: &mut R) -> Option<Self> {
78        let factored_n = FactoredInteger::new(n)?;
79        let num_prime_powers = factored_n.factors.len();
80
81        let mut order = (0..num_prime_powers).collect::<Vec<_>>();
82        for a in 0..num_prime_powers {
83            let b = rng.random_range(a..num_prime_powers);
84            order.swap(a, b);
85        }
86
87        let sub_perms = (0..num_prime_powers)
88            .map(|i| {
89                let (p, k) = factored_n.factors[order[i]];
90                let pk = (p as u64).pow(k as u32);
91                let mut vec = (0..pk).collect::<Vec<_>>();
92
93                let pk = pk as usize;
94                for a in 0..pk {
95                    let b = rng.random_range(a..pk);
96                    vec.swap(a, b);
97                }
98
99                vec
100            })
101            .collect();
102
103        Some(Self {
104            num_points: n,
105            sub_perms,
106        })
107    }
108
109    #[must_use]
110    pub fn inverse(&self) -> Inverse<'_> {
111        Inverse { perm: self }
112    }
113}
114
115impl Permutation for RandomPermutation {
116    fn num_points(&self) -> u64 {
117        self.num_points
118    }
119
120    fn nth(&self, mut n: u64) -> Option<u64> {
121        if n >= self.num_points {
122            return None;
123        }
124
125        let remainders = self.sub_perms.iter().fold(Vec::new(), |mut rem, perm| {
126            let pk = perm.len() as u64;
127            rem.push(perm[(n % pk) as usize]);
128            n /= pk;
129            rem
130        });
131
132        let moduli = self
133            .sub_perms
134            .iter()
135            .map(|perm| perm.len() as u64)
136            .collect::<Vec<_>>();
137
138        Some(crt::chinese_remainder(&remainders, &moduli).unwrap())
139    }
140}
141
142pub struct Inverse<'a> {
143    perm: &'a RandomPermutation,
144}
145
146impl Permutation for Inverse<'_> {
147    fn num_points(&self) -> u64 {
148        self.perm.num_points()
149    }
150
151    fn nth(&self, n: u64) -> Option<u64> {
152        if n >= self.num_points() {
153            None
154        } else {
155            Some(self.perm.sub_perms.iter().rev().fold(0, |idx, perm| {
156                let pk = perm.len() as u64;
157                let pos = perm.iter().position(|&a| a == n % pk).unwrap() as u64;
158                idx * pk + pos
159            }))
160        }
161    }
162}
163
164pub struct PermutationIter<'a, P: Permutation> {
165    perm: &'a P,
166    idx: u64,
167}
168
169impl<P: Permutation> Iterator for PermutationIter<'_, P> {
170    type Item = u64;
171
172    fn next(&mut self) -> Option<Self::Item> {
173        let a = self.perm.nth(self.idx);
174        self.idx += 1;
175        a
176    }
177
178    fn nth(&mut self, n: usize) -> Option<Self::Item> {
179        self.idx += n as u64;
180        self.perm.nth(self.idx)
181    }
182}
183
184pub struct Composition<'a> {
185    perms: &'a [RandomPermutation],
186}
187
188impl<'a> Composition<'a> {
189    #[must_use]
190    pub fn new(perms: &'a [RandomPermutation]) -> Option<Self> {
191        if perms.is_empty() {
192            return None;
193        }
194
195        if !perms
196            .iter()
197            .all(|p| p.num_points() == perms[0].num_points())
198        {
199            return None;
200        }
201
202        Some(Self { perms })
203    }
204}
205
206impl Permutation for Composition<'_> {
207    fn num_points(&self) -> u64 {
208        self.perms[0].num_points()
209    }
210
211    fn nth(&self, n: u64) -> Option<u64> {
212        self.perms.iter().try_fold(n, |n, perm| perm.nth(n))
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use rand::SeedableRng;
219    use rand_xoshiro::Xoshiro256StarStar;
220
221    use crate::*;
222
223    const SEED: [u8; 32] = [
224        144, 115, 104, 224, 226, 59, 231, 208, 100, 18, 137, 138, 234, 236, 129, 82, 184, 196, 19,
225        43, 145, 94, 60, 77, 184, 198, 244, 164, 174, 224, 59, 152,
226    ];
227
228    mod factored_integer {
229        use super::*;
230
231        #[test]
232        fn test_new_1() {
233            let n = FactoredInteger::new(14237396402848819200);
234            assert_eq!(
235                n,
236                Some(FactoredInteger {
237                    factors: vec![
238                        (2, 11),
239                        (3, 3),
240                        (5, 2),
241                        (7, 3),
242                        (11, 4),
243                        (13, 1),
244                        (19, 3),
245                        (23, 1)
246                    ]
247                })
248            );
249        }
250
251        #[test]
252        fn test_new_2() {
253            let n = FactoredInteger::new(8929777156897433877);
254            assert_eq!(
255                n,
256                Some(FactoredInteger {
257                    factors: vec![(3, 25), (199, 1), (211, 1), (251, 1)]
258                })
259            );
260        }
261
262        #[test]
263        fn test_new_3() {
264            let n = FactoredInteger::new(2u64.pow(63));
265            assert_eq!(
266                n,
267                Some(FactoredInteger {
268                    factors: vec![(2, 63)]
269                })
270            );
271        }
272
273        #[test]
274        fn test_new_4() {
275            let n = FactoredInteger::new(257);
276            assert_eq!(n, None);
277        }
278
279        #[test]
280        fn test_new_5() {
281            let n = FactoredInteger::new(1297068779 * 3196491187);
282            assert_eq!(n, None);
283        }
284    }
285
286    mod random_permutation {
287        use super::*;
288
289        #[test]
290        fn test_random_permutation() {
291            let mut rng = Xoshiro256StarStar::from_seed(SEED);
292
293            let p = RandomPermutation::with_rng(362880, &mut rng).unwrap();
294            let mut vec = p.iter().collect::<Vec<_>>();
295            vec.sort();
296
297            assert!(vec.iter().copied().eq(0..362880));
298        }
299
300        #[test]
301        fn test_nth() {
302            let mut rng = Xoshiro256StarStar::from_seed(SEED);
303            let p = RandomPermutation::with_rng((1..=20).product(), &mut rng).unwrap();
304
305            assert_eq!(p.nth(0), Some(1319942158626894223));
306            assert_eq!(p.nth(1), Some(2412265509236814223));
307            assert_eq!(p.nth(2), Some(2263312325062734223));
308            assert_eq!(p.nth(1000000000000000000), Some(320224176323398978));
309            assert_eq!(p.nth(1146506666437230775), Some(0));
310            assert_eq!(p.nth(2432902008176639999), Some(725036512928904481));
311            assert_eq!(p.nth(2432902008176640000), None);
312            assert_eq!(p.nth(u64::MAX), None);
313        }
314
315        #[test]
316        fn test_nth_2() {
317            let mut rng = Xoshiro256StarStar::from_seed(SEED);
318            let p = RandomPermutation::with_rng(300, &mut rng).unwrap();
319            let v = p.iter().collect::<Vec<_>>();
320
321            assert_eq!(
322                v,
323                &[
324                    176, 276, 76, 101, 201, 1, 26, 126, 226, 251, 51, 151, 200, 0, 100, 125, 225,
325                    25, 50, 150, 250, 275, 75, 175, 152, 252, 52, 77, 177, 277, 2, 102, 202, 227,
326                    27, 127, 224, 24, 124, 149, 249, 49, 74, 174, 274, 299, 99, 199, 8, 108, 208,
327                    233, 33, 133, 158, 258, 58, 83, 183, 283, 284, 84, 184, 209, 9, 109, 134, 234,
328                    34, 59, 159, 259, 236, 36, 136, 161, 261, 61, 86, 186, 286, 11, 111, 211, 104,
329                    204, 4, 29, 129, 229, 254, 54, 154, 179, 279, 79, 212, 12, 112, 137, 237, 37,
330                    62, 162, 262, 287, 87, 187, 272, 72, 172, 197, 297, 97, 122, 222, 22, 47, 147,
331                    247, 68, 168, 268, 293, 93, 193, 218, 18, 118, 143, 243, 43, 56, 156, 256, 281,
332                    81, 181, 206, 6, 106, 131, 231, 31, 80, 180, 280, 5, 105, 205, 230, 30, 130,
333                    155, 255, 55, 44, 144, 244, 269, 69, 169, 194, 294, 94, 119, 219, 19, 128, 228,
334                    28, 53, 153, 253, 278, 78, 178, 203, 3, 103, 20, 120, 220, 245, 45, 145, 170,
335                    270, 70, 95, 195, 295, 164, 264, 64, 89, 189, 289, 14, 114, 214, 239, 39, 139,
336                    296, 96, 196, 221, 21, 121, 146, 246, 46, 71, 171, 271, 260, 60, 160, 185, 285,
337                    85, 110, 210, 10, 35, 135, 235, 116, 216, 16, 41, 141, 241, 266, 66, 166, 191,
338                    291, 91, 92, 192, 292, 17, 117, 217, 242, 42, 142, 167, 267, 67, 248, 48, 148,
339                    173, 273, 73, 98, 198, 298, 23, 123, 223, 32, 132, 232, 257, 57, 157, 182, 282,
340                    82, 107, 207, 7, 140, 240, 40, 65, 165, 265, 290, 90, 190, 215, 15, 115, 188,
341                    288, 88, 113, 213, 13, 38, 138, 238, 263, 63, 163
342                ]
343            );
344        }
345    }
346
347    mod inverse {
348        use super::*;
349
350        #[test]
351        fn test_nth_1() {
352            let mut rng = Xoshiro256StarStar::from_seed(SEED);
353            let p = RandomPermutation::with_rng((1..=20).product(), &mut rng).unwrap();
354            let inv = p.inverse();
355
356            assert_eq!(inv.nth(1319942158626894223), Some(0));
357            assert_eq!(inv.nth(2412265509236814223), Some(1));
358            assert_eq!(inv.nth(2263312325062734223), Some(2));
359            assert_eq!(inv.nth(320224176323398978), Some(1000000000000000000));
360            assert_eq!(inv.nth(0), Some(1146506666437230775));
361            assert_eq!(inv.nth(725036512928904481), Some(2432902008176639999));
362            assert_eq!(inv.nth(2432902008176640000), None);
363            assert_eq!(inv.nth(u64::MAX), None);
364        }
365
366        #[test]
367        fn test_nth_2() {
368            let mut rng = Xoshiro256StarStar::from_seed(SEED);
369            let p = RandomPermutation::with_rng(300, &mut rng).unwrap();
370            let inv = p.inverse();
371            let v = inv.iter().collect::<Vec<_>>();
372
373            assert_eq!(
374                v,
375                &[
376                    13, 5, 30, 178, 86, 147, 139, 275, 48, 64, 224, 81, 97, 293, 198, 286, 230,
377                    243, 127, 167, 180, 208, 116, 261, 37, 17, 6, 34, 170, 87, 151, 143, 264, 52,
378                    68, 225, 73, 101, 294, 202, 278, 231, 247, 131, 156, 184, 212, 117, 253, 41,
379                    18, 10, 26, 171, 91, 155, 132, 268, 56, 69, 217, 77, 102, 298, 194, 279, 235,
380                    251, 120, 160, 188, 213, 109, 257, 42, 22, 2, 27, 175, 95, 144, 136, 272, 57,
381                    61, 221, 78, 106, 290, 195, 283, 239, 240, 124, 164, 189, 205, 113, 258, 46,
382                    14, 3, 31, 179, 84, 148, 140, 273, 49, 65, 222, 82, 98, 291, 199, 287, 228,
383                    244, 128, 165, 181, 209, 114, 262, 38, 15, 7, 35, 168, 88, 152, 141, 265, 53,
384                    66, 226, 74, 99, 295, 203, 276, 232, 248, 129, 157, 185, 210, 118, 254, 39, 19,
385                    11, 24, 172, 92, 153, 133, 269, 54, 70, 218, 75, 103, 299, 192, 280, 236, 249,
386                    121, 161, 186, 214, 110, 255, 43, 23, 0, 28, 176, 93, 145, 137, 270, 58, 62,
387                    219, 79, 107, 288, 196, 284, 237, 241, 125, 162, 190, 206, 111, 259, 47, 12, 4,
388                    32, 177, 85, 149, 138, 274, 50, 63, 223, 83, 96, 292, 200, 285, 229, 245, 126,
389                    166, 182, 207, 115, 263, 36, 16, 8, 33, 169, 89, 150, 142, 266, 51, 67, 227,
390                    72, 100, 296, 201, 277, 233, 246, 130, 158, 183, 211, 119, 252, 40, 20, 9, 25,
391                    173, 90, 154, 134, 267, 55, 71, 216, 76, 104, 297, 193, 281, 234, 250, 122,
392                    159, 187, 215, 108, 256, 44, 21, 1, 29, 174, 94, 146, 135, 271, 59, 60, 220,
393                    80, 105, 289, 197, 282, 238, 242, 123, 163, 191, 204, 112, 260, 45
394                ]
395            );
396        }
397
398        #[test]
399        fn test_nth_3() {
400            let mut rng = Xoshiro256StarStar::from_seed(SEED);
401            let p = RandomPermutation::with_rng((1..=20).product(), &mut rng).unwrap();
402            let inv = p.inverse();
403
404            for i in 0..1000 {
405                assert_eq!(p.nth(inv.nth(i).unwrap()), Some(i));
406            }
407        }
408    }
409
410    mod iterator {
411        use super::*;
412
413        #[test]
414        fn test_next() {
415            let mut rng = Xoshiro256StarStar::from_seed(SEED);
416            let p = RandomPermutation::with_rng(3113510400, &mut rng).unwrap();
417            let mut iter = p.iter();
418
419            for i in 0..1000 {
420                assert_eq!(iter.next(), p.nth(i));
421            }
422        }
423
424        #[test]
425        fn test_nth() {
426            let mut rng = Xoshiro256StarStar::from_seed(SEED);
427            let p = RandomPermutation::with_rng(3113510400, &mut rng).unwrap();
428            let mut iter = p.iter();
429
430            for i in 0..1000 {
431                assert_eq!(iter.nth(1000000), p.nth((i + 1) * 1000000));
432            }
433        }
434    }
435
436    mod composition {
437        use super::*;
438
439        #[test]
440        fn test_new_1() {
441            let comp = Composition::new(&[]);
442            assert!(comp.is_none());
443        }
444
445        #[test]
446        fn test_new_2() {
447            let mut rng = Xoshiro256StarStar::from_seed(SEED);
448            let p1 = RandomPermutation::with_rng(300, &mut rng).unwrap();
449            let p2 = RandomPermutation::with_rng(400, &mut rng).unwrap();
450
451            let v = vec![p1, p2];
452            let comp = Composition::new(&v);
453
454            assert!(comp.is_none());
455        }
456
457        #[test]
458        fn test_nth() {
459            let mut rng = Xoshiro256StarStar::from_seed(SEED);
460            let p1 = RandomPermutation::with_rng(300, &mut rng).unwrap();
461            let p2 = RandomPermutation::with_rng(300, &mut rng).unwrap();
462
463            let v = vec![p1.clone(), p2.clone()];
464            let comp = Composition::new(&v).unwrap();
465
466            for i in 0..300 {
467                assert_eq!(comp.nth(i), p2.nth(p1.nth(i).unwrap()));
468            }
469        }
470    }
471}