@@ -10,3 +10,4 @@ [workspace.dependencies]
 anyhow = "1.0.56"
 rand = "0.8.5"
 rand_distr = "0.4.3"
+itertools = "0.14.0"

@@ -1,4 +1,4 @@
 # fhe-study
-Code done while studying some FHE papers.
+Code done while studying some FHE papers, with the idea of doing implementations from scratch.
 
 - arithmetic: contains $\mathbb{Z}_q$ and $\mathbb{Z}_q[X]/(X^N+1)$ arithmetic implementations, together with the NTT implementation.

@@ -4,11 +4,13 @@ #![allow(non_camel_case_types)]
 #![allow(clippy::upper_case_acronyms)]
 #![allow(dead_code)] // TMP
 
-mod naive; // TODO rm
+mod naive_ntt; // TODO rm
 pub mod ntt;
 pub mod ring;
+pub mod ringq;
 pub mod zq;
 
 pub use ntt::NTT;
-pub use ring::PR;
+pub use ring::R;
+pub use ringq::Rq;
 pub use zq::Zq;

@@ -102,7 +102,7 @@ use super::*;
     use rand_distr::Uniform;
 
     use crate::ring::matrix_vec_product;
-    use crate::ring::PR;
+    use crate::ring::Rq;
 
     #[test]
     fn roots_of_unity() -> Result<()> {
@@ -136,7 +136,7 @@ let v_inv = NTT::<Q, N>::invert_vandermonde(&v);
 
         let mut rng = rand::thread_rng();
         let uniform_distr = Uniform::new(0_f64, Q as f64);
-        let a = PR::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
+        let a = Rq::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
         // let a = PR::<Q, N>::new_from_u64(vec![36, 21, 9, 19]);
 
         // let a_padded_coeffs: [Zq<Q>; 2 * N] =
@@ -148,7 +148,7 @@ let a_ntt = matrix_vec_product(&v, &a_padded)?;
         let a_intt: Vec<Zq<Q>> = matrix_vec_product(&v_inv, &a_ntt)?;
         assert_eq!(a_intt, a_padded);
         let a_intt_arr: [Zq<Q>; N] = std::array::from_fn(|i| a_intt[i]);
-        assert_eq!(PR::new(a_intt_arr, None), a);
+        assert_eq!(Rq::new(a_intt_arr, None), a);
 
         Ok(())
     }
@@ -162,7 +162,7 @@ let ntt = NTT::<Q, N>::new()?;
 
         let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0]
             .iter()
-            .map(|&e| Zq::new(e))
+            .map(|&e| Zq::from_u64(e))
             .collect();
         let a_ntt = matrix_vec_product(&ntt.ntt, &a)?;
         let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
@@ -181,7 +181,7 @@ // let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2*N) as u64)?;
         let ntt = NTT::<Q, N>::new()?;
 
         let rng = rand::thread_rng();
-        let a = PR::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
+        let a = Rq::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
         let a = a.coeffs;
         dbg!(&a);
         let a_ntt = matrix_vec_product(&ntt.ntt, &a.to_vec())?;

@@ -115,19 +115,20 @@ }
 
 /// returns x^k mod Q
 const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
-    let mut r = 1u64;
-    let mut x = x;
-    let mut k = k;
-    x = x % Q;
+    // work on u128 to avoid overflow
+    let mut r = 1u128;
+    let mut x = x as u128;
+    let mut k = k as u128;
+    x = x % Q as u128;
     // exponentiation by square strategy
     while k > 0 {
         if k % 2 == 1 {
-            r = (r * x) % Q;
+            r = (r * x) % Q as u128;
         }
-        x = (x * x) % Q;
+        x = (x * x) % Q as u128;
         k /= 2;
     }
-    r
+    r as u64
 }
 
 /// returns x^-1 mod Q
@@ -149,7 +150,7 @@ const Q: u64 = 2u64.pow(16) + 1;
         const N: usize = 4;
 
         let a: [u64; N] = [1u64, 2, 3, 4];
-        let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
+        let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
 
         let a_ntt = NTT::<Q, N>::ntt(a);
 

@@ -1,3 +1,7 @@
+//! Polynomial ring Z[X]/(X^N+1)
+//!
+
+use anyhow::{anyhow, Result};
 use rand::{distributions::Distribution, Rng};
 use std::array;
 use std::fmt;
@@ -5,432 +9,176 @@ use std::ops;
 
 use crate::ntt::NTT;
 use crate::zq::Zq;
-use anyhow::{anyhow, Result};
 
-// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
-#[derive(Clone, Copy)]
-pub struct PR<const Q: u64, const N: usize> {
-    pub(crate) coeffs: [Zq<Q>; N],
+// PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
+#[derive(Clone, Copy, Debug)]
+pub struct R<const N: usize>([i64; N]);
 
-    // evals are set when doig a PRxPR multiplication, so it can be reused in future
-    // multiplications avoiding recomputing it
-    pub(crate) evals: Option<[Zq<Q>; N]>,
+impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {
+    fn from(rq: crate::ringq::Rq<Q, N>) -> Self {
+        Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
+    }
 }
 
-// TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just
-// use the trait and not need to add '<Q, N>' to their params
-
-// apply mod (X^N+1)
-pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
-    if p.len() < N {
-        return;
+impl<const N: usize> R<N> {
+    pub fn coeffs(&self) -> [i64; N] {
+        self.0
     }
-    for i in N..p.len() {
-        p[i - N] = p[i - N].clone() - p[i].clone();
-        p[i] = Zq(0);
+    pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
+        crate::Rq::<Q, N>::from(self)
     }
-    p.truncate(N);
-}
 
-// PR stands for PolynomialRing
-impl<const Q: u64, const N: usize> PR<Q, N> {
-    pub fn coeffs(&self) -> [Zq<Q>; N] {
-        self.coeffs
-    }
-
-    pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
+    pub fn from_vec(coeffs: Vec<i64>) -> Self {
         let mut p = coeffs;
-        modulus::<Q, N>(&mut p);
-        let coeffs = array::from_fn(|i| p[i]);
-        Self {
-            coeffs,
-            evals: None,
-        }
+        modulus::<N>(&mut p);
+        Self(array::from_fn(|i| p[i]))
     }
     // this method is mostly for tests
     pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
-        let coeffs_mod_q = coeffs.iter().map(|c| Zq::new(*c)).collect();
-        Self::from_vec(coeffs_mod_q)
+        let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
+        Self::from_vec(coeffs_i64)
     }
-    pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
-        Self { coeffs, evals }
+    pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
+        let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
+        Self::from_vec(coeffs_i64)
     }
-
-    pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
-        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
-        Ok(Self {
-            coeffs,
-            evals: None,
-        })
-    }
-    pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
-        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
-        Ok(Self {
-            coeffs,
-            evals: None,
-        })
-    }
-    pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Result<Self> {
-        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::new(dist.sample(&mut rng)));
-        Ok(Self {
-            coeffs,
-            evals: None,
-        })
-    }
-    // WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
-    pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
-        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
-        Ok(PR {
-            coeffs,
-            evals: None,
-        })
-    }
-    // Warning: this method assumes Q < P
-    pub fn remodule<const P: u64>(&self) -> PR<P, N> {
-        assert!(Q < P);
-        PR::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
-    }
-
-    // TODO review if needed, or if with this interface
-    pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
-        matrix_vec_product(m, &self.coeffs.to_vec())
-    }
-    pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
-        Self {
-            coeffs: array::from_fn(|i| self.coeffs[i] * *s),
-            evals: None,
-        }
-    }
-    pub fn mul_by_u64(&self, s: u64) -> Self {
-        let s = Zq::new(s);
-        Self {
-            coeffs: array::from_fn(|i| self.coeffs[i] * s),
-            // coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
-            evals: None,
-        }
-    }
-    pub fn mul_by_f64(&self, s: f64) -> Self {
-        Self {
-            coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
-            evals: None,
-        }
+    pub fn new(coeffs: [i64; N]) -> Self {
+        Self(coeffs)
     }
-
-    pub fn mul(&mut self, rhs: &mut Self) -> Self {
-        mul_mut(self, rhs)
+    pub fn mul_by_i64(&self, s: i64) -> Self {
+        Self(array::from_fn(|i| self.0[i] * s))
     }
-
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        // TODO simplify
-        let mut str = "";
-        let mut zero = true;
-        for (i, coeff) in self.coeffs.iter().enumerate().rev() {
-            if coeff.0 == 0 {
-                continue;
-            }
-            zero = false;
-            f.write_str(str)?;
-            if coeff.0 != 1 {
-                f.write_str(coeff.0.to_string().as_str())?;
-                if i > 0 {
-                    f.write_str("*")?;
-                }
-            }
-            if coeff.0 == 1 && i == 0 {
-                f.write_str(coeff.0.to_string().as_str())?;
-            }
-            if i == 1 {
-                f.write_str("x")?;
-            } else if i > 1 {
-                f.write_str("x^")?;
-                f.write_str(i.to_string().as_str())?;
-            }
-            str = " + ";
-        }
-        if zero {
-            f.write_str("0")?;
-        }
-
-        f.write_str(" mod Z_")?;
-        f.write_str(Q.to_string().as_str())?;
-        f.write_str("/(X^")?;
-        f.write_str(N.to_string().as_str())?;
-        f.write_str("+1)")?;
-        Ok(())
+    // performs the multiplication and division over f64, and then it rounds the
+    // result, only applying the mod Q at the end
+    pub fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
+        let r: Vec<f64> = self
+            .coeffs()
+            .iter()
+            .map(|e| ((num as f64 * *e as f64) / den as f64).round())
+            .collect();
+        crate::Rq::<Q, N>::from_vec_f64(r)
     }
 }
-pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
-    // assert_eq!(m.len(), m[0].len()); // TODO change to returning err
-    // assert_eq!(m.len(), v.len());
-    if m.len() != m[0].len() {
-        return Err(anyhow!("expected 'm' to be a square matrix"));
-    }
-    if m.len() != v.len() {
-        return Err(anyhow!(
-            "m.len: {} should be equal to v.len(): {}",
-            m.len(),
-            v.len(),
-        ));
-    }
-
-    Ok(m.iter()
-        .map(|row| {
-            row.iter()
-                .zip(v.iter())
-                .map(|(&row_i, &v_i)| row_i * v_i)
-                .sum()
-        })
-        .collect::<Vec<Zq<Q>>>())
+pub fn mul_div_round<const Q: u64, const N: usize>(
+    v: Vec<i64>,
+    num: u64,
+    den: u64,
+) -> crate::Rq<Q, N> {
+    // dbg!(&v);
+    let r: Vec<f64> = v
+        .iter()
+        .map(|e| ((num as f64 * *e as f64) / den as f64).round())
+        .collect();
+    // dbg!(&r);
+    crate::Rq::<Q, N>::from_vec_f64(r)
 }
-pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
-    // TODO case when m[0].len()=0
-    // TODO non square matrix
-    let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
-    for (i, m_row) in m.iter().enumerate() {
-        for (j, m_ij) in m_row.iter().enumerate() {
-            r[j][i] = *m_ij;
-        }
+
+// apply mod (X^N+1)
+pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
+    if p.len() < N {
+        return;
+    }
+    for i in N..p.len() {
+        p[i - N] = p[i - N].clone() - p[i].clone();
+        p[i] = 0;
     }
-    r
+    p.truncate(N);
 }
 
-impl<const Q: u64, const N: usize> PartialEq for PR<Q, N> {
+impl<const N: usize> PartialEq for R<N> {
     fn eq(&self, other: &Self) -> bool {
-        self.coeffs == other.coeffs
+        self.0 == other.0
     }
 }
-impl<const Q: u64, const N: usize> ops::Add<PR<Q, N>> for PR<Q, N> {
+impl<const N: usize> ops::Add<R<N>> for R<N> {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self {
-        Self {
-            coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
-            evals: None,
-        }
-        // Self {
-        //     coeffs: self
-        //         .coeffs
-        //         .iter()
-        //         .zip(rhs.coeffs)
-        //         .map(|(a, b)| *a + b)
-        //         .collect(),
-        //     evals: None,
-        // }
-        // Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
+        Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
     }
 }
-impl<const Q: u64, const N: usize> ops::Add<&PR<Q, N>> for &PR<Q, N> {
-    type Output = PR<Q, N>;
+impl<const N: usize> ops::Add<&R<N>> for &R<N> {
+    type Output = R<N>;
 
-    fn add(self, rhs: &PR<Q, N>) -> Self::Output {
-        PR {
-            coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
-            evals: None,
-        }
+    fn add(self, rhs: &R<N>) -> Self::Output {
+        R(array::from_fn(|i| self.0[i] + rhs.0[i]))
     }
 }
-impl<const Q: u64, const N: usize> ops::Sub<PR<Q, N>> for PR<Q, N> {
+impl<const N: usize> ops::Sub<R<N>> for R<N> {
     type Output = Self;
 
     fn sub(self, rhs: Self) -> Self {
-        Self {
-            coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
-            evals: None,
-        }
+        Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
     }
 }
-impl<const Q: u64, const N: usize> ops::Sub<&PR<Q, N>> for &PR<Q, N> {
-    type Output = PR<Q, N>;
+impl<const N: usize> ops::Sub<&R<N>> for &R<N> {
+    type Output = R<N>;
 
-    fn sub(self, rhs: &PR<Q, N>) -> Self::Output {
-        PR {
-            coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
-            evals: None,
-        }
+    fn sub(self, rhs: &R<N>) -> Self::Output {
+        R(array::from_fn(|i| self.0[i] - rhs.0[i]))
     }
 }
-impl<const Q: u64, const N: usize> ops::Mul<PR<Q, N>> for PR<Q, N> {
+impl<const N: usize> ops::Mul<R<N>> for R<N> {
     type Output = Self;
 
     fn mul(self, rhs: Self) -> Self {
-        mul(&self, &rhs)
+        naive_poly_mul(&self, &rhs)
     }
 }
-impl<const Q: u64, const N: usize> ops::Mul<&PR<Q, N>> for &PR<Q, N> {
-    type Output = PR<Q, N>;
+impl<const N: usize> ops::Mul<&R<N>> for &R<N> {
+    type Output = R<N>;
 
-    fn mul(self, rhs: &PR<Q, N>) -> Self::Output {
-        mul(self, rhs)
+    fn mul(self, rhs: &R<N>) -> Self::Output {
+        naive_poly_mul(self, rhs)
     }
 }
 
-// mul by Zq element
-impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for PR<Q, N> {
-    type Output = Self;
+// TODO with NTT(?)
+pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
+    let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
+    let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
+    let mut result: Vec<i128> = vec![0; (N * 2) - 1];
+    for i in 0..N {
+        for j in 0..N {
+            result[i + j] = result[i + j] + poly1[i] * poly2[j];
+        }
+    }
 
-    fn mul(self, s: Zq<Q>) -> Self {
-        self.mul_by_zq(&s)
+    // apply mod (X^N + 1))
+    R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
+}
+pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
+    let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
+    let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
+    let mut result = vec![0; (N * 2) - 1];
+    for i in 0..N {
+        for j in 0..N {
+            result[i + j] = result[i + j] + poly1[i] * poly2[j];
+        }
     }
+    result.iter().map(|c| *c as i64).collect()
 }
-impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &PR<Q, N> {
-    type Output = PR<Q, N>;
 
-    fn mul(self, s: &Zq<Q>) -> Self::Output {
-        self.mul_by_zq(s)
-    }
-}
 // mul by u64
-impl<const Q: u64, const N: usize> ops::Mul<u64> for PR<Q, N> {
+impl<const N: usize> ops::Mul<u64> for R<N> {
     type Output = Self;
 
     fn mul(self, s: u64) -> Self {
-        self.mul_by_u64(s)
+        self.mul_by_i64(s as i64)
     }
 }
-impl<const Q: u64, const N: usize> ops::Mul<&u64> for &PR<Q, N> {
-    type Output = PR<Q, N>;
+impl<const N: usize> ops::Mul<&u64> for &R<N> {
+    type Output = R<N>;
 
     fn mul(self, s: &u64) -> Self::Output {
-        self.mul_by_u64(*s)
+        self.mul_by_i64(*s as i64)
     }
 }
 
-impl<const Q: u64, const N: usize> ops::Neg for PR<Q, N> {
+impl<const N: usize> ops::Neg for R<N> {
     type Output = Self;
 
     fn neg(self) -> Self::Output {
-        Self {
-            coeffs: array::from_fn(|i| -self.coeffs[i]),
-            evals: None,
-        }
-    }
-}
-
-fn mul_mut<const Q: u64, const N: usize>(lhs: &mut PR<Q, N>, rhs: &mut PR<Q, N>) -> PR<Q, N> {
-    // reuse evaluations if already computed
-    if !lhs.evals.is_some() {
-        lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
-    };
-    if !rhs.evals.is_some() {
-        rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
-    };
-    let lhs_evals = lhs.evals.unwrap();
-    let rhs_evals = rhs.evals.unwrap();
-
-    let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
-    let c = NTT::<Q, { N }>::intt(c_ntt);
-    PR::new(c, Some(c_ntt))
-}
-fn mul<const Q: u64, const N: usize>(lhs: &PR<Q, N>, rhs: &PR<Q, N>) -> PR<Q, N> {
-    // reuse evaluations if already computed
-    let lhs_evals = if lhs.evals.is_some() {
-        lhs.evals.unwrap()
-    } else {
-        NTT::<Q, N>::ntt(lhs.coeffs)
-    };
-    let rhs_evals = if rhs.evals.is_some() {
-        rhs.evals.unwrap()
-    } else {
-        NTT::<Q, N>::ntt(rhs.coeffs)
-    };
-
-    let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
-    let c = NTT::<Q, { N }>::intt(c_ntt);
-    PR::new(c, Some(c_ntt))
-}
-
-impl<const Q: u64, const N: usize> fmt::Display for PR<Q, N> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        self.fmt(f)?;
-        Ok(())
-    }
-}
-impl<const Q: u64, const N: usize> fmt::Debug for PR<Q, N> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        self.fmt(f)?;
-        Ok(())
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn poly_ring() {
-        // the test values used are generated with SageMath
-        const Q: u64 = 7;
-        const N: usize = 3;
-
-        // p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
-        let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
-        assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
-
-        // try with coefficients bigger than Q
-        let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
-        assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
-
-        // try with other ring
-        let p = PR::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
-        assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
-
-        let p = PR::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
-        assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
-
-        let p = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
-        assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
-
-        let a = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
-        assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
-
-        let b = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
-        assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
-
-        // add
-        assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
-        assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
-        // assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
-
-        // sub
-        assert_eq!(
-            (a.clone() - b.clone()).to_string(),
-            "x^2 + x + 1 mod Z_7/(X^3+1)"
-        );
-    }
-
-    fn test_mul_opt<const Q: u64, const N: usize>(
-        a: [u64; N],
-        b: [u64; N],
-        expected_c: [u64; N],
-    ) -> Result<()> {
-        let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
-        let mut a = PR::new(a, None);
-        let b: [Zq<Q>; N] = array::from_fn(|i| Zq::new(b[i]));
-        let mut b = PR::new(b, None);
-        let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::new(expected_c[i]));
-        let expected_c = PR::new(expected_c, None);
-
-        let c = mul_mut(&mut a, &mut b);
-        assert_eq!(c, expected_c);
-        Ok(())
-    }
-    #[test]
-    fn test_mul() -> Result<()> {
-        const Q: u64 = 2u64.pow(16) + 1;
-        const N: usize = 4;
-
-        let a: [u64; N] = [1u64, 2, 3, 4];
-        let b: [u64; N] = [1u64, 2, 3, 4];
-        let c: [u64; N] = [65513, 65517, 65531, 20];
-        test_mul_opt::<Q, N>(a, b, c)?;
-
-        let a: [u64; N] = [0u64, 0, 0, 2];
-        let b: [u64; N] = [0u64, 0, 0, 2];
-        let c: [u64; N] = [0u64, 0, 65533, 0];
-        test_mul_opt::<Q, N>(a, b, c)?;
-
-        // TODO more testvectors
-
-        Ok(())
+        Self(array::from_fn(|i| -self.0[i]))
     }
 }

@@ -0,0 +1,504 @@
+//! Polynomial ring Z_q[X]/(X^N+1)
+//!
+
+use rand::{distributions::Distribution, Rng};
+use std::array;
+use std::fmt;
+use std::ops;
+
+use crate::ntt::NTT;
+use crate::zq::{modulus_u64, Zq};
+use anyhow::{anyhow, Result};
+
+/// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
+/// The implementation assumes that q is prime.
+#[derive(Clone, Copy)]
+pub struct Rq<const Q: u64, const N: usize> {
+    pub(crate) coeffs: [Zq<Q>; N],
+
+    // evals are set when doig a PRxPR multiplication, so it can be reused in future
+    // multiplications avoiding recomputing it
+    pub(crate) evals: Option<[Zq<Q>; N]>,
+}
+
+// TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just
+// use the trait and not need to add '<Q, N>' to their params
+
+impl<const Q: u64, const N: usize> From<crate::ring::R<N>> for Rq<Q, N> {
+    fn from(r: crate::ring::R<N>) -> Self {
+        Self::from_vec(
+            r.coeffs()
+                .iter()
+                .map(|e| Zq::<Q>::from_f64(*e as f64))
+                .collect(),
+        )
+    }
+}
+
+// apply mod (X^N+1)
+pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
+    if p.len() < N {
+        return;
+    }
+    for i in N..p.len() {
+        p[i - N] = p[i - N].clone() - p[i].clone();
+        p[i] = Zq(0);
+    }
+    p.truncate(N);
+}
+
+// PR stands for PolynomialRing
+impl<const Q: u64, const N: usize> Rq<Q, N> {
+    pub fn coeffs(&self) -> [Zq<Q>; N] {
+        self.coeffs
+    }
+    pub fn to_r(self) -> crate::R<N> {
+        crate::R::<N>::from(self)
+    }
+
+    pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
+        let mut p = coeffs;
+        modulus::<Q, N>(&mut p);
+        let coeffs = array::from_fn(|i| p[i]);
+        Self {
+            coeffs,
+            evals: None,
+        }
+    }
+    // this method is mostly for tests
+    pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
+        let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_u64(*c)).collect();
+        Self::from_vec(coeffs_mod_q)
+    }
+    pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
+        let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c)).collect();
+        Self::from_vec(coeffs_mod_q)
+    }
+    pub fn from_vec_i64(coeffs: Vec<i64>) -> Self {
+        let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c as f64)).collect();
+        Self::from_vec(coeffs_mod_q)
+    }
+    pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
+        Self { coeffs, evals }
+    }
+
+    pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
+        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
+        Ok(Self {
+            coeffs,
+            evals: None,
+        })
+    }
+    pub fn rand_f64_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
+        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
+        Ok(Self {
+            coeffs,
+            evals: None,
+        })
+    }
+    pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
+        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
+        Ok(Self {
+            coeffs,
+            evals: None,
+        })
+    }
+    pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Result<Self> {
+        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
+        Ok(Self {
+            coeffs,
+            evals: None,
+        })
+    }
+    // WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
+    pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
+        let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
+        Ok(Rq {
+            coeffs,
+            evals: None,
+        })
+    }
+    // Warning: this method will behave differently depending on the values P and Q:
+    // if Q<P, it just 'renames' the modulus parameter to P
+    // if Q>=P, it crops to mod P
+    pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
+        Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
+    }
+    // applies mod(T) to all coefficients of self
+    pub fn coeffs_mod<const T: u64>(&self) -> Self {
+        Rq::<Q, N>::from_vec_u64(
+            self.coeffs()
+                .iter()
+                .map(|m_i| modulus_u64::<T>(m_i.0))
+                .collect(),
+        )
+    }
+
+    // TODO review if needed, or if with this interface
+    pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
+        matrix_vec_product(m, &self.coeffs.to_vec())
+    }
+    pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
+        Self {
+            coeffs: array::from_fn(|i| self.coeffs[i] * *s),
+            evals: None,
+        }
+    }
+    pub fn mul_by_u64(&self, s: u64) -> Self {
+        let s = Zq::from_u64(s);
+        Self {
+            coeffs: array::from_fn(|i| self.coeffs[i] * s),
+            // coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
+            evals: None,
+        }
+    }
+    pub fn mul_by_f64(&self, s: f64) -> Self {
+        Self {
+            coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
+            evals: None,
+        }
+    }
+
+    pub fn mul(&mut self, rhs: &mut Self) -> Self {
+        mul_mut(self, rhs)
+    }
+    // divides by the given scalar 's' and rounds, returning a Rq<Q,N>
+    // TODO rm
+    pub fn div_round(&self, s: u64) -> Self {
+        let r: Vec<f64> = self
+            .coeffs()
+            .iter()
+            .map(|e| (e.0 as f64 / s as f64).round())
+            .collect();
+        Rq::<Q, N>::from_vec_f64(r)
+    }
+    // returns [ [(num/den) * self].round() ] mod q
+    // ie. performs the multiplication and division over f64, and then it rounds the
+    // result, only applying the mod Q at the end
+    pub fn mul_div_round(&self, num: u64, den: u64) -> Self {
+        let r: Vec<f64> = self
+            .coeffs()
+            .iter()
+            .map(|e| ((num as f64 * e.0 as f64) / den as f64).round())
+            .collect();
+        Rq::<Q, N>::from_vec_f64(r)
+    }
+
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        // TODO simplify
+        let mut str = "";
+        let mut zero = true;
+        for (i, coeff) in self.coeffs.iter().enumerate().rev() {
+            if coeff.0 == 0 {
+                continue;
+            }
+            zero = false;
+            f.write_str(str)?;
+            if coeff.0 != 1 {
+                f.write_str(coeff.0.to_string().as_str())?;
+                if i > 0 {
+                    f.write_str("*")?;
+                }
+            }
+            if coeff.0 == 1 && i == 0 {
+                f.write_str(coeff.0.to_string().as_str())?;
+            }
+            if i == 1 {
+                f.write_str("x")?;
+            } else if i > 1 {
+                f.write_str("x^")?;
+                f.write_str(i.to_string().as_str())?;
+            }
+            str = " + ";
+        }
+        if zero {
+            f.write_str("0")?;
+        }
+
+        f.write_str(" mod Z_")?;
+        f.write_str(Q.to_string().as_str())?;
+        f.write_str("/(X^")?;
+        f.write_str(N.to_string().as_str())?;
+        f.write_str("+1)")?;
+        Ok(())
+    }
+
+    pub fn infinity_norm(&self) -> u64 {
+        self.coeffs().iter().map(|x| x.0).fold(0, |a, b| a.max(b))
+    }
+}
+pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
+    // assert_eq!(m.len(), m[0].len()); // TODO change to returning err
+    // assert_eq!(m.len(), v.len());
+    if m.len() != m[0].len() {
+        return Err(anyhow!("expected 'm' to be a square matrix"));
+    }
+    if m.len() != v.len() {
+        return Err(anyhow!(
+            "m.len: {} should be equal to v.len(): {}",
+            m.len(),
+            v.len(),
+        ));
+    }
+
+    Ok(m.iter()
+        .map(|row| {
+            row.iter()
+                .zip(v.iter())
+                .map(|(&row_i, &v_i)| row_i * v_i)
+                .sum()
+        })
+        .collect::<Vec<Zq<Q>>>())
+}
+pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
+    // TODO case when m[0].len()=0
+    // TODO non square matrix
+    let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
+    for (i, m_row) in m.iter().enumerate() {
+        for (j, m_ij) in m_row.iter().enumerate() {
+            r[j][i] = *m_ij;
+        }
+    }
+    r
+}
+
+impl<const Q: u64, const N: usize> PartialEq for Rq<Q, N> {
+    fn eq(&self, other: &Self) -> bool {
+        self.coeffs == other.coeffs
+    }
+}
+impl<const Q: u64, const N: usize> ops::Add<Rq<Q, N>> for Rq<Q, N> {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self {
+        Self {
+            coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
+            evals: None,
+        }
+        // Self {
+        //     coeffs: self
+        //         .coeffs
+        //         .iter()
+        //         .zip(rhs.coeffs)
+        //         .map(|(a, b)| *a + b)
+        //         .collect(),
+        //     evals: None,
+        // }
+        // Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
+    }
+}
+impl<const Q: u64, const N: usize> ops::Add<&Rq<Q, N>> for &Rq<Q, N> {
+    type Output = Rq<Q, N>;
+
+    fn add(self, rhs: &Rq<Q, N>) -> Self::Output {
+        Rq {
+            coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
+            evals: None,
+        }
+    }
+}
+impl<const Q: u64, const N: usize> ops::Sub<Rq<Q, N>> for Rq<Q, N> {
+    type Output = Self;
+
+    fn sub(self, rhs: Self) -> Self {
+        Self {
+            coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
+            evals: None,
+        }
+    }
+}
+impl<const Q: u64, const N: usize> ops::Sub<&Rq<Q, N>> for &Rq<Q, N> {
+    type Output = Rq<Q, N>;
+
+    fn sub(self, rhs: &Rq<Q, N>) -> Self::Output {
+        Rq {
+            coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
+            evals: None,
+        }
+    }
+}
+impl<const Q: u64, const N: usize> ops::Mul<Rq<Q, N>> for Rq<Q, N> {
+    type Output = Self;
+
+    fn mul(self, rhs: Self) -> Self {
+        mul(&self, &rhs)
+    }
+}
+impl<const Q: u64, const N: usize> ops::Mul<&Rq<Q, N>> for &Rq<Q, N> {
+    type Output = Rq<Q, N>;
+
+    fn mul(self, rhs: &Rq<Q, N>) -> Self::Output {
+        mul(self, rhs)
+    }
+}
+
+// mul by Zq element
+impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for Rq<Q, N> {
+    type Output = Self;
+
+    fn mul(self, s: Zq<Q>) -> Self {
+        self.mul_by_zq(&s)
+    }
+}
+impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &Rq<Q, N> {
+    type Output = Rq<Q, N>;
+
+    fn mul(self, s: &Zq<Q>) -> Self::Output {
+        self.mul_by_zq(s)
+    }
+}
+// mul by u64
+impl<const Q: u64, const N: usize> ops::Mul<u64> for Rq<Q, N> {
+    type Output = Self;
+
+    fn mul(self, s: u64) -> Self {
+        self.mul_by_u64(s)
+    }
+}
+impl<const Q: u64, const N: usize> ops::Mul<&u64> for &Rq<Q, N> {
+    type Output = Rq<Q, N>;
+
+    fn mul(self, s: &u64) -> Self::Output {
+        self.mul_by_u64(*s)
+    }
+}
+
+impl<const Q: u64, const N: usize> ops::Neg for Rq<Q, N> {
+    type Output = Self;
+
+    fn neg(self) -> Self::Output {
+        Self {
+            coeffs: array::from_fn(|i| -self.coeffs[i]),
+            evals: None,
+        }
+    }
+}
+
+fn mul_mut<const Q: u64, const N: usize>(lhs: &mut Rq<Q, N>, rhs: &mut Rq<Q, N>) -> Rq<Q, N> {
+    // reuse evaluations if already computed
+    if !lhs.evals.is_some() {
+        lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
+    };
+    if !rhs.evals.is_some() {
+        rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
+    };
+    let lhs_evals = lhs.evals.unwrap();
+    let rhs_evals = rhs.evals.unwrap();
+
+    let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
+    let c = NTT::<Q, { N }>::intt(c_ntt);
+    Rq::new(c, Some(c_ntt))
+}
+fn mul<const Q: u64, const N: usize>(lhs: &Rq<Q, N>, rhs: &Rq<Q, N>) -> Rq<Q, N> {
+    // reuse evaluations if already computed
+    let lhs_evals = if lhs.evals.is_some() {
+        lhs.evals.unwrap()
+    } else {
+        NTT::<Q, N>::ntt(lhs.coeffs)
+    };
+    let rhs_evals = if rhs.evals.is_some() {
+        rhs.evals.unwrap()
+    } else {
+        NTT::<Q, N>::ntt(rhs.coeffs)
+    };
+
+    let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
+    let c = NTT::<Q, { N }>::intt(c_ntt);
+    Rq::new(c, Some(c_ntt))
+}
+
+impl<const Q: u64, const N: usize> fmt::Display for Rq<Q, N> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.fmt(f)?;
+        Ok(())
+    }
+}
+impl<const Q: u64, const N: usize> fmt::Debug for Rq<Q, N> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        self.fmt(f)?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn poly_ring() {
+        // the test values used are generated with SageMath
+        const Q: u64 = 7;
+        const N: usize = 3;
+
+        // p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
+        let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
+        assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
+
+        // try with coefficients bigger than Q
+        let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
+        assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
+
+        // try with other ring
+        let p = Rq::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
+        assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
+
+        let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
+        assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
+
+        let p = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
+        assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
+
+        let a = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
+        assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
+
+        let b = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
+        assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
+
+        // add
+        assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
+        assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
+        // assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
+
+        // sub
+        assert_eq!(
+            (a.clone() - b.clone()).to_string(),
+            "x^2 + x + 1 mod Z_7/(X^3+1)"
+        );
+    }
+
+    fn test_mul_opt<const Q: u64, const N: usize>(
+        a: [u64; N],
+        b: [u64; N],
+        expected_c: [u64; N],
+    ) -> Result<()> {
+        let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
+        let mut a = Rq::new(a, None);
+        let b: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(b[i]));
+        let mut b = Rq::new(b, None);
+        let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(expected_c[i]));
+        let expected_c = Rq::new(expected_c, None);
+
+        let c = mul_mut(&mut a, &mut b);
+        assert_eq!(c, expected_c);
+        Ok(())
+    }
+    #[test]
+    fn test_mul() -> Result<()> {
+        const Q: u64 = 2u64.pow(16) + 1;
+        const N: usize = 4;
+
+        let a: [u64; N] = [1u64, 2, 3, 4];
+        let b: [u64; N] = [1u64, 2, 3, 4];
+        let c: [u64; N] = [65513, 65517, 65531, 20];
+        test_mul_opt::<Q, N>(a, b, c)?;
+
+        let a: [u64; N] = [0u64, 0, 0, 2];
+        let b: [u64; N] = [0u64, 0, 0, 2];
+        let c: [u64; N] = [0u64, 0, 65533, 0];
+        test_mul_opt::<Q, N>(a, b, c)?;
+
+        // TODO more testvectors
+
+        Ok(())
+    }
+}

@@ -12,22 +12,36 @@ //         v.into_iter().map(Zq::new).collect()
 //     }
 // }
 
+pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {
+    (e % Q + Q) % Q
+}
 impl<const Q: u64> Zq<Q> {
-    pub fn new(e: u64) -> Self {
+    pub fn from_u64(e: u64) -> Self {
         if e >= Q {
-            return Zq(e % Q);
+            // (e % Q + Q) % Q
+            return Zq(modulus_u64::<Q>(e));
+            // return Zq(e % Q);
         }
         Zq(e)
     }
     pub fn from_f64(e: f64) -> Self {
         // WIP method
         let e: i64 = e.round() as i64;
-        if e < 0 {
-            return Zq((Q as i64 + e) as u64);
-        } else if e >= Q as i64 {
-            return Zq((e % Q as i64) as u64);
+        let q = Q as i64;
+        if e < 0 || e >= q {
+            return Zq(((e % q + q) % q) as u64);
         }
         Zq(e as u64)
+
+        // if e < 0 {
+        //     // dbg!(&e);
+        //     // dbg!(Zq::<Q>(((Q as i64 + e) % Q as i64) as u64));
+        //     // return Zq(((Q as i64 + e) % Q as i64) as u64);
+        //     return Zq(e as u64 % Q);
+        // } else if e >= Q as i64 {
+        //     return Zq((e % Q as i64) as u64);
+        // }
+        // Zq(e as u64)
     }
     pub fn from_bool(b: bool) -> Self {
         if b {
@@ -83,7 +97,7 @@ }
         // if t < 0 {
         //     t = t + q;
         // }
-        return Zq::new(t);
+        return Zq::from_u64(t);
     }
     pub fn inv(self) -> Zq<Q> {
         let (g, x, _) = Self::egcd(self.0 as i128, Q as i128);

@@ -7,5 +7,6 @@ [dependencies]
 anyhow = { workspace = true }
 rand = { workspace = true }
 rand_distr = { workspace = true }
+itertools = { workspace = true }
 
 arithmetic = { path="../arithmetic" }

@@ -10,26 +10,100 @@ use rand::Rng;
 use rand_distr::{Normal, Uniform};
 use std::ops;
 
-use arithmetic::{Zq, PR};
+use arithmetic::{Rq, Zq, R};
 
 // error deviation for the Gaussian(Normal) distribution
 // sigma=3.2 from: https://eprint.iacr.org/2022/162.pdf page 5
 const ERR_SIGMA: f64 = 3.2;
+// const ERR_SIGMA: f64 = 0.0; // TODO WIP
 
 #[derive(Clone, Debug)]
-pub struct SecretKey<const Q: u64, const N: usize>(PR<Q, N>);
+pub struct SecretKey<const Q: u64, const N: usize>(Rq<Q, N>);
 
 #[derive(Clone, Debug)]
-pub struct PublicKey<const Q: u64, const N: usize>(PR<Q, N>, PR<Q, N>);
+pub struct PublicKey<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
+
+/// Relinearization key
+#[derive(Clone, Debug)]
+pub struct RLK<const PQ: u64, const N: usize>(Rq<PQ, N>, Rq<PQ, N>);
+
+// impl<const PQ: u64, const N: usize> RLK<Q, PQ, N> {
+//     // const P: u64 = PQ / Q;
+//
+//     // const PQ: u64 = P * Q;
+// }
 
 // RLWE ciphertext
 #[derive(Clone, Debug)]
-pub struct RLWE<const Q: u64, const N: usize>(PR<Q, N>, PR<Q, N>);
+pub struct RLWE<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
 
 impl<const Q: u64, const N: usize> RLWE<Q, N> {
     fn add(lhs: Self, rhs: Self) -> Self {
         RLWE::<Q, N>(lhs.0 + rhs.0, lhs.1 + rhs.1)
     }
+    pub fn remodule<const P: u64>(&self) -> RLWE<P, N> {
+        let x = self.0.remodule::<P>();
+        let y = self.1.remodule::<P>();
+        RLWE::<P, N>(x, y)
+    }
+    fn tensor<const PQ: u64, const T: u64>(a: &Self, b: &Self) -> (Rq<Q, N>, Rq<Q, N>, Rq<Q, N>) {
+        // expand Q->PQ // TODO rm
+        // get the coefficients in Z, ie. interpret a,b \in R (instead of R_q)
+        let a0: R<N> = a.0.to_r();
+        let a1: R<N> = a.1.to_r();
+        let b0: R<N> = b.0.to_r();
+        let b1: R<N> = b.1.to_r();
+
+        // tensor (\in R)
+        use arithmetic::ring::naive_mul;
+        let c0: Vec<i64> = naive_mul(&a0, &b0);
+        let c1_l: Vec<i64> = naive_mul(&a0, &b1);
+        let c1_r = naive_mul(&a1, &b0);
+        let c1: Vec<i64> = itertools::zip_eq(c1_l, c1_r).map(|(l, r)| l + r).collect();
+        let c2: Vec<i64> = naive_mul(&a1, &b1);
+
+        // scale down, module Q, so result is \in R_q
+        let c0: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c0, T, Q);
+        let c1: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c1, T, Q);
+        let c2: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2, T, Q);
+
+        (c0, c1, c2)
+    }
+    fn tensor_DBG<const PQ: u64, const T: u64>(
+        a: &Self,
+        b: &Self,
+    ) -> (Rq<Q, N>, Rq<Q, N>, Rq<Q, N>) {
+        // iacr 2021-204:
+        // expand Q->PQ
+        // let a: RLWE<PQ, N> = a.remodule::<PQ>();
+        // let b: RLWE<PQ, N> = b.remodule::<PQ>();
+
+        // tensor
+        let c0: Rq<Q, N> = a.0 * b.0; // NTT mul
+        let c1: Rq<Q, N> = a.0 * b.1 + a.1 * b.0; // NTT mul
+        let c2: Rq<Q, N> = a.1 * b.1; // NTT mul
+
+        // scale down
+        let c0: Rq<Q, N> = c0.mul_div_round(T, Q);
+        let c1: Rq<Q, N> = c1.mul_div_round(T, Q);
+        let c2: Rq<Q, N> = c2.mul_div_round(T, Q);
+
+        // expand^-1 PQ->Q
+        // let c0: Rq<Q, N> = c0.remodule::<Q>();
+        // let c1: Rq<Q, N> = c1.remodule::<Q>();
+        // let c2: Rq<Q, N> = c2.remodule::<Q>();
+
+        (c0, c1, c2)
+    }
+    /// ciphertext multiplication
+    fn mul<const PQ: u64, const T: u64>(rlk: &RLK<PQ, N>, a: &Self, b: &Self) -> Self {
+        let (c0, c1, c2) = Self::tensor::<PQ, T>(a, b);
+        BFV::<Q, N, T>::relinearize_204::<PQ>(&rlk, &c0, &c1, &c2)
+    }
+}
+// naive mul in the ring Rq, reusing the ring::naive_mul and then applying mod(X^N +1)
+fn tmp_naive_mul<const Q: u64, const N: usize>(a: Rq<Q, N>, b: Rq<Q, N>) -> Rq<Q, N> {
+    Rq::<Q, N>::from_vec_i64(arithmetic::ring::naive_mul(&a.to_r(), &b.to_r()))
 }
 
 impl<const Q: u64, const N: usize> ops::Add<RLWE<Q, N>> for RLWE<Q, N> {
@@ -39,17 +113,15 @@ Self::add(self, rhs)
     }
 }
 
-impl<const Q: u64, const N: usize, const T: u64> ops::Add<&PR<T, N>> for &RLWE<Q, N> {
+impl<const Q: u64, const N: usize, const T: u64> ops::Add<&Rq<T, N>> for &RLWE<Q, N> {
     type Output = RLWE<Q, N>;
-    fn add(self, rhs: &PR<T, N>) -> Self::Output {
-        // todo!()
+    fn add(self, rhs: &Rq<T, N>) -> Self::Output {
         BFV::<Q, N, T>::add_const(self, rhs)
     }
 }
-impl<const Q: u64, const N: usize, const T: u64> ops::Mul<&PR<T, N>> for &RLWE<Q, N> {
+impl<const Q: u64, const N: usize, const T: u64> ops::Mul<&Rq<T, N>> for &RLWE<Q, N> {
     type Output = RLWE<Q, N>;
-    fn mul(self, rhs: &PR<T, N>) -> Self::Output {
-        // todo!()
+    fn mul(self, rhs: &Rq<T, N>) -> Self::Output {
         BFV::<Q, N, T>::mul_const(&self, rhs)
     }
 }
@@ -71,50 +143,175 @@ // let Xi_key = Bernoulli::new(0.5)?;
         let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 
         // secret key
-        let s = PR::<Q, N>::rand_u64(&mut rng, Xi_key)?;
+        let s = Rq::<Q, N>::rand_u64(&mut rng, Xi_key)?;
+
+        #[cfg(test)] // sanity check
+        assert!(s.infinity_norm() <= 1, "{:?}", s.coeffs());
 
         // pk = (-a * s + e, a)
-        let a = PR::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, Q))?;
-        let e = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?;
+        let a = Rq::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, Q))?;
+        let e = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
         let pk: PublicKey<Q, N> = PublicKey((&(-a) * &s) + e, a.clone());
         Ok((SecretKey(s), pk))
     }
 
-    pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N>, m: &PR<T, N>) -> Result<RLWE<Q, N>> {
-        let Xi_key = Uniform::new(-1_f64, 1_f64);
+    pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N>, m: &Rq<T, N>) -> Result<RLWE<Q, N>> {
+        // let Xi_key = Uniform::new(-1_f64, 1_f64);
+        // let Xi_key = Uniform::new(0_f64, 2_f64);
+        let Xi_key = Uniform::new(0_u64, 2_u64);
         let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 
-        let u = PR::<Q, N>::rand_f64(&mut rng, Xi_key)?;
-        let e_1 = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?;
-        let e_2 = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?;
+        let u = Rq::<Q, N>::rand_u64(&mut rng, Xi_key)?;
+        let e_1 = Rq::<Q, N>::rand_f64_abs(&mut rng, Xi_err)?;
+        let e_2 = Rq::<Q, N>::rand_f64_abs(&mut rng, Xi_err)?;
+
+        // println!("{:?}", &e_1.coeffs());
+        // println!("{:?}", &e_2.coeffs());
+
+        #[cfg(test)] // sanity check
+        assert!(u.infinity_norm() <= 1, "{:?}", u.coeffs());
 
         // migrate m's coeffs to the bigger modulus Q (from T)
-        let m = PR::<Q, N>::from_vec_u64(m.coeffs().iter().map(|m_i| m_i.0).collect());
-        let c0 = &pk.0 * &u + e_1 + m * Self::DELTA;
+        let m = m.remodule::<Q>();
+        #[cfg(test)]
+        {
+            // sanity check // TODO rm
+            let m_remod_naive =
+                Rq::<Q, N>::from_vec_u64(m.coeffs().iter().map(|m_i| m_i.0).collect());
+            assert_eq!(m_remod_naive, m);
+        }
+
+        // let c0 = &pk.0 * &u + e_1 + m * Self::DELTA;
+        let c0 = &pk.0 * &u + e_1 + m.mul_div_round(Q, T); // TODO use DELTA?
         let c1 = &pk.1 * &u + e_2;
+        // let c0 = tmp_naive_mul(pk.0, u) + e_1 + m * Self::DELTA;
+        // let c0 = tmp_naive_mul(pk.0, u) + e_1 + m.mul_div_round(Q, T);
+        // let c1 = tmp_naive_mul(pk.1, u)
+        //     // &pk.1 * &u
+        //     + e_2;
         Ok(RLWE::<Q, N>(c0, c1))
     }
 
-    pub fn decrypt(sk: &SecretKey<Q, N>, c: &RLWE<Q, N>) -> PR<T, N> {
+    pub fn decrypt(sk: &SecretKey<Q, N>, c: &RLWE<Q, N>) -> Rq<T, N> {
         let cs = c.0 + c.1 * sk.0; // done in mod q
-        let r: Vec<u64> = cs
-            .coeffs()
-            .iter()
-            .map(|e| ((T as f64 * e.0 as f64) / Q as f64).round() as u64)
-            .collect();
-        PR::<T, N>::from_vec_u64(r)
+
+        // let c1s = tmp_naive_mul(c.1, sk.0);
+        // // let c1s = arithmetic::ring::naive_mul(&c.1.to_r(), &sk.0.to_r()); // TODO rm
+        // // let c1s = Rq::<Q, N>::from_vec_i64(c1s);
+        // let cs = c.0 + c1s;
+
+        // let r: Vec<u64> = cs
+        //     .coeffs()
+        //     .iter()
+        //     .map(|e| ((T as f64 * e.0 as f64) / Q as f64).round() as u64)
+        //     .collect();
+        // Rq::<T, N>::from_vec_u64(r)
+
+        let r: Rq<Q, N> = cs.mul_div_round(T, Q);
+        r.remodule::<T>()
     }
 
-    fn add_const(c: &RLWE<Q, N>, m: &PR<T, N>) -> RLWE<Q, N> {
+    fn add_const(c: &RLWE<Q, N>, m: &Rq<T, N>) -> RLWE<Q, N> {
         // assuming T<Q, move m from Zq<T> to Zq<Q>
         let m = m.remodule::<Q>();
         RLWE::<Q, N>(c.0 + m * Self::DELTA, c.1)
     }
-    fn mul_const(c: &RLWE<Q, N>, m: &PR<T, N>) -> RLWE<Q, N> {
+    fn mul_const(c: &RLWE<Q, N>, m: &Rq<T, N>) -> RLWE<Q, N> {
         // assuming T<Q, move m from Zq<T> to Zq<Q>
         let m = m.remodule::<Q>();
         RLWE::<Q, N>(c.0 * m * Self::DELTA, c.1)
     }
+
+    fn rlk_key<const PQ: u64>(mut rng: impl Rng, s: &SecretKey<Q, N>) -> Result<RLK<PQ, N>> {
+        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; // TODO review Xi' instead of Xi
+
+        let s = s.0.remodule::<PQ>();
+        let a = Rq::<PQ, N>::rand_u64(&mut rng, Uniform::new(0_u64, PQ))?;
+        let e = Rq::<PQ, N>::rand_f64(&mut rng, Xi_err)?;
+
+        // let rlk_1: Rq<PQ, N> = (&(-a) * &s) - e + (s * s) * P;
+        let P = PQ / Q;
+
+        // let rlk: RLK<PQ, N> = RLK::<PQ, N>((&(-a) * &s) - e + (s * s) * P, a.clone());
+        let rlk: RLK<PQ, N> = RLK::<PQ, N>(
+            -(tmp_naive_mul(a, s) + e) + tmp_naive_mul(s, s) * P,
+            a.clone(),
+        );
+        // let rlk: RLK<PQ, N> = RLK::<PQ, N>(-(&a * &s + e) + (s * s) * P, a.clone());
+
+        // let pq = P * Q;
+        // let a = Rq::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, pq))?;
+        // let e = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
+        //
+        // let rlk_0: Rq<Q, N> = (&(-a) * &s) - e + (s * s) * P;
+        // let rlk_0 = rlk_0.remodule::<>();
+        // let rlk: RLK<Q, N> = RLK(rlk_0, a);
+        Ok(rlk)
+    }
+
+    fn relinearize<const PQ: u64>(
+        rlk: &RLK<PQ, N>,
+        c0: &Rq<Q, N>,
+        c1: &Rq<Q, N>,
+        c2: &Rq<Q, N>,
+    ) -> RLWE<Q, N> {
+        let P = PQ / Q;
+
+        // let c2 = c2.remodule::<PQ>();
+        // let c2 = c2.to_r();
+
+        let c2rlk0: Vec<f64> = (c2.to_r() * rlk.0.to_r())
+            .coeffs()
+            .iter()
+            .map(|e| (*e as f64 / P as f64).round())
+            .collect();
+
+        let c2rlk1: Vec<f64> = (c2.to_r() * rlk.1.to_r())
+            .coeffs()
+            .iter()
+            .map(|e| (*e as f64 / P as f64).round())
+            .collect();
+
+        let r0 = Rq::<Q, N>::from_vec_f64(c2rlk0);
+        let r1 = Rq::<Q, N>::from_vec_f64(c2rlk1);
+
+        let res = RLWE::<Q, N>(c0 + &r0, c1 + &r1);
+        res
+    }
+    fn relinearize_204<const PQ: u64>(
+        rlk: &RLK<PQ, N>,
+        c0: &Rq<Q, N>,
+        c1: &Rq<Q, N>,
+        c2: &Rq<Q, N>,
+    ) -> RLWE<Q, N> {
+        let P = PQ / Q;
+
+        // let c2 = c2.remodule::<PQ>();
+        // let c2 = c2.to_r();
+
+        // let c2rlk0: Vec<f64> = (c2.remodule::<PQ>() * rlk.0)
+        use arithmetic::ring::naive_mul;
+        let c2rlk0: Vec<i64> = naive_mul(&c2.to_r(), &rlk.0.to_r());
+        // .coeffs()
+        // .iter()
+        // .map(|e| (*e as f64 / P as f64).round())
+        // .collect();
+
+        let c2rlk1: Vec<i64> = naive_mul(&c2.to_r(), &rlk.1.to_r());
+        // .coeffs()
+        // .iter()
+        // .map(|e| (*e as f64 / P as f64).round())
+        // .collect();
+        //
+        let r0: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2rlk0, 1, P);
+        let r1: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2rlk1, 1, P);
+
+        // let r0 = Rq::<Q, N>::from_vec_f64(c2rlk0);
+        // let r1 = Rq::<Q, N>::from_vec_f64(c2rlk1);
+
+        let res = RLWE::<Q, N>(c0 + &r0, c1 + &r1);
+        res
+    }
 }
 
 #[cfg(test)]
@@ -133,15 +330,17 @@ type S = BFV<Q, N, T>;
 
         let mut rng = rand::thread_rng();
 
-        let (sk, pk) = S::new_key(&mut rng)?;
+        for _ in 0..1000 {
+            let (sk, pk) = S::new_key(&mut rng)?;
 
-        let msg_dist = Uniform::new(0_u64, T);
-        let m = PR::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let msg_dist = Uniform::new(0_u64, T);
+            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
 
-        let c = S::encrypt(rng, &pk, &m)?;
-        let m_recovered = S::decrypt(&sk, &c);
+            let c = S::encrypt(&mut rng, &pk, &m)?;
+            let m_recovered = S::decrypt(&sk, &c);
 
-        assert_eq!(m, m_recovered);
+            assert_eq!(m, m_recovered);
+        }
 
         Ok(())
     }
@@ -149,7 +348,8 @@
     #[test]
     fn test_addition() -> Result<()> {
         const Q: u64 = 2u64.pow(16) + 1;
-        const N: usize = 32;
+        // const N: usize = 32;
+        const N: usize = 4;
         const T: u64 = 4; // plaintext modulus
         type S = BFV<Q, N, T>;
 
@@ -158,8 +358,8 @@
         let (sk, pk) = S::new_key(&mut rng)?;
 
         let msg_dist = Uniform::new(0_u64, T);
-        let m1 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?;
-        let m2 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?;
+        let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+        let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
 
         let c1 = S::encrypt(&mut rng, &pk, &m1)?;
         let c2 = S::encrypt(&mut rng, &pk, &m2)?;
@@ -174,7 +374,7 @@ Ok(())
     }
 
     #[test]
-    fn test_constant_add_mul() -> Result<()> {
+    fn test_constant_add() -> Result<()> {
         const Q: u64 = 2u64.pow(16) + 1;
         const N: usize = 32;
         const T: u64 = 4; // plaintext modulus
@@ -185,66 +385,271 @@
         let (sk, pk) = S::new_key(&mut rng)?;
 
         let msg_dist = Uniform::new(0_u64, T);
-        let m1 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?;
-        let m2_const = PR::<T, N>::rand_u64(&mut rng, msg_dist)?;
+        let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+        let m2_const = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
 
         let c1 = S::encrypt(&mut rng, &pk, &m1)?;
 
         let c3_add = &c1 + &m2_const;
-        let c3_mul = &c1 * &m2_const;
+        // let c3_mul = &c1 * &m2_const;
 
         let m3_add_recovered = S::decrypt(&sk, &c3_add);
-        let m3_mul_recovered = S::decrypt(&sk, &c3_mul);
+        // let m3_mul_recovered = S::decrypt(&sk, &c3_mul);
 
         assert_eq!(m1 + m2_const, m3_add_recovered);
+        //
+        // let mut mul_res = naive_poly_mul::<T>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
+        // arithmetic::ring::modulus::<T, N>(&mut mul_res);
+        // dbg!(&mul_res);
+        // let mul_res_2 =
+        //     naive_poly_mul_2::<T, N>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
+        // assert_eq!(mul_res, mul_res_2);
+        // let mul_res = PR::<T, N>::from_vec(mul_res);
+        // assert_eq!(mul_res.coeffs(), m3_mul_recovered.coeffs());
 
-        let mut mul_res = naive_poly_mul::<T>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
-        arithmetic::ring::modulus::<T, N>(&mut mul_res);
-        dbg!(&mul_res);
-        let mul_res_2 =
-            naive_poly_mul_2::<T, N>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
-        assert_eq!(mul_res, mul_res_2);
-        let mul_res = PR::<T, N>::from_vec(mul_res);
-        assert_eq!(mul_res.coeffs(), m3_mul_recovered.coeffs());
+        Ok(())
+    }
+
+    #[test]
+    fn test_tensor() -> Result<()> {
+        const Q: u64 = 2u64.pow(16) + 1; // q prime, and 2^q + 1 shape
+        const N: usize = 8;
+        const T: u64 = 4; // plaintext modulus
+
+        // const P: u64 = Q;
+        const P: u64 = Q * Q;
+        // const P: u64 = 2_u64.pow(13) * Q + 1;
+        const PQ: u64 = P * Q;
+
+        let mut rng = rand::thread_rng();
+
+        let msg_dist = Uniform::new(0_u64, T);
+        for _ in 0..10_000 {
+            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+
+            test_tensor_opt::<Q, N, T, PQ>(&mut rng, m1, m2)?;
+        }
 
         Ok(())
     }
+    fn test_tensor_opt<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
+        mut rng: impl Rng,
+        m1: Rq<T, N>,
+        m2: Rq<T, N>,
+    ) -> Result<()> {
+        let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
 
-    fn naive_poly_mul<const T: u64>(a: &[Zq<T>], b: &[Zq<T>]) -> Vec<Zq<T>> {
-        let mut result: Vec<Zq<T>> = vec![Zq::zero(); a.len() + b.len() - 1];
-        for (i, &ai) in a.iter().enumerate() {
-            for (j, &bj) in b.iter().enumerate() {
-                result[i + j] = result[i + j] + (ai * bj);
-            }
+        let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
+        let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
+
+        let (c_a, c_b, c_c) = RLWE::<Q, N>::tensor::<PQ, T>(&c1, &c2);
+
+        // decrypt non-relinearized mul result
+        let m3: Rq<Q, N> = c_a + c_b * sk.0 + c_c * sk.0 * sk.0;
+        let m3: Rq<Q, N> = m3.mul_div_round(T, Q);
+        let m3 = m3.remodule::<T>();
+
+        let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
+        assert_eq!(
+            m3.coeffs().to_vec(),
+            naive.coeffs().to_vec(),
+            "\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
+            m1.coeffs(),
+            m2.coeffs()
+        );
+        if m3.coeffs().to_vec() != naive.coeffs().to_vec() {
+            return Err(anyhow!("not eq"));
         }
-        result
+
+        Ok(())
     }
-    fn naive_poly_mul_2<const T: u64, const N: usize>(
-        poly1: &[Zq<T>],
-        poly2: &[Zq<T>],
-    ) -> Vec<Zq<T>> {
-        let degree1 = poly1.len();
-        let degree2 = poly2.len();
+
+    fn test_tensor_opt_DBG<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
+        mut rng: impl Rng,
+        m1: Rq<T, N>,
+        m2: Rq<T, N>,
+    ) -> Result<()> {
+        let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
+
+        let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
+        let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
 
-        // The degree of the resulting polynomial will be degree1 + degree2 - 1
-        let mut result = vec![Zq::zero(); degree1 + degree2 - 1];
+        let (c_a, c_b, c_c) = RLWE::<Q, N>::tensor::<PQ, T>(&c1, &c2);
 
-        // Perform the multiplication
-        for i in 0..degree1 {
-            for j in 0..degree2 {
-                result[i + j] = result[i + j] + poly1[i] * poly2[j];
-            }
+        // decrypt non-relinearized mul result
+        let m3: Rq<Q, N> = c_a + c_b * sk.0 + c_c * sk.0 * sk.0;
+        dbg!(m3);
+        let m3: Rq<Q, N> = m3.mul_div_round(T, Q);
+        dbg!(m3);
+        let m3 = m3.remodule::<T>();
+        dbg!(m3);
+
+        // let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
+        // let naive = m1.remodule::<Q>() * m2.remodule::<Q>();
+        let naive = (m1.remodule::<Q>() * m2.remodule::<Q>()).remodule::<T>();
+        dbg!(naive);
+        assert_eq!(
+            m3.coeffs().to_vec(),
+            naive.coeffs().to_vec(),
+            "\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
+            m1.coeffs(),
+            m2.coeffs()
+        );
+        // if m3.coeffs().to_vec() != naive.coeffs().to_vec() {
+        //     return Err(anyhow!("not eq"));
+        // }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_mul_relin() -> Result<()> {
+        const Q: u64 = 2u64.pow(16) + 1;
+        const N: usize = 32;
+        const T: u64 = 4; // plaintext modulus
+        type S = BFV<Q, N, T>;
+
+        const P: u64 = Q * Q;
+        const PQ: u64 = P * Q;
+
+        let mut rng = rand::thread_rng();
+        let msg_dist = Uniform::new(0_u64, T);
+
+        for _ in 0..100 {
+            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+
+            test_mul_relin_opt::<Q, N, T, PQ>(&mut rng, m1, m2)?;
         }
 
-        // Reduce the result modulo x^N + 1
-        let mut reduced_result = vec![Zq::zero(); N];
+        Ok(())
+    }
+
+    fn test_mul_relin_opt<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
+        mut rng: impl Rng,
+        m1: Rq<T, N>,
+        m2: Rq<T, N>,
+    ) -> Result<()> {
+        let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
+
+        let rlk = BFV::<Q, N, T>::rlk_key::<PQ>(&mut rng, &sk)?;
+
+        let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
+        let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
+
+        let c3 = RLWE::<Q, N>::mul::<PQ, T>(&rlk, &c1, &c2);
+
+        let m3 = BFV::<Q, N, T>::decrypt(&sk, &c3);
 
-        for i in 0..result.len() {
-            let mod_index = i % N; // wrap around using modulo N
-            reduced_result[mod_index] += result[i];
+        let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
+        assert_eq!(m3.coeffs().to_vec(), naive.coeffs().to_vec(),
+            "\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
+            m1.coeffs(),
+            m2.coeffs()
+            );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_naive_mul() -> Result<()> {
+        const Q: u64 = 2u64.pow(16) + 1; // prime, and 2^q + 1 shape
+        const N: usize = 4;
+        const T: u64 = 4; // plaintext modulus
+
+        let mut rng = rand::thread_rng();
+        let msg_dist = Uniform::new(0_u64, T);
+        // for _ in 0..10_000 {
+        for _ in 0..2 {
+            println!("---");
+            // let a = Rq::<Q, N>::rand_u64(&mut rng, msg_dist)?;
+            // let b = Rq::<Q, N>::rand_u64(&mut rng, msg_dist)?;
+            // let a = Rq::<Q, N>::from_vec_u64(vec![Q - 1, Q - 2, Q - 3, Q - 3]);
+            // let b = Rq::<Q, N>::from_vec_u64(vec![Q - 3, Q - 3, Q - 2, Q - 1]);
+            let a = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            let b = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
+            dbg!(&a);
+            dbg!(&b);
+            let (_, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
+            let ciph1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &a)?;
+            let ciph2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &b)?;
+            let a = ciph1.0;
+            let b = ciph2.0;
+            dbg!(&a);
+            dbg!(&b);
+
+            let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
+            let c0 = Rq::<Q, N>::from_vec_i64(c0);
+            let c1 = tmp_naive_mul(a, b); // naive mul
+            let c2: Rq<Q, N> = a * b; // NTT mul
+
+            println!("{:?}", c0.coeffs());
+            println!("{:?}", c1.coeffs());
+            println!("{:?}", c2.coeffs());
+            assert_eq!(c0, c2);
+            assert_eq!(c1, c2);
+
+            // scale by Delta=Q/T
+            let a = a.mul_div_round(Q, T);
+            let b = b.mul_div_round(Q, T);
+            dbg!(&a);
+            dbg!(&b);
+
+            let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
+            let c0 = Rq::<Q, N>::from_vec_i64(c0);
+            let c1 = tmp_naive_mul(a, b); // naive mul
+            let c2: Rq<Q, N> = a * b; // NTT mul
+
+            println!("{:?}", c0.coeffs());
+            println!("{:?}", c1.coeffs());
+            println!("{:?}", c2.coeffs());
+            assert_eq!(c0, c2);
+            assert_eq!(c1, c2);
+
+            let c0 = c0.mul_div_round(T, Q);
+            let c1 = c1.mul_div_round(T, Q);
+            let c2 = c2.mul_div_round(T, Q);
+            println!("{:?}", c0.coeffs());
+            println!("{:?}", c1.coeffs());
+            println!("{:?}", c2.coeffs());
+            assert_eq!(c0, c2);
+            assert_eq!(c1, c2);
+
+            /*
+            // now same as before, but multiplying by T/Q
+            let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
+            let c0: Vec<f64> = c0
+                .iter()
+                .map(|e| ((T as f64 * *e as f64) / Q as f64).round())
+                .collect();
+            let c0 = Rq::<Q, N>::from_vec_f64(c0);
+            dbg!(&c0.coeffs());
+
+            let a = a.mul_div_round(T, Q);
+            let b = b.mul_div_round(T, Q);
+            println!("a{:?}", a.coeffs());
+            println!("b{:?}", b.coeffs());
+
+            let c4: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
+            let c4 = Rq::<Q, N>::from_vec_i64(c4);
+            let c4 = c4.mul_div_round(T, Q);
+
+            let c1 = tmp_naive_mul(a, b); // naive mul
+            let c1 = c1.mul_div_round(T, Q);
+            let c2 = a * b; // NTT mul
+            let c2 = c2.mul_div_round(T, Q);
+
+            println!("{:?}", c0.coeffs());
+            println!("{:?}", c1.coeffs());
+            println!("{:?}", c2.coeffs());
+            println!("{:?}", c4.coeffs());
+            assert_eq!(c0, c2);
+            assert_eq!(c4, c2);
+            assert_eq!(c1, c2);
+             */
         }
 
-        // Return the reduced polynomial
-        reduced_result
+        Ok(())
     }
 }