1pub mod poly;
  2use std::cmp::min;
  3use std::cmp::PartialEq;
  4
  5#[derive(Clone, Debug, Eq)]
  6pub struct Poly<O>(pub Vec<O>);
  7
  8impl<O> Poly<O>
  9where
 10    O: Default + Clone + Copy,
 11{
 12    pub fn new(n: usize) -> Self {
 13        Self(vec![O::default(); n])
 14    }
 15
 16    pub fn buffer_size(&self) -> usize {
 17        return self.0.len();
 18    }
 19
 20    pub fn from_buffer(&mut self, n: usize, buf: &mut [O]) {
 21        assert!(
 22            buf.len() >= n,
 23            "invalid buffer: buf.len()={} < n={}",
 24            buf.len(),
 25            n
 26        );
 27        self.0 = Vec::from(&buf[..n]);
 28    }
 29
 30    pub fn log_n(&self) -> usize {
 31        (usize::BITS - (self.n() - 1).leading_zeros()) as usize
 32    }
 33
 34    pub fn n(&self) -> usize {
 35        self.0.len()
 36    }
 37
 38    pub fn resize(&mut self, n: usize) {
 39        self.0.resize(n, O::default());
 40    }
 41
 42    pub fn set(&mut self, v: &[O]) {
 43        let size: usize = min(v.len(), self.n());
 44        self.0[..size].copy_from_slice(&v[..size]);
 45    }
 46
 47    pub fn fill(&mut self, v: &O) {
 48        self.0.fill(*v)
 49    }
 50
 51    pub fn zero(&mut self) {
 52        self.fill(&O::default())
 53    }
 54
 55    pub fn copy_from(&mut self, other: &Poly<O>) {
 56        if std::ptr::eq(self, other) {
 57            return;
 58        }
 59        self.resize(other.n());
 60        self.0.copy_from_slice(&other.0)
 61    }
 62}
 63
 64impl<O: PartialEq> PartialEq for Poly<O> {
 65    fn eq(&self, other: &Self) -> bool {
 66        std::ptr::eq(self, other) || (self.0 == other.0)
 67    }
 68}
 69
 70impl<O> Default for Poly<O> {
 71    fn default() -> Self {
 72        Poly(Vec::new())
 73    }
 74}
 75
 76#[derive(Clone, Debug, Eq)]
 77pub struct PolyRNS<O>(pub Vec<Poly<O>>);
 78
 79impl<O> PolyRNS<O>
 80where
 81    O: Default + Clone + Copy,
 82{
 83    pub fn new(n: usize, level: usize) -> Self {
 84        let mut polyrns: PolyRNS<O> = PolyRNS::<O>::default();
 85        let mut buf: Vec<O> = vec![O::default(); polyrns.buffer_size(n, level)];
 86        polyrns.from_buffer(n, level, &mut buf[..]);
 87        polyrns
 88    }
 89
 90    pub fn n(&self) -> usize {
 91        self.0[0].n()
 92    }
 93
 94    pub fn log_n(&self) -> usize {
 95        self.0[0].log_n()
 96    }
 97
 98    pub fn level(&self) -> usize {
 99        self.0.len() - 1
100    }
101
102    pub fn buffer_size(&self, n: usize, level: usize) -> usize {
103        n * (level + 1)
104    }
105
106    pub fn from_buffer(&mut self, n: usize, level: usize, buf: &mut [O]) {
107        assert!(
108            buf.len() >= n * (level + 1),
109            "invalid buffer: buf.len()={} < n * (level+1)={}",
110            buf.len(),
111            level + 1
112        );
113        self.0.clear();
114        for chunk in buf.chunks_mut(n).take(level + 1) {
115            let mut poly: Poly<O> = Poly(Vec::new());
116            poly.from_buffer(n, chunk);
117            self.0.push(poly);
118        }
119    }
120
121    pub fn resize(&mut self, level: usize) {
122        self.0.resize(level + 1, Poly::<O>::new(self.n()));
123    }
124
125    pub fn split_at_mut(&mut self, level: usize) -> (&mut [Poly<O>], &mut [Poly<O>]) {
126        self.0.split_at_mut(level)
127    }
128
129    pub fn at(&self, level: usize) -> &Poly<O> {
130        assert!(
131            level <= self.level(),
132            "invalid argument level: level={} > self.level()={}",
133            level,
134            self.level()
135        );
136        &self.0[level]
137    }
138
139    pub fn at_mut(&mut self, level: usize) -> &mut Poly<O> {
140        &mut self.0[level]
141    }
142
143    pub fn fill(&mut self, v: &O) {
144        (0..self.level() + 1).for_each(|i| self.at_mut(i).fill(v))
145    }
146
147    pub fn zero(&mut self) {
148        self.fill(&O::default())
149    }
150
151    pub fn copy(&mut self, other: &PolyRNS<O>) {
152        if std::ptr::eq(self, other) {
153            return;
154        }
155        self.resize(other.level());
156        self.copy_level(other.level(), other);
157    }
158
159    pub fn copy_level(&mut self, level: usize, other: &PolyRNS<O>) {
160        assert!(
161            self.level() <= level,
162            "invalid argument level: level={} > self.level()={}",
163            level,
164            self.level()
165        );
166        assert!(
167            other.level() <= level,
168            "invalid argument level: level={} > other.level()={}",
169            level,
170            other.level()
171        );
172        (0..level + 1).for_each(|i| self.at_mut(i).copy_from(other.at(i)))
173    }
174}
175
176impl<O: PartialEq> PartialEq for PolyRNS<O> {
177    fn eq(&self, other: &Self) -> bool {
178        std::ptr::eq(self, other) && (self.0 == other.0)
179    }
180}
181
182impl<O> Default for PolyRNS<O> {
183    fn default() -> Self {
184        Self { 0: Vec::new() }
185    }
186}