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