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