implement strong coeff

mixnet-rs/ordering/src/bin/coeff.rs

@@ -38,21 +38,46 @@ fn parse_data_msg(value: &str) -> DataMessage {
     }
 }
 
+pub fn strong_and_casual_coeff(seq1: &[Entry], seq2: &[Entry]) -> (u64, u64) {
+    let coeffs = coeff(seq1, seq2, CoefficientType::StrongAndCasual);
+    assert_eq!(coeffs.len(), 2);
+    (coeffs[0], coeffs[1])
+}
+
+pub fn weak_coeff(seq1: &[Entry], seq2: &[Entry]) -> u64 {
+    let coeffs = coeff(seq1, seq2, CoefficientType::Weak);
+    assert_eq!(coeffs.len(), 1);
+    coeffs[0]
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum CoefficientType {
-    Strong,
-    Casual,
+enum CoefficientType {
+    StrongAndCasual,
     Weak,
 }
 
-pub fn coeff(seq1: &[Entry], seq2: &[Entry], coeff_type: CoefficientType) -> u64 {
-    let mut coeff = 0;
+impl CoefficientType {
+    fn zero(&self) -> Vec<u64> {
+        match self {
+            CoefficientType::StrongAndCasual => vec![0, 0],
+            CoefficientType::Weak => vec![0],
+        }
+    }
+}
+
+fn coeff(seq1: &[Entry], seq2: &[Entry], coeff_type: CoefficientType) -> Vec<u64> {
+    let mut coeffs = coeff_type.zero();
     let mut i = 0;
 
     while i < seq1.len() {
         if let Entry::Data(_) = &seq1[i] {
             let (c, next_i) = coeff_from(seq1, i, seq2, coeff_type);
-            coeff += c;
+            assert_eq!(coeffs.len(), c.len());
+            coeffs = coeffs
+                .iter()
+                .zip(c.iter())
+                .map(|(a, b)| a.checked_add(*b).unwrap())
+                .collect();
 
             if next_i != i {
                 i = next_i;
@@ -64,7 +89,7 @@ pub fn coeff(seq1: &[Entry], seq2: &[Entry], coeff_type: CoefficientType) -> u64
         }
     }
 
-    coeff
+    coeffs
 }
 
 fn coeff_from(
@@ -72,10 +97,10 @@ fn coeff_from(
     start_idx: usize,
     seq2: &[Entry],
     coeff_type: CoefficientType,
-) -> (u64, usize) {
+) -> (Vec<u64>, usize) {
     let msg1 = match seq1[start_idx] {
         Entry::Data(msg) => msg,
-        _ => panic!("Entry at {start_idx} must be Message"),
+        _ => panic!("Entry at index {start_idx} must be Message"),
     };
 
     for (j, entry) in seq2.iter().enumerate() {
@@ -83,18 +108,42 @@ fn coeff_from(
             if msg1 == *msg2 {
                 // Found the 1st matching msg. Start finding the next adjacent matching msg.
                 match coeff_type {
-                    CoefficientType::Strong => todo!(),
-                    CoefficientType::Casual => {
-                        return casual_coeff_from(seq1, start_idx, seq2, j);
+                    CoefficientType::StrongAndCasual => {
+                        return strong_and_casual_coeff_from(seq1, start_idx, seq2, j);
                     }
                     CoefficientType::Weak => {
-                        return weak_coeff_from(seq1, start_idx, seq2, j);
+                        let (coeff, next_start_idx) = weak_coeff_from(seq1, start_idx, seq2, j);
+                        return (vec![coeff], next_start_idx);
                     }
                 }
             }
         }
     }
-    (0, start_idx)
+
+    // Couldn't find any matching msg in seq2. Returning the zero coefficients and the same start index.
+    (coeff_type.zero(), start_idx)
+}
+
+fn strong_and_casual_coeff_from(
+    seq1: &[Entry],
+    start_idx: usize,
+    seq2: &[Entry],
+    seq2_start_idx: usize,
+) -> (Vec<u64>, usize) {
+    let (casual_coeff, next_start_idx) = casual_coeff_from(seq1, start_idx, seq2, seq2_start_idx);
+
+    // The partially calcuated casual coefficient means the number of consecutive matching pairs,
+    // which is also what should be calculated for the strong coefficient.
+    // So, we reuse it.
+    let num_consecutive_pairs = casual_coeff;
+    let strong_coeff = if num_consecutive_pairs == 0 {
+        0
+    } else {
+        (num_consecutive_pairs)
+            .checked_pow(num_consecutive_pairs.try_into().unwrap())
+            .unwrap()
+    };
+    (vec![strong_coeff, casual_coeff], next_start_idx)
 }
 
 fn casual_coeff_from(
@@ -103,6 +152,7 @@ fn casual_coeff_from(
     seq2: &[Entry],
     seq2_start_idx: usize,
 ) -> (u64, usize) {
+    // Find the number of consecutive matching pairs while accounting for noises.
     let mut coeff = 0;
     let mut i = start_idx + 1;
     let mut j = seq2_start_idx + 1;
@@ -137,6 +187,7 @@ fn weak_coeff_from(
     seq2: &[Entry],
     seq2_start_idx: usize,
 ) -> (u64, usize) {
+    // Find the number of consecutive matching pairs with ignoring noises.
     let mut coeff = 0;
     let mut i = start_idx + 1;
     let mut j = seq2_start_idx + 1;
@@ -167,8 +218,10 @@ fn skip_noise(seq: &[Entry], mut index: usize) -> usize {
 fn main() {
     let seq1 = load_sequence("/Users/yjlee/repos/nomos-simulations/mixnet-rs/results/ordering_e5s1_PureCoinFlipping_2024-09-01T19:30:03.957310+00:00_0d0h12m25s/paramset_1/iteration_0_0d0h0m2s/sent_seq_0.csv");
     let seq2 = load_sequence("/Users/yjlee/repos/nomos-simulations/mixnet-rs/results/ordering_e5s1_PureCoinFlipping_2024-09-01T19:30:03.957310+00:00_0d0h12m25s/paramset_1/iteration_0_0d0h0m2s/recv_seq_0.csv");
-    println!("casual:{:?}", coeff(&seq1, &seq2, CoefficientType::Casual));
-    println!("weak:{:?}", coeff(&seq1, &seq2, CoefficientType::Weak));
+    let (strong, casual) = strong_and_casual_coeff(&seq1, &seq2);
+    println!("strong:{:?}", strong);
+    println!("casual:{:?}", casual);
+    println!("weak:{:?}", weak_coeff(&seq1, &seq2));
 }
 
 #[cfg(test)]
@@ -176,116 +229,175 @@ mod tests {
     use super::*;
 
     #[test]
-    fn test_casual_coeff() {
+    fn test_strong_coeff() {
         // Empty sequences
-        assert_eq!(coeff(&[], &[], CoefficientType::Casual), 0);
+        assert_eq!(strong_coeff(&[], &[]), 0);
 
         // One matching pair without noise
         let seq = vec![data(1), data(2)];
-        assert_eq!(coeff(&seq, &seq, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq, &seq), 1);
 
         // One matching pair with noise
         let seq = vec![data(1), noise(10), data(2)];
-        assert_eq!(coeff(&seq, &seq, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq, &seq), 1);
 
         // One matching pair without noise from different sequences
         let seq1 = vec![data(1), data(2), data(3)];
         let seq2 = vec![data(1), data(2), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq1, &seq2), 1);
         let seq1 = vec![data(4), data(2), data(3)];
         let seq2 = vec![data(1), data(2), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq1, &seq2), 1);
 
         // One matching pair with noise from different sequences
         let seq1 = vec![data(1), noise(10), data(2), data(3)];
         let seq2 = vec![data(1), noise(10), data(2), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq1, &seq2), 1);
         let seq1 = vec![data(4), data(2), noise(10), data(3)];
         let seq2 = vec![data(1), data(2), noise(10), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 1);
+        assert_eq!(strong_coeff(&seq1, &seq2), 1);
 
         // Two pairs with noise
         let seq1 = vec![data(1), noise(10), data(2), data(3)];
         let seq2 = vec![data(1), noise(10), data(2), data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 2);
+        assert_eq!(strong_coeff(&seq1, &seq2), 4);
 
         // No match
         let seq1 = vec![data(1), data(2)];
         let seq2 = vec![data(2), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 0);
+        assert_eq!(strong_coeff(&seq1, &seq2), 0);
         let seq1 = vec![data(1), data(2)];
         let seq2 = vec![data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 0);
+        assert_eq!(strong_coeff(&seq1, &seq2), 0);
 
         // No match because of noise
         let seq1 = vec![data(1), noise(10), data(2)];
         let seq2 = vec![data(1), data(2)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 0);
+        assert_eq!(strong_coeff(&seq1, &seq2), 0);
         let seq1 = vec![data(1), noise(10), data(2)];
         let seq2 = vec![data(1), noise(5), data(2)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 0);
+        assert_eq!(strong_coeff(&seq1, &seq2), 0);
 
         // Matching pairs in different indexes
         let seq1 = vec![data(1), data(2), data(3), data(4)];
         let seq2 = vec![data(2), data(3), data(4), data(1)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 2);
+        assert_eq!(strong_coeff(&seq1, &seq2), 4);
         let seq1 = vec![data(1), data(2), data(3), data(4)];
         let seq2 = vec![data(1), data(2), data(5), data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Casual), 2);
+        assert_eq!(strong_coeff(&seq1, &seq2), 2);
+    }
+
+    #[test]
+    fn test_casual_coeff() {
+        // Empty sequences
+        assert_eq!(casual_coeff(&[], &[]), 0);
+
+        // One matching pair without noise
+        let seq = vec![data(1), data(2)];
+        assert_eq!(casual_coeff(&seq, &seq), 1);
+
+        // One matching pair with noise
+        let seq = vec![data(1), noise(10), data(2)];
+        assert_eq!(casual_coeff(&seq, &seq), 1);
+
+        // One matching pair without noise from different sequences
+        let seq1 = vec![data(1), data(2), data(3)];
+        let seq2 = vec![data(1), data(2), data(4)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 1);
+        let seq1 = vec![data(4), data(2), data(3)];
+        let seq2 = vec![data(1), data(2), data(3)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 1);
+
+        // One matching pair with noise from different sequences
+        let seq1 = vec![data(1), noise(10), data(2), data(3)];
+        let seq2 = vec![data(1), noise(10), data(2), data(4)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 1);
+        let seq1 = vec![data(4), data(2), noise(10), data(3)];
+        let seq2 = vec![data(1), data(2), noise(10), data(3)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 1);
+
+        // Two pairs with noise
+        let seq1 = vec![data(1), noise(10), data(2), data(3)];
+        let seq2 = vec![data(1), noise(10), data(2), data(3), data(4)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 2);
+
+        // No match
+        let seq1 = vec![data(1), data(2)];
+        let seq2 = vec![data(2), data(3)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 0);
+        let seq1 = vec![data(1), data(2)];
+        let seq2 = vec![data(3), data(4)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 0);
+
+        // No match because of noise
+        let seq1 = vec![data(1), noise(10), data(2)];
+        let seq2 = vec![data(1), data(2)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 0);
+        let seq1 = vec![data(1), noise(10), data(2)];
+        let seq2 = vec![data(1), noise(5), data(2)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 0);
+
+        // Matching pairs in different indexes
+        let seq1 = vec![data(1), data(2), data(3), data(4)];
+        let seq2 = vec![data(2), data(3), data(4), data(1)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 2);
+        let seq1 = vec![data(1), data(2), data(3), data(4)];
+        let seq2 = vec![data(1), data(2), data(5), data(3), data(4)];
+        assert_eq!(casual_coeff(&seq1, &seq2), 2);
     }
 
     #[test]
     fn test_weak_coeff() {
         // Empty sequences
-        assert_eq!(coeff(&[], &[], CoefficientType::Weak), 0);
+        assert_eq!(weak_coeff(&[], &[]), 0);
 
         // One matching pair without noise
         let seq = vec![data(1), data(2)];
-        assert_eq!(coeff(&seq, &seq, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq, &seq), 1);
 
         // One matching pair with noise
         let seq = vec![data(1), noise(10), data(2)];
-        assert_eq!(coeff(&seq, &seq, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq, &seq), 1);
 
         // One matching pair without noise from different sequences
         let seq1 = vec![data(1), data(2), data(3)];
         let seq2 = vec![data(1), data(2), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq1, &seq2), 1);
         let seq1 = vec![data(4), data(2), data(3)];
         let seq2 = vec![data(1), data(2), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq1, &seq2), 1);
 
         // One matching pair with noise from different sequences
         let seq1 = vec![data(1), noise(10), data(2), data(3)];
         let seq2 = vec![data(1), noise(5), data(2), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq1, &seq2), 1);
         let seq1 = vec![data(4), data(2), noise(10), data(3)];
         let seq2 = vec![data(1), data(2), noise(5), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq1, &seq2), 1);
         let seq1 = vec![data(4), data(2), noise(10), data(3)];
         let seq2 = vec![data(1), data(2), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 1);
+        assert_eq!(weak_coeff(&seq1, &seq2), 1);
 
         // Two pairs with noise
         let seq1 = vec![data(1), noise(10), data(2), data(3)];
         let seq2 = vec![data(1), noise(5), data(2), data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 2);
+        assert_eq!(weak_coeff(&seq1, &seq2), 2);
 
         // No match
         let seq1 = vec![data(1), data(2)];
         let seq2 = vec![data(2), data(3)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 0);
+        assert_eq!(weak_coeff(&seq1, &seq2), 0);
         let seq1 = vec![data(1), data(2)];
         let seq2 = vec![data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 0);
+        assert_eq!(weak_coeff(&seq1, &seq2), 0);
 
         // Matching pairs in different indexes
         let seq1 = vec![data(1), data(2), data(3), data(4)];
         let seq2 = vec![data(2), data(3), data(4), data(1)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 2);
+        assert_eq!(weak_coeff(&seq1, &seq2), 2);
         let seq1 = vec![data(1), data(2), data(3), data(4)];
         let seq2 = vec![data(1), data(2), data(5), data(3), data(4)];
-        assert_eq!(coeff(&seq1, &seq2, CoefficientType::Weak), 2);
+        assert_eq!(weak_coeff(&seq1, &seq2), 2);
     }
 
     fn data(msg_id: u32) -> Entry {
@@ -295,4 +407,14 @@ mod tests {
     fn noise(count: u32) -> Entry {
         Entry::Noise(count)
     }
+
+    fn strong_coeff(seq1: &[Entry], seq2: &[Entry]) -> u64 {
+        let (strong, _) = strong_and_casual_coeff(seq1, seq2);
+        strong
+    }
+
+    fn casual_coeff(seq1: &[Entry], seq2: &[Entry]) -> u64 {
+        let (_, casual) = strong_and_casual_coeff(seq1, seq2);
+        casual
+    }
 }