From 8e2e522f8a8c5eb503da9d4a8ef65de1c51fa20b Mon Sep 17 00:00:00 2001
From: arnaucube <git@arnaucube.com>
Date: Sat, 2 Aug 2025 21:42:36 +0000
Subject: [PATCH] arith/torus: add left_rotate & mod_switch

---
 arith/src/ring_torus.rs | 62 +++++++++++++++++++++++++++++++++++++++++
 arith/src/torus.rs      |  9 +++++-
 arith/src/tuple_ring.rs | 20 +++++++++++++
 arith/src/zq.rs         |  3 ++
 4 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/arith/src/ring_torus.rs b/arith/src/ring_torus.rs
index 1f35adf..1a8192d 100644
--- a/arith/src/ring_torus.rs
+++ b/arith/src/ring_torus.rs
@@ -82,6 +82,31 @@ impl<const N: usize> Ring for Tn<N> {
     }
 }
 
+impl<const N: usize> Tn<N> {
+    // multiply self by X^-h
+    pub fn left_rotate(&self, h: usize) -> Self {
+        dbg!(&h);
+        dbg!(&N);
+        let h = h % N;
+        assert!(h < N);
+        let c = self.0;
+        // c[h], c[h+1], c[h+2], ..., c[n-1],  -c[0], -c[1], ..., -c[h-1]
+        // let r: Vec<T64> = vec![c[h..N], c[0..h].iter().map(|&c_i| -c_i).collect()].concat();
+        dbg!(&h);
+        let r: Vec<T64> = c[h..N]
+            .iter()
+            .copied()
+            .chain(c[0..h].iter().map(|&x| -x))
+            .collect();
+        Self::from_vec(r)
+    }
+
+    pub fn from_vec_u64(v: Vec<u64>) -> Self {
+        let coeffs = v.iter().map(|c| T64(*c)).collect();
+        Self::from_vec(coeffs)
+    }
+}
+
 // apply mod (X^N+1)
 pub fn modulus<const N: usize>(p: &mut Vec<T64>) {
     if p.len() < N {
@@ -227,3 +252,40 @@ impl<const N: usize> Mul<&u64> for &Tn<N> {
         Tn::<N>(array::from_fn(|i| self.0[i] * *s))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_left_rotate() {
+        const N: usize = 4;
+        let f = Tn::<N>::from_vec(
+            vec![2i64, 3, -4, -1]
+                .iter()
+                .map(|c| T64(*c as u64))
+                .collect(),
+        );
+
+        // expect f*x^-3 == -1 -2x -3x^2 +4x^3
+        assert_eq!(
+            f.left_rotate(3),
+            Tn::<N>::from_vec(
+                vec![-1i64, -2, -3, 4]
+                    .iter()
+                    .map(|c| T64(*c as u64))
+                    .collect(),
+            )
+        );
+        // expect f*x^-1 == 3 -4x -1x^2 -2x^3
+        assert_eq!(
+            f.left_rotate(1),
+            Tn::<N>::from_vec(
+                vec![3i64, -4, -1, -2]
+                    .iter()
+                    .map(|c| T64(*c as u64))
+                    .collect(),
+            )
+        );
+    }
+}
diff --git a/arith/src/torus.rs b/arith/src/torus.rs
index aa706fd..fa48ad7 100644
--- a/arith/src/torus.rs
+++ b/arith/src/torus.rs
@@ -50,8 +50,15 @@ impl Ring for T64 {
         todo!()
     }
 
+    // modulus switch from Q to Q2: self * Q2/Q
     fn mod_switch<const Q2: u64>(&self) -> T64 {
-        todo!()
+        // for the moment we assume Q|Q2, since Q=2^64, check that Q2 is a power
+        // of two:
+        assert!(Q2.is_power_of_two());
+        // since Q=2^64, dividing Q2/Q is equivalent to dividing 2^log2(Q2)/2^64
+        // which would be like right-shifting 64-log2(Q2).
+        let log2_q2 = 63 - Q2.leading_zeros();
+        T64(self.0 >> (64 - log2_q2))
     }
 
     fn mul_div_round(&self, num: u64, den: u64) -> Self {
diff --git a/arith/src/tuple_ring.rs b/arith/src/tuple_ring.rs
index 01456e8..1da5e08 100644
--- a/arith/src/tuple_ring.rs
+++ b/arith/src/tuple_ring.rs
@@ -17,8 +17,14 @@ use crate::Ring;
 /// since if using a fixed-size array it would overflow the stack.
 #[derive(Clone, Debug)]
 pub struct TR<R: Ring, const K: usize>(pub Vec<R>);
+// TODO rm pub from Vec<R>, so that TR can not be created from a Vec with
+// invalid length, since it has to be created using the `new` method.
 
 impl<R: Ring, const K: usize> TR<R, K> {
+    pub fn new(v: Vec<R>) -> Self {
+        assert_eq!(v.len(), K);
+        Self(v)
+    }
     pub fn zero() -> Self {
         Self((0..K).into_iter().map(|_| R::zero()).collect())
     }
@@ -37,6 +43,20 @@ impl<R: Ring, const K: usize> TR<R, K> {
     }
 }
 
+impl<const K: usize> TR<crate::torus::T64, K> {
+    pub fn mod_switch<const Q2: u64>(&self) -> TR<crate::torus::T64, K> {
+        TR(self.0.iter().map(|c_i| c_i.mod_switch::<Q2>()).collect())
+    }
+    // pub fn mod_switch(&self, Q2: u64) -> TR<crate::torus::T64, K> {
+    //     TR(self.0.iter().map(|c_i| c_i.mod_switch(Q2)).collect())
+    // }
+}
+impl<const N: usize, const K: usize> TR<crate::ring_torus::Tn<N>, K> {
+    pub fn left_rotate(&self, h: usize) -> Self {
+        TR(self.0.iter().map(|c_i| c_i.left_rotate(h)).collect())
+    }
+}
+
 impl<R: Ring, const K: usize> TR<R, K> {
     pub fn iter(&self) -> std::slice::Iter<R> {
         self.0.iter()
diff --git a/arith/src/zq.rs b/arith/src/zq.rs
index 3f42887..f6b5904 100644
--- a/arith/src/zq.rs
+++ b/arith/src/zq.rs
@@ -269,6 +269,9 @@ impl<const Q: u64> Neg for Zq<Q> {
     type Output = Self;
 
     fn neg(self) -> Self::Output {
+        if self.0 == 0 {
+            return self;
+        }
         Zq(Q - self.0)
     }
 }