Merge pull request #9 from TrAyZeN/cleanup

Cleanup and document

benches/snr.rs

@@ -1,5 +1,5 @@
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
-use muscat::processors::{compute_snr, Snr};
+use muscat::processors::{snr, Snr};
 use ndarray::{Array1, Array2};
 use ndarray_rand::rand::{rngs::StdRng, SeedableRng};
 use ndarray_rand::rand_distr::Uniform;
@@ -16,7 +16,7 @@ fn snr_sequential(leakages: &Array2<i64>, plaintexts: &Array2<u8>) -> Array1<f64
 }
 
 fn snr_parallel(leakages: &Array2<i64>, plaintexts: &Array2<u8>) -> Array1<f64> {
-    compute_snr(leakages.view(), 256, |i| plaintexts.row(i)[0].into(), 500)
+    snr(leakages.view(), 256, |i| plaintexts.row(i)[0].into(), 500)
 }
 
 fn bench_snr(c: &mut Criterion) {

src/cpa.rs

@@ -1,3 +1,4 @@
+use crate::util::max_per_row;
 use ndarray::{concatenate, s, Array1, Array2, ArrayView1, ArrayView2, Axis};
 use rayon::{
     iter::ParallelBridge,
@@ -10,14 +11,17 @@ use std::{iter::zip, ops::Add};
 /// # Panics
 /// - Panic if `leakages.shape()[0] != plaintexts.shape()[0]`
 /// - Panic if `chunk_size` is 0.
-pub fn cpa(
-    leakages: ArrayView2<usize>,
-    plaintexts: ArrayView2<usize>,
+pub fn cpa<T>(
+    leakages: ArrayView2<T>,
+    plaintexts: ArrayView2<T>,
     guess_range: usize,
     target_byte: usize,
     leakage_func: fn(usize, usize) -> usize,
     chunk_size: usize,
-) -> Cpa {
+) -> Cpa
+where
+    T: Into<usize> + Copy + Sync,
+{
     assert_eq!(leakages.shape()[0], plaintexts.shape()[0]);
     assert!(chunk_size > 0);
 
@@ -46,20 +50,27 @@ pub fn cpa(
 }
 
 pub struct Cpa {
-    sum_leakages: Array1<usize>,
-    sig_leakages: Array1<usize>,
-    sum_keys: Array1<usize>,
-    sig_keys: Array1<usize>,
-    values: Array1<usize>,
-    a_l: Array2<usize>,
+    /// Number of samples per trace
+    len_samples: usize,
     target_byte: usize,
-    len_leakages: usize,
+    /// Guess range upper excluded bound
     guess_range: usize,
+    /// Sum of traces
+    sum_leakages: Array1<usize>,
+    /// Sum of square of traces
+    sum_squares_leakages: Array1<usize>,
+    /// Sum of traces per key guess
+    guess_sum_leakages: Array1<usize>,
+    /// Sum of square of traces per key guess
+    guess_sum_squares_leakages: Array1<usize>,
+    a_l: Array2<usize>,
     corr: Array2<f32>,
-    max_corr: Array2<f32>,
+    max_corr: Array1<f32>,
     rank_slice: Array2<f32>,
+    /// Leakage model
     leakage_func: fn(usize, usize) -> usize,
-    len_samples: usize,
+    /// Number of traces processed
+    len_leakages: usize,
 }
 
 /* This class implements the CPA shown in this paper: https://eprint.iacr.org/2013/794.pdf */
@@ -72,125 +83,106 @@ impl Cpa {
     ) -> Self {
         Self {
             len_samples: size,
-            a_l: Array2::zeros((guess_range, size)),
             target_byte,
             guess_range,
             sum_leakages: Array1::zeros(size),
-            sig_leakages: Array1::zeros(size),
-            sum_keys: Array1::zeros(guess_range),
-            sig_keys: Array1::zeros(guess_range),
-            values: Array1::zeros(guess_range),
+            sum_squares_leakages: Array1::zeros(size),
+            guess_sum_leakages: Array1::zeros(guess_range),
+            guess_sum_squares_leakages: Array1::zeros(guess_range),
+            a_l: Array2::zeros((guess_range, size)),
             corr: Array2::zeros((guess_range, size)),
-            max_corr: Array2::zeros((guess_range, 1)),
+            max_corr: Array1::zeros(guess_range),
             rank_slice: Array2::zeros((guess_range, 1)),
             leakage_func,
             len_leakages: 0,
         }
     }
 
-    pub fn update(&mut self, trace: ArrayView1<usize>, plaintext: ArrayView1<usize>) {
+    /// # Panics
+    /// Panic in debug if `trace.shape()[0] != self.len_samples`.
+    pub fn update<T>(&mut self, trace: ArrayView1<T>, plaintext: ArrayView1<T>)
+    where
+        T: Into<usize> + Copy,
+    {
+        debug_assert_eq!(trace.shape()[0], self.len_samples);
+
         /* This function updates the main arrays of the CPA, as shown in Alg. 4
         in the paper.*/
-        self.len_leakages += 1;
-        self.gen_values(plaintext, self.guess_range, self.target_byte);
-        self.go(trace, plaintext, self.guess_range);
-    }
-
-    pub fn gen_values(
-        &mut self,
-        metadata: ArrayView1<usize>,
-        guess_range: usize,
-        target_key: usize,
-    ) {
-        for guess in 0..guess_range {
-            self.values[guess] = (self.leakage_func)(metadata[target_key], guess);
-        }
-    }
 
-    pub fn go(
-        &mut self,
-        trace: ArrayView1<usize>,
-        metadata: ArrayView1<usize>,
-        guess_range: usize,
-    ) {
         for i in 0..self.len_samples {
-            self.sum_leakages[i] += trace[i];
-            self.sig_leakages[i] += trace[i] * trace[i];
+            self.sum_leakages[i] += trace[i].into();
+            self.sum_squares_leakages[i] += trace[i].into() * trace[i].into();
         }
 
-        for guess in 0..guess_range {
-            self.sum_keys[guess] += self.values[guess];
-            self.sig_keys[guess] += self.values[guess] * self.values[guess];
+        for guess in 0..self.guess_range {
+            let value = (self.leakage_func)(plaintext[self.target_byte].into(), guess);
+            self.guess_sum_leakages[guess] += value;
+            self.guess_sum_squares_leakages[guess] += value * value;
         }
-        let partition: usize = metadata[self.target_byte];
+
+        let partition = plaintext[self.target_byte].into();
         for i in 0..self.len_samples {
-            self.a_l[[partition, i]] += trace[i];
+            self.a_l[[partition, i]] += trace[i].into();
         }
+
+        self.len_leakages += 1;
     }
 
     pub fn finalize(&mut self) {
         /* This function finalizes the calculation after feeding the
         overall traces */
 
-        let shape_p = self.guess_range;
-        let mut p = Array2::zeros((shape_p, shape_p));
-        for i in 0..self.guess_range {
+        let mut p = Array2::zeros((self.guess_range, self.guess_range));
+        for guess in 0..self.guess_range {
             for x in 0..self.guess_range {
-                p[[x, i]] = (self.leakage_func)(x, i);
+                p[[x, guess]] = (self.leakage_func)(x, guess);
             }
         }
-        for i in 0..self.guess_range {
-            let _sigkeys = self.sig_keys[i] as f32 / self.len_leakages as f32;
-            let _sumkeys = self.sum_keys[i] as f32 / self.len_leakages as f32;
-            let lower1 = _sigkeys - (_sumkeys * _sumkeys);
+
+        for guess in 0..self.guess_range {
+            let mean_key = self.guess_sum_leakages[guess] as f32 / self.len_leakages as f32;
+            let mean_squares_key =
+                self.guess_sum_squares_leakages[guess] as f32 / self.len_leakages as f32;
+            let var_key = mean_squares_key - (mean_key * mean_key);
 
             /* Parallel operation using multi-threading */
             let tmp: Vec<f32> = (0..self.len_samples)
                 .into_par_iter()
                 .map(|x| {
-                    let _sumleakages = self.sum_leakages[x] as f32 / self.len_leakages as f32;
-                    let _sigleakages = self.sig_leakages[x] as f32 / self.len_leakages as f32;
-                    let slice_a = self.a_l.slice(s![.., x]);
-                    let slice_b = p.slice(s![.., i]);
-                    let summult: i32 = self.sum_mult(slice_a, slice_b);
-                    let upper1: f32 = summult as f32 / self.len_leakages as f32;
-                    let upper: f32 = upper1 - (_sumkeys * _sumleakages);
-                    let lower2: f32 = _sigleakages - (_sumleakages * _sumleakages);
-                    let lower = f32::sqrt(lower1 * lower2);
+                    let mean_leakages = self.sum_leakages[x] as f32 / self.len_leakages as f32;
+                    let summult = self.sum_mult(self.a_l.slice(s![.., x]), p.slice(s![.., guess]));
+                    let upper1 = summult as f32 / self.len_leakages as f32;
+                    let upper = upper1 - (mean_key * mean_leakages);
+
+                    let mean_squares_leakages =
+                        self.sum_squares_leakages[x] as f32 / self.len_leakages as f32;
+                    let var_leakages = mean_squares_leakages - (mean_leakages * mean_leakages);
+                    let lower = f32::sqrt(var_key * var_leakages);
+
                     f32::abs(upper / lower)
                 })
                 .collect();
 
             #[allow(clippy::needless_range_loop)]
-            for z in 0..self.len_samples {
-                self.corr[[i, z]] = tmp[z];
+            for u in 0..self.len_samples {
+                self.corr[[guess, u]] = tmp[u];
             }
         }
-        self.calculation();
-    }
 
-    pub fn calculation(&mut self) {
-        // let mut max_256: Array2<f32> = Array2::zeros((self.guess_range as usize, 1));
-        for i in 0..self.guess_range {
-            let row = self.corr.row(i);
-            // Calculating the max value in the row
-            let max_value = row
-                .into_iter()
-                .reduce(|a, b| {
-                    let mut tmp = a;
-                    if tmp < b {
-                        tmp = b;
-                    }
-                    tmp
-                })
-                .unwrap();
-            self.max_corr[[i, 0]] = *max_value;
-        }
-        self.rank_slice = concatenate![Axis(1), self.rank_slice, self.max_corr];
+        self.max_corr = max_per_row(self.corr.view());
+
+        self.rank_slice = concatenate![
+            Axis(1),
+            self.rank_slice,
+            self.max_corr
+                .clone()
+                .into_shape((self.max_corr.shape()[0], 1))
+                .unwrap()
+        ];
     }
 
     pub fn pass_rank(&self) -> ArrayView2<f32> {
-        self.rank_slice.slice(s![.., 1..])
+        self.rank_slice.view()
     }
 
     pub fn pass_corr_array(&self) -> ArrayView2<f32> {
@@ -202,8 +194,8 @@ impl Cpa {
         let mut guess = 0;
 
         for i in 0..self.guess_range {
-            if self.max_corr[[i, 0]] > init_value {
-                init_value = self.max_corr[[i, 0]];
+            if self.max_corr[i] > init_value {
+                init_value = self.max_corr[i];
                 guess = i;
             }
         }
@@ -220,21 +212,26 @@ impl Add for Cpa {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
+        debug_assert_eq!(self.target_byte, rhs.target_byte);
+        debug_assert_eq!(self.guess_range, rhs.guess_range);
+        debug_assert_eq!(self.len_samples, rhs.len_samples);
+        debug_assert_eq!(self.leakage_func, rhs.leakage_func);
+
         Self {
+            len_samples: self.len_samples,
+            target_byte: self.target_byte,
+            guess_range: self.guess_range,
             sum_leakages: self.sum_leakages + rhs.sum_leakages,
-            sig_leakages: self.sig_leakages + rhs.sig_leakages,
-            sum_keys: self.sum_keys + rhs.sum_keys,
-            sig_keys: self.sig_keys + rhs.sig_keys,
-            values: self.values + rhs.values,
+            sum_squares_leakages: self.sum_squares_leakages + rhs.sum_squares_leakages,
+            guess_sum_leakages: self.guess_sum_leakages + rhs.guess_sum_leakages,
+            guess_sum_squares_leakages: self.guess_sum_squares_leakages
+                + rhs.guess_sum_squares_leakages,
             a_l: self.a_l + rhs.a_l,
-            target_byte: rhs.target_byte,
-            len_leakages: self.len_leakages + rhs.len_leakages,
-            guess_range: rhs.guess_range,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: rhs.len_samples,
             leakage_func: self.leakage_func,
+            len_leakages: self.len_leakages + rhs.len_leakages,
         }
     }
 }