adele_ring/
basis.rs

1//! Adaptive multimodular basis.
2//!
3//! Where the original engine used a *fixed* 32-prime channel set (a fixed dynamic
4//! range that silently aliased on overflow), `Basis` provisions **as many primes
5//! as the computation's height demands** (see [`crate::bounds`]), then CRT +
6//! rational-reconstructs once at the boundary.
7//!
8//! All primes live in the open interval `(2^15, 2^16)`. That makes them
9//! GPU-eligible: the product of two residues is `< 2^32`, so the naive WGSL
10//! `(a*b) % m` is safe, and each channel contributes ~16 bits of range.
11
12use std::sync::{Arc, OnceLock};
13
14use num_bigint::{BigInt, BigUint};
15use num_traits::One;
16
17use crate::primes::primes_up_to;
18
19/// Lower bound (exclusive) for basis primes: `2^15`.
20pub const PRIME_MIN: u32 = 1 << 15;
21/// Upper bound (exclusive) for basis primes: `2^16`.
22pub const PRIME_MAX: u32 = 1 << 16;
23
24/// All GPU-eligible primes in `(2^15, 2^16)`, computed once and shared.
25fn gpu_primes() -> &'static Vec<u32> {
26    static PRIMES: OnceLock<Vec<u32>> = OnceLock::new();
27    PRIMES.get_or_init(|| {
28        primes_up_to((PRIME_MAX - 1) as usize)
29            .into_iter()
30            .map(|p| p as u32)
31            .filter(|&p| p > PRIME_MIN)
32            .collect()
33    })
34}
35
36/// An ordered, extensible set of pairwise-coprime moduli.
37#[derive(Clone, Debug)]
38pub struct Basis {
39    primes: Arc<Vec<u32>>,
40}
41
42impl Basis {
43    /// Build directly from a prime count (mainly for tests / fixed-size needs).
44    pub fn from_count(n: usize) -> Self {
45        let all = gpu_primes();
46        assert!(n <= all.len(), "requested more primes than are GPU-eligible");
47        Basis { primes: Arc::new(all[..n].to_vec()) }
48    }
49
50    /// A general-purpose default basis (~512 bits ≈ 32 channels). Useful for
51    /// open-ended Level 0/1 work and as the seed the multimodular pipeline grows.
52    pub fn standard() -> Self {
53        Self::with_bits(512)
54    }
55
56    /// The smallest basis whose product exceeds `2^bits`.
57    pub fn with_bits(bits: u64) -> Self {
58        let all = gpu_primes();
59        let mut product = BigUint::one();
60        let target = BigUint::one() << bits as usize;
61        let mut take = 0;
62        for &p in all.iter() {
63            if product > target {
64                break;
65            }
66            product *= BigUint::from(p);
67            take += 1;
68        }
69        Basis { primes: Arc::new(all[..take].to_vec()) }
70    }
71
72    /// Number of channels.
73    #[inline]
74    pub fn len(&self) -> usize {
75        self.primes.len()
76    }
77
78    /// Whether the basis is empty.
79    #[inline]
80    pub fn is_empty(&self) -> bool {
81        self.primes.is_empty()
82    }
83
84    /// The modulus of channel `i`.
85    #[inline]
86    pub fn modulus(&self, i: usize) -> u32 {
87        self.primes[i]
88    }
89
90    /// All moduli as a slice.
91    #[inline]
92    pub fn moduli(&self) -> &[u32] {
93        &self.primes
94    }
95
96    /// `floor(log2(∏ primes))` — the usable bit capacity.
97    pub fn capacity_bits(&self) -> u64 {
98        let product = self.modulus_product();
99        if product <= BigUint::one() {
100            0
101        } else {
102            product.bits() - 1
103        }
104    }
105
106    /// The product `M = ∏ primes`.
107    pub fn modulus_product(&self) -> BigUint {
108        self.primes.iter().map(|&p| BigUint::from(p)).product()
109    }
110
111    /// Alias for [`Basis::modulus_product`] — total dynamic range `M`.
112    pub fn capacity(&self) -> BigUint {
113        self.modulus_product()
114    }
115
116    /// Signed range bound `⌊M/2⌋`: balanced values live in `(-bound, bound]`.
117    pub fn signed_capacity(&self) -> BigInt {
118        BigInt::from(self.modulus_product() / BigUint::from(2u8))
119    }
120
121    /// A basis with at least `bits` of capacity (extends with more primes if
122    /// needed; the returned basis shares no state but is cheap to build).
123    pub fn extend_to_bits(&self, bits: u64) -> Basis {
124        if self.capacity_bits() >= bits {
125            return self.clone();
126        }
127        Basis::with_bits(bits)
128    }
129}
130
131impl PartialEq for Basis {
132    fn eq(&self, other: &Self) -> bool {
133        Arc::ptr_eq(&self.primes, &other.primes) || self.primes == other.primes
134    }
135}
136impl Eq for Basis {}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn primes_are_gpu_eligible() {
144        let b = Basis::from_count(10);
145        for i in 0..b.len() {
146            let p = b.modulus(i);
147            assert!(p > PRIME_MIN && p < PRIME_MAX);
148        }
149    }
150
151    #[test]
152    fn with_bits_exceeds_target() {
153        let b = Basis::with_bits(166); // ~10^50
154        assert!(b.capacity_bits() >= 166);
155        // each prime is ~16 bits, so ~11 channels
156        assert!(b.len() <= 12 && b.len() >= 10);
157    }
158
159    #[test]
160    fn extend_grows() {
161        let small = Basis::with_bits(50);
162        let big = small.extend_to_bits(500);
163        assert!(big.capacity_bits() >= 500);
164        assert!(big.len() > small.len());
165    }
166}
adele_ring/basis.rs

adele_ring/
basis.rs
