From 6a0ba33a47db4570df6aa8f15eeaca362bcbe155 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Fri, 12 Apr 2024 15:47:24 +0200
Subject: [PATCH 1/9] Add debug assertions

---
 src/cpa.rs        | 15 ++++++++++++---
 src/cpa_normal.rs | 20 +++++++++++++++-----
 src/dpa.rs        | 15 ++++++++++++---
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index b59af8d..338a18a 100644
--- a/src/cpa.rs
+++ b/src/cpa.rs
@@ -88,7 +88,11 @@ impl Cpa {
         }
     }
 
+    /// # Panics
+    /// Panic in debug if `trace.shape()[0] != self.len_samples`.
     pub fn update(&mut self, trace: ArrayView1<usize>, plaintext: ArrayView1<usize>) {
+        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;
@@ -220,6 +224,11 @@ 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 {
             sum_leakages: self.sum_leakages + rhs.sum_leakages,
             sig_leakages: self.sig_leakages + rhs.sig_leakages,
@@ -227,13 +236,13 @@ impl Add for Cpa {
             sig_keys: self.sig_keys + rhs.sig_keys,
             values: self.values + rhs.values,
             a_l: self.a_l + rhs.a_l,
-            target_byte: rhs.target_byte,
+            target_byte: self.target_byte,
             len_leakages: self.len_leakages + rhs.len_leakages,
-            guess_range: rhs.guess_range,
+            guess_range: self.guess_range,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: rhs.len_samples,
+            len_samples: self.len_samples,
             leakage_func: self.leakage_func,
         }
     }
diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index 1bf04dc..eef3796 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -89,17 +89,23 @@ impl Cpa {
         }
     }
 
+    /// # Panics
+    /// - Panic in debug if `trace_batch.shape()[0] != plaintext_batch.shape()[0]`.
+    /// - Panic in debug if `trace_batch.shape()[1] != self.len_samples`.
     pub fn update<T, U>(&mut self, trace_batch: ArrayView2<T>, plaintext_batch: ArrayView2<U>)
     where
         T: Into<f32> + Copy,
         U: Into<usize> + Copy,
     {
+        debug_assert_eq!(trace_batch.shape()[0], plaintext_batch.shape()[0]);
+        debug_assert_eq!(trace_batch.shape()[1], self.len_samples);
+
         /* This function updates the internal arrays of the CPA
         It accepts trace_batch and plaintext_batch to update them*/
         let trace_batch = trace_batch.map(|&t| t.into());
         let plaintext_batch = plaintext_batch.map(|&m| m.into());
 
-        self.update_values(plaintext_batch.view(), trace_batch.view(), self.guess_range);
+        self.update_values(trace_batch.view(), plaintext_batch.view(), self.guess_range);
         self.update_key_leakages(trace_batch.view(), self.guess_range);
 
         self.len_leakages += self.chunk;
@@ -108,8 +114,8 @@ impl Cpa {
     pub fn update_values(
         /* This function generates the values and cov arrays */
         &mut self,
-        metadata: ArrayView2<usize>,
         trace: ArrayView2<f32>,
+        metadata: ArrayView2<usize>,
         guess_range: usize,
     ) {
         for row in 0..self.chunk {
@@ -227,6 +233,10 @@ impl Add for Cpa {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
+        debug_assert_eq!(self.len_samples, rhs.len_samples);
+        debug_assert_eq!(self.chunk, rhs.chunk);
+        debug_assert_eq!(self.guess_range, rhs.guess_range);
+
         Self {
             sum_leakages: self.sum_leakages + rhs.sum_leakages,
             sum2_leakages: self.sum2_leakages + rhs.sum2_leakages,
@@ -234,13 +244,13 @@ impl Add for Cpa {
             sum2_keys: self.sum2_keys + rhs.sum2_keys,
             values: self.values + rhs.values,
             len_leakages: self.len_leakages + rhs.len_leakages,
-            guess_range: rhs.guess_range,
-            chunk: rhs.chunk,
+            guess_range: self.guess_range,
+            chunk: self.chunk,
             cov: self.cov + rhs.cov,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: rhs.len_samples,
+            len_samples: self.len_samples,
             leakage_func: self.leakage_func,
             rank_traces: self.rank_traces,
         }
diff --git a/src/dpa.rs b/src/dpa.rs
index 1254758..63bb0c7 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -1,5 +1,6 @@
 use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
 use std::ops::Add;
+
 pub struct Dpa<T> {
     /* List of internal class variables */
     sum_0: Array2<f32>,
@@ -38,11 +39,14 @@ impl<T: Clone> Dpa<T> {
         }
     }
 
-    //
+    /// # Panics
+    /// Panic in debug if `trace.shape()[0] != self.len_samples`.
     pub fn update<U>(&mut self, trace: ArrayView1<U>, metadata: T)
     where
         U: Into<f32> + Copy,
     {
+        debug_assert_eq!(trace.shape()[0], self.len_samples);
+
         /* This function updates the internal arrays of the DPA
         It accepts trace_batch and plaintext_batch to update them*/
         for guess in 0..self.guess_range as i16 {
@@ -85,6 +89,7 @@ impl<T: Clone> Dpa<T> {
             }
         }
     }
+
     pub fn assign_rank_traces(&mut self, value: usize) {
         self.rank_traces = value;
     }
@@ -154,16 +159,20 @@ impl<T> Add for Dpa<T> {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
+        debug_assert_eq!(self.len_samples, rhs.len_samples);
+        debug_assert_eq!(self.guess_range, rhs.guess_range);
+        debug_assert_eq!(self.leakage_func, rhs.leakage_func);
+
         Self {
             sum_0: self.sum_0 + rhs.sum_0,
             sum_1: self.sum_1 + rhs.sum_1,
             count_0: self.count_0 + rhs.count_0,
             count_1: self.count_1 + rhs.count_1,
-            guess_range: rhs.guess_range,
+            guess_range: self.guess_range,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: rhs.len_samples,
+            len_samples: self.len_samples,
             leakage_func: self.leakage_func,
             rank_traces: self.rank_traces,
             len_leakages: self.len_leakages + rhs.len_leakages,

From 214f604e87562dd78eb81bb0ce5f06d72bd68a32 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Fri, 12 Apr 2024 15:47:24 +0200
Subject: [PATCH 2/9] Make internal functions private

---
 src/cpa.rs        | 16 +++-------------
 src/cpa_normal.rs |  6 +++---
 src/dpa.rs        |  2 +-
 3 files changed, 7 insertions(+), 17 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index 338a18a..24468b6 100644
--- a/src/cpa.rs
+++ b/src/cpa.rs
@@ -100,23 +100,13 @@ impl Cpa {
         self.go(trace, plaintext, self.guess_range);
     }
 
-    pub fn gen_values(
-        &mut self,
-        metadata: ArrayView1<usize>,
-        guess_range: usize,
-        target_key: usize,
-    ) {
+    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,
-    ) {
+    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];
@@ -173,7 +163,7 @@ impl Cpa {
         self.calculation();
     }
 
-    pub fn calculation(&mut self) {
+    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);
diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index eef3796..0a7f7a3 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -111,7 +111,7 @@ impl Cpa {
         self.len_leakages += self.chunk;
     }
 
-    pub fn update_values(
+    fn update_values(
         /* This function generates the values and cov arrays */
         &mut self,
         trace: ArrayView2<f32>,
@@ -128,7 +128,7 @@ impl Cpa {
         self.cov = self.cov.clone() + self.values.t().dot(&trace);
     }
 
-    pub fn update_key_leakages(&mut self, trace: ArrayView2<f32>, guess_range: usize) {
+    fn update_key_leakages(&mut self, trace: ArrayView2<f32>, guess_range: usize) {
         for i in 0..self.len_samples {
             self.sum_leakages[i] += trace.column(i).sum(); // trace[i] as usize;
             self.sum2_leakages[i] += trace.column(i).dot(&trace.column(i)); // (trace[i] * trace[i]) as usize;
@@ -186,7 +186,7 @@ impl Cpa {
         self.select_max();
     }
 
-    pub fn select_max(&mut self) {
+    fn select_max(&mut self) {
         for i in 0..self.guess_range {
             let row = self.corr.row(i);
             // Calculating the max value in the row
diff --git a/src/dpa.rs b/src/dpa.rs
index 63bb0c7..225b21a 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -112,7 +112,7 @@ impl<T: Clone> Dpa<T> {
         self.select_max();
     }
 
-    pub fn select_max(&mut self) {
+    fn select_max(&mut self) {
         for i in 0..self.guess_range {
             let row = self.corr.row(i as usize);
             // Calculating the max value in the row

From c4026f2d367e62b2627a9d258f4955e7f9793218 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Fri, 12 Apr 2024 15:47:24 +0200
Subject: [PATCH 3/9] Cleanup cpa code

---
 src/cpa.rs | 144 ++++++++++++++++++++++++++++-------------------------
 1 file changed, 77 insertions(+), 67 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index 24468b6..19fb37c 100644
--- a/src/cpa.rs
+++ b/src/cpa.rs
@@ -46,20 +46,25 @@ 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 observations per trace
+    len_samples: usize,
     target_byte: usize,
-    len_leakages: usize,
     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,
+    len_leakages: usize,
 }
 
 /* This class implements the CPA shown in this paper: https://eprint.iacr.org/2013/794.pdf */
@@ -72,16 +77,15 @@ 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,
@@ -95,96 +99,102 @@ impl Cpa {
 
         /* 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);
-    }
-
-    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);
-        }
-    }
 
-    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_squares_leakages[i] += trace[i] * trace[i];
         }
 
-        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], 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];
         for i in 0..self.len_samples {
             self.a_l[[partition, i]] += trace[i];
         }
+
+        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();
     }
 
     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);
+        for guess in 0..self.guess_range {
             // Calculating the max value in the row
-            let max_value = row
+            let max_value = self
+                .corr
+                .row(guess)
                 .into_iter()
                 .reduce(|a, b| {
-                    let mut tmp = a;
-                    if tmp < b {
-                        tmp = b;
+                    let mut max = a;
+                    if max < b {
+                        max = b;
                     }
-                    tmp
+
+                    max
                 })
                 .unwrap();
-            self.max_corr[[i, 0]] = *max_value;
+
+            self.max_corr[guess] = *max_value;
         }
-        self.rank_slice = concatenate![Axis(1), self.rank_slice, self.max_corr];
+
+        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> {
@@ -196,8 +206,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;
             }
         }
@@ -221,10 +231,10 @@ impl Add for Cpa {
 
         Self {
             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: self.target_byte,
             len_leakages: self.len_leakages + rhs.len_leakages,

From 743d235f5bbb89b61fcfb76b81e85cdce99d91be Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Fri, 12 Apr 2024 15:47:24 +0200
Subject: [PATCH 4/9] Make cpa update generic

---
 src/cpa.rs | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index 19fb37c..83ef33c 100644
--- a/src/cpa.rs
+++ b/src/cpa.rs
@@ -10,14 +10,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);
 
@@ -94,26 +97,29 @@ impl Cpa {
 
     /// # Panics
     /// Panic in debug if `trace.shape()[0] != self.len_samples`.
-    pub fn update(&mut self, trace: ArrayView1<usize>, plaintext: ArrayView1<usize>) {
+    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.*/
 
         for i in 0..self.len_samples {
-            self.sum_leakages[i] += trace[i];
-            self.sum_squares_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..self.guess_range {
-            let value = (self.leakage_func)(plaintext[self.target_byte], guess);
+            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 = plaintext[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;

From 251989c1d195658ce8cc8caf9c832b756f777c4b Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Thu, 18 Apr 2024 10:52:46 +0200
Subject: [PATCH 5/9] Add util function to compute maximum per row

---
 src/cpa.rs        | 24 ++----------------------
 src/cpa_normal.rs | 23 +++++------------------
 src/dpa.rs        | 21 ++++-----------------
 src/util.rs       | 19 ++++++++++++++++++-
 4 files changed, 29 insertions(+), 58 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index 83ef33c..886fd24 100644
--- a/src/cpa.rs
+++ b/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,
@@ -166,28 +167,7 @@ impl Cpa {
             }
         }
 
-        self.calculation();
-    }
-
-    fn calculation(&mut self) {
-        for guess in 0..self.guess_range {
-            // Calculating the max value in the row
-            let max_value = self
-                .corr
-                .row(guess)
-                .into_iter()
-                .reduce(|a, b| {
-                    let mut max = a;
-                    if max < b {
-                        max = b;
-                    }
-
-                    max
-                })
-                .unwrap();
-
-            self.max_corr[guess] = *max_value;
-        }
+        self.max_corr = max_per_row(self.corr.view());
 
         self.rank_slice = concatenate![
             Axis(1),
diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index 0a7f7a3..891d702 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -1,7 +1,9 @@
-use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
+use ndarray::{concatenate, s, Array1, Array2, ArrayView1, ArrayView2, Axis};
 use rayon::iter::{ParallelBridge, ParallelIterator};
 use std::{iter::zip, ops::Add};
 
+use crate::util::max_per_row;
+
 /// Computes the [`Cpa`] of the given traces.
 ///
 /// # Panics
@@ -183,24 +185,9 @@ impl Cpa {
                 }
             }
         }
-        self.select_max();
-    }
 
-    fn select_max(&mut self) {
-        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;
+        for (i, max) in max_per_row(self.corr.view()).into_iter().enumerate() {
+            self.max_corr[[i, 0]] = max;
         }
     }
 
diff --git a/src/dpa.rs b/src/dpa.rs
index 225b21a..5755493 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -1,6 +1,8 @@
 use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
 use std::ops::Add;
 
+use crate::util::max_per_row;
+
 pub struct Dpa<T> {
     /* List of internal class variables */
     sum_0: Array2<f32>,
@@ -109,24 +111,9 @@ impl<T: Clone> Dpa<T> {
         let diff = tmp_avg_0.clone() - tmp_avg_1;
 
         self.corr = diff.map(|e| f32::abs(*e));
-        self.select_max();
-    }
 
-    fn select_max(&mut self) {
-        for i in 0..self.guess_range {
-            let row = self.corr.row(i as usize);
-            // 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 as usize, 0]] = *max_value;
+        for (i, max) in max_per_row(self.corr.view()).into_iter().enumerate() {
+            self.max_corr[[i, 0]] = max;
         }
     }
 
diff --git a/src/util.rs b/src/util.rs
index dc8316f..27a6b47 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -3,7 +3,7 @@
 use std::{io::Read, path::Path, time::Duration};
 
 use indicatif::{ProgressBar, ProgressStyle};
-use ndarray::{Array, Array1, Array2, ArrayView2};
+use ndarray::{Array, Array1, Array2, ArrayView2, Axis};
 use ndarray_npy::{read_npy, write_npy, ReadNpyError, ReadableElement, WriteNpyError};
 use npyz::{Deserialize, NpyFile};
 
@@ -56,3 +56,20 @@ pub fn read_array2_from_npy_file<T: ReadableElement>(
 pub fn save_array2(path: impl AsRef<Path>, array: ArrayView2<f32>) -> Result<(), WriteNpyError> {
     write_npy(path, &array)
 }
+
+/// Return an array where the i-th element contains the maximum of the i-th row of the input array.
+pub fn max_per_row(arr: ArrayView2<f32>) -> Array1<f32> {
+    arr.axis_iter(Axis(0))
+        .map(|row| {
+            *row.into_iter()
+                .reduce(|a, b| {
+                    let mut tmp = a;
+                    if tmp < b {
+                        tmp = b;
+                    }
+                    tmp
+                })
+                .unwrap()
+        })
+        .collect()
+}

From f9f1bd0710755da0072aaf1fe87b01058909a599 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Thu, 18 Apr 2024 11:04:14 +0200
Subject: [PATCH 6/9] Remove unneeded dimension of max_corr

---
 src/cpa_normal.rs | 31 +++++++++++++++++++++----------
 src/dpa.rs        | 27 ++++++++++++++++++---------
 2 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index 891d702..3b5bd56 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -1,4 +1,4 @@
-use ndarray::{concatenate, s, Array1, Array2, ArrayView1, ArrayView2, Axis};
+use ndarray::{concatenate, Array1, Array2, ArrayView1, ArrayView2, Axis};
 use rayon::iter::{ParallelBridge, ParallelIterator};
 use std::{iter::zip, ops::Add};
 
@@ -54,7 +54,7 @@ pub struct Cpa {
     guess_range: usize,
     cov: Array2<f32>,
     corr: Array2<f32>,
-    max_corr: Array2<f32>,
+    max_corr: Array1<f32>,
     rank_slice: Array2<f32>,
     leakage_func: fn(ArrayView1<usize>, usize) -> usize,
     len_samples: usize,
@@ -83,7 +83,7 @@ impl Cpa {
             values: Array2::zeros((batch, guess_range)),
             cov: 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,
@@ -156,9 +156,20 @@ impl Cpa {
         if self.len_leakages % self.rank_traces == 0 {
             self.finalize();
             if self.len_leakages == self.rank_traces {
-                self.rank_slice = self.max_corr.clone();
+                self.rank_slice = self
+                    .max_corr
+                    .clone()
+                    .into_shape((self.max_corr.shape()[0], 1))
+                    .unwrap();
             } else {
-                self.rank_slice = concatenate![Axis(1), self.rank_slice, self.max_corr];
+                self.rank_slice = concatenate![
+                    Axis(1),
+                    self.rank_slice,
+                    self.max_corr
+                        .clone()
+                        .into_shape((self.max_corr.shape()[0], 1))
+                        .unwrap()
+                ];
             }
         }
     }
@@ -186,9 +197,7 @@ impl Cpa {
             }
         }
 
-        for (i, max) in max_per_row(self.corr.view()).into_iter().enumerate() {
-            self.max_corr[[i, 0]] = max;
-        }
+        self.max_corr = max_per_row(self.corr.view());
     }
 
     pub fn success_traces(&mut self, traces_no: usize) {
@@ -206,12 +215,14 @@ impl Cpa {
     pub fn pass_guess(&self) -> usize {
         let mut init_value = 0.0;
         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;
             }
         }
+
         guess
     }
 }
diff --git a/src/dpa.rs b/src/dpa.rs
index 5755493..cf245b2 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -11,7 +11,7 @@ pub struct Dpa<T> {
     count_1: Array1<usize>,
     guess_range: i32,
     corr: Array2<f32>,
-    max_corr: Array2<f32>,
+    max_corr: Array1<f32>,
     rank_slice: Array2<f32>,
     leakage_func: fn(T, usize) -> usize,
     len_samples: usize,
@@ -33,7 +33,7 @@ impl<T: Clone> Dpa<T> {
             count_0: Array1::zeros(guess_range as usize),
             count_1: Array1::zeros(guess_range as usize),
             corr: Array2::zeros((guess_range as usize, size)),
-            max_corr: Array2::zeros((guess_range as usize, 1)),
+            max_corr: Array1::zeros(guess_range as usize),
             rank_slice: Array2::zeros((guess_range as usize, 1)),
             leakage_func: f,
             rank_traces: 0,
@@ -85,9 +85,20 @@ impl<T: Clone> Dpa<T> {
             self.finalize();
 
             if self.len_leakages == self.rank_traces {
-                self.rank_slice = self.max_corr.clone();
+                self.rank_slice = self
+                    .max_corr
+                    .clone()
+                    .into_shape((self.max_corr.shape()[0], 1))
+                    .unwrap();
             } else {
-                self.rank_slice = concatenate![Axis(1), self.rank_slice, self.max_corr];
+                self.rank_slice = concatenate![
+                    Axis(1),
+                    self.rank_slice,
+                    self.max_corr
+                        .clone()
+                        .into_shape((self.max_corr.shape()[0], 1))
+                        .unwrap()
+                ];
             }
         }
     }
@@ -112,9 +123,7 @@ impl<T: Clone> Dpa<T> {
 
         self.corr = diff.map(|e| f32::abs(*e));
 
-        for (i, max) in max_per_row(self.corr.view()).into_iter().enumerate() {
-            self.max_corr[[i, 0]] = max;
-        }
+        self.max_corr = max_per_row(self.corr.view());
     }
 
     pub fn success_traces(&mut self, traces_no: usize) {
@@ -133,8 +142,8 @@ impl<T: Clone> Dpa<T> {
         let mut init_value: f32 = 0.0;
         let mut guess: i32 = 0;
         for i in 0..self.guess_range {
-            if self.max_corr[[i as usize, 0]] > init_value {
-                init_value = self.max_corr[[i as usize, 0]];
+            if self.max_corr[i as usize] > init_value {
+                init_value = self.max_corr[i as usize];
                 guess = i;
             }
         }

From a3d994c9232a358801dabbad59d09ec80e425d04 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Thu, 18 Apr 2024 12:50:10 +0200
Subject: [PATCH 7/9] Rename SNR helper function

---
 benches/snr.rs    | 4 ++--
 src/cpa_normal.rs | 1 -
 src/dpa.rs        | 1 -
 src/processors.rs | 2 +-
 4 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/benches/snr.rs b/benches/snr.rs
index ee3da8d..e6f5a74 100644
--- a/benches/snr.rs
+++ b/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) {
diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index 3b5bd56..e56e7fb 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -44,7 +44,6 @@ where
 }
 
 pub struct Cpa {
-    /* List of internal class variables */
     sum_leakages: Array1<f32>,
     sum2_leakages: Array1<f32>,
     sum_keys: Array1<f32>,
diff --git a/src/dpa.rs b/src/dpa.rs
index cf245b2..7d4f877 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -4,7 +4,6 @@ use std::ops::Add;
 use crate::util::max_per_row;
 
 pub struct Dpa<T> {
-    /* List of internal class variables */
     sum_0: Array2<f32>,
     sum_1: Array2<f32>,
     count_0: Array1<usize>,
diff --git a/src/processors.rs b/src/processors.rs
index aaa84d7..98846ed 100644
--- a/src/processors.rs
+++ b/src/processors.rs
@@ -85,7 +85,7 @@ impl Add for MeanVar {
 ///
 /// # Panics
 /// Panic if `chunk_size` is 0.
-pub fn compute_snr<T, F>(
+pub fn snr<T, F>(
     leakages: ArrayView2<T>,
     classes: usize,
     get_class: F,

From 080fa17ed8ad3de40d87a1a7228cf84b298c0049 Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Thu, 18 Apr 2024 13:10:56 +0200
Subject: [PATCH 8/9] Document {Cpa,Dpa} fields

---
 src/cpa.rs        | 12 +++++----
 src/cpa_normal.rs | 64 +++++++++++++++++++++++++++--------------------
 src/dpa.rs        | 24 ++++++++++++------
 3 files changed, 60 insertions(+), 40 deletions(-)

diff --git a/src/cpa.rs b/src/cpa.rs
index 886fd24..ac2e870 100644
--- a/src/cpa.rs
+++ b/src/cpa.rs
@@ -50,9 +50,10 @@ where
 }
 
 pub struct Cpa {
-    /// Number of observations per trace
+    /// Number of samples per trace
     len_samples: usize,
     target_byte: usize,
+    /// Guess range upper excluded bound
     guess_range: usize,
     /// Sum of traces
     sum_leakages: Array1<usize>,
@@ -68,6 +69,7 @@ pub struct Cpa {
     rank_slice: Array2<f32>,
     /// Leakage model
     leakage_func: fn(usize, usize) -> usize,
+    /// Number of traces processed
     len_leakages: usize,
 }
 
@@ -216,20 +218,20 @@ impl Add for Cpa {
         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,
             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: self.target_byte,
-            len_leakages: self.len_leakages + rhs.len_leakages,
-            guess_range: self.guess_range,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: self.len_samples,
             leakage_func: self.leakage_func,
+            len_leakages: self.len_leakages + rhs.len_leakages,
         }
     }
 }
diff --git a/src/cpa_normal.rs b/src/cpa_normal.rs
index e56e7fb..98a7a60 100644
--- a/src/cpa_normal.rs
+++ b/src/cpa_normal.rs
@@ -44,21 +44,30 @@ where
 }
 
 pub struct Cpa {
+    /// Number of samples per trace
+    len_samples: usize,
+    /// Guess range upper excluded bound
+    guess_range: usize,
+    /// Sum of traces
     sum_leakages: Array1<f32>,
+    /// Sum of square of traces
     sum2_leakages: Array1<f32>,
-    sum_keys: Array1<f32>,
-    sum2_keys: Array1<f32>,
+    /// Sum of traces per key guess
+    guess_sum_leakages: Array1<f32>,
+    /// Sum of square of traces per key guess
+    guess_sum2_leakages: Array1<f32>,
     values: Array2<f32>,
-    len_leakages: usize,
-    guess_range: usize,
     cov: Array2<f32>,
     corr: Array2<f32>,
     max_corr: Array1<f32>,
     rank_slice: Array2<f32>,
-    leakage_func: fn(ArrayView1<usize>, usize) -> usize,
-    len_samples: usize,
-    chunk: usize,
     rank_traces: usize, // Number of traces to calculate succes rate
+    /// Batch size
+    batch_size: usize,
+    /// Leakage model
+    leakage_func: fn(ArrayView1<usize>, usize) -> usize,
+    /// Number of traces processed
+    len_leakages: usize,
 }
 
 /* This class implements the CPA algorithm shown in:
@@ -67,26 +76,26 @@ https://www.iacr.org/archive/ches2004/31560016/31560016.pdf */
 impl Cpa {
     pub fn new(
         size: usize,
-        batch: usize,
+        batch_size: usize,
         guess_range: usize,
         leakage_func: fn(ArrayView1<usize>, usize) -> usize,
     ) -> Self {
         Self {
             len_samples: size,
-            chunk: batch,
             guess_range,
             sum_leakages: Array1::zeros(size),
             sum2_leakages: Array1::zeros(size),
-            sum_keys: Array1::zeros(guess_range),
-            sum2_keys: Array1::zeros(guess_range),
-            values: Array2::zeros((batch, guess_range)),
+            guess_sum_leakages: Array1::zeros(guess_range),
+            guess_sum2_leakages: Array1::zeros(guess_range),
+            values: Array2::zeros((batch_size, guess_range)),
             cov: Array2::zeros((guess_range, size)),
             corr: Array2::zeros((guess_range, size)),
             max_corr: Array1::zeros(guess_range),
             rank_slice: Array2::zeros((guess_range, 1)),
+            rank_traces: 0,
+            batch_size,
             leakage_func,
             len_leakages: 0,
-            rank_traces: 0,
         }
     }
 
@@ -109,7 +118,7 @@ impl Cpa {
         self.update_values(trace_batch.view(), plaintext_batch.view(), self.guess_range);
         self.update_key_leakages(trace_batch.view(), self.guess_range);
 
-        self.len_leakages += self.chunk;
+        self.len_leakages += self.batch_size;
     }
 
     fn update_values(
@@ -119,7 +128,7 @@ impl Cpa {
         metadata: ArrayView2<usize>,
         guess_range: usize,
     ) {
-        for row in 0..self.chunk {
+        for row in 0..self.batch_size {
             for guess in 0..guess_range {
                 let pass_to_leakage = metadata.row(row);
                 self.values[[row, guess]] = (self.leakage_func)(pass_to_leakage, guess) as f32;
@@ -136,8 +145,9 @@ impl Cpa {
         }
 
         for guess in 0..guess_range {
-            self.sum_keys[guess] += self.values.column(guess).sum(); //self.values[guess] as usize;
-            self.sum2_keys[guess] += self.values.column(guess).dot(&self.values.column(guess));
+            self.guess_sum_leakages[guess] += self.values.column(guess).sum(); //self.values[guess] as usize;
+            self.guess_sum2_leakages[guess] +=
+                self.values.column(guess).dot(&self.values.column(guess));
             // (self.values[guess] * self.values[guess]) as usize;
         }
     }
@@ -177,8 +187,8 @@ impl Cpa {
         /* This function finalizes the calculation after
         feeding all stored acc arrays */
         let cov_n = self.cov.clone() / self.len_leakages as f32;
-        let avg_keys = self.sum_keys.clone() / self.len_leakages as f32;
-        let std_key = self.sum2_keys.clone() / self.len_leakages as f32;
+        let avg_keys = self.guess_sum_leakages.clone() / self.len_leakages as f32;
+        let std_key = self.guess_sum2_leakages.clone() / self.len_leakages as f32;
         let avg_leakages = self.sum_leakages.clone() / self.len_leakages as f32;
         let std_leakages = self.sum2_leakages.clone() / self.len_leakages as f32;
 
@@ -231,25 +241,25 @@ impl Add for Cpa {
 
     fn add(self, rhs: Self) -> Self::Output {
         debug_assert_eq!(self.len_samples, rhs.len_samples);
-        debug_assert_eq!(self.chunk, rhs.chunk);
+        debug_assert_eq!(self.batch_size, rhs.batch_size);
         debug_assert_eq!(self.guess_range, rhs.guess_range);
 
         Self {
+            len_samples: self.len_samples,
+            guess_range: self.guess_range,
             sum_leakages: self.sum_leakages + rhs.sum_leakages,
             sum2_leakages: self.sum2_leakages + rhs.sum2_leakages,
-            sum_keys: self.sum_keys + rhs.sum_keys,
-            sum2_keys: self.sum2_keys + rhs.sum2_keys,
+            guess_sum_leakages: self.guess_sum_leakages + rhs.guess_sum_leakages,
+            guess_sum2_leakages: self.guess_sum2_leakages + rhs.guess_sum2_leakages,
             values: self.values + rhs.values,
-            len_leakages: self.len_leakages + rhs.len_leakages,
-            guess_range: self.guess_range,
-            chunk: self.chunk,
             cov: self.cov + rhs.cov,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: self.len_samples,
-            leakage_func: self.leakage_func,
             rank_traces: self.rank_traces,
+            batch_size: self.batch_size,
+            leakage_func: self.leakage_func,
+            len_leakages: self.len_leakages + rhs.len_leakages,
         }
     }
 }
diff --git a/src/dpa.rs b/src/dpa.rs
index 7d4f877..3de07c2 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -4,17 +4,25 @@ use std::ops::Add;
 use crate::util::max_per_row;
 
 pub struct Dpa<T> {
+    /// Number of samples per trace
+    len_samples: usize,
+    /// Guess range upper excluded bound
+    guess_range: i32,
+    /// Sum of traces for which the selection function equals 0
     sum_0: Array2<f32>,
+    /// Sum of traces for which the selection function equals 1
     sum_1: Array2<f32>,
+    /// Number of traces processed for which the selection function equals 0
     count_0: Array1<usize>,
+    /// Number of traces processed for which the selection function equals 1
     count_1: Array1<usize>,
-    guess_range: i32,
     corr: Array2<f32>,
     max_corr: Array1<f32>,
     rank_slice: Array2<f32>,
-    leakage_func: fn(T, usize) -> usize,
-    len_samples: usize,
     rank_traces: usize, // Number of traces to calculate succes rate
+    /// Selection function
+    leakage_func: fn(T, usize) -> usize,
+    /// Number of traces processed
     len_leakages: usize,
 }
 
@@ -26,7 +34,7 @@ impl<T: Clone> Dpa<T> {
     pub fn new(size: usize, guess_range: i32, f: fn(T, usize) -> usize) -> Self {
         Self {
             len_samples: size,
-            guess_range, //fixing clippy warning
+            guess_range,
             sum_0: Array2::zeros((guess_range as usize, size)),
             sum_1: Array2::zeros((guess_range as usize, size)),
             count_0: Array1::zeros(guess_range as usize),
@@ -34,8 +42,8 @@ impl<T: Clone> Dpa<T> {
             corr: Array2::zeros((guess_range as usize, size)),
             max_corr: Array1::zeros(guess_range as usize),
             rank_slice: Array2::zeros((guess_range as usize, 1)),
-            leakage_func: f,
             rank_traces: 0,
+            leakage_func: f,
             len_leakages: 0,
         }
     }
@@ -159,17 +167,17 @@ impl<T> Add for Dpa<T> {
         debug_assert_eq!(self.leakage_func, rhs.leakage_func);
 
         Self {
+            len_samples: self.len_samples,
+            guess_range: self.guess_range,
             sum_0: self.sum_0 + rhs.sum_0,
             sum_1: self.sum_1 + rhs.sum_1,
             count_0: self.count_0 + rhs.count_0,
             count_1: self.count_1 + rhs.count_1,
-            guess_range: self.guess_range,
             corr: self.corr + rhs.corr,
             max_corr: self.max_corr,
             rank_slice: self.rank_slice,
-            len_samples: self.len_samples,
-            leakage_func: self.leakage_func,
             rank_traces: self.rank_traces,
+            leakage_func: self.leakage_func,
             len_leakages: self.len_leakages + rhs.len_leakages,
         }
     }

From 834562ee0d04bb4720b504dc018765b3f014ecbc Mon Sep 17 00:00:00 2001
From: TrAyZeN <1810leo@gmail.com>
Date: Thu, 18 Apr 2024 13:12:14 +0200
Subject: [PATCH 9/9] Make Dpa::guess_range usize

---
 src/dpa.rs | 56 ++++++++++++++++++++++++++++--------------------------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/src/dpa.rs b/src/dpa.rs
index 3de07c2..fdec0a8 100644
--- a/src/dpa.rs
+++ b/src/dpa.rs
@@ -7,7 +7,7 @@ pub struct Dpa<T> {
     /// Number of samples per trace
     len_samples: usize,
     /// Guess range upper excluded bound
-    guess_range: i32,
+    guess_range: usize,
     /// Sum of traces for which the selection function equals 0
     sum_0: Array2<f32>,
     /// Sum of traces for which the selection function equals 1
@@ -31,17 +31,17 @@ https://paulkocher.com/doc/DifferentialPowerAnalysis.pdf
 https://web.mit.edu/6.857/OldStuff/Fall03/ref/kocher-DPATechInfo.pdf */
 
 impl<T: Clone> Dpa<T> {
-    pub fn new(size: usize, guess_range: i32, f: fn(T, usize) -> usize) -> Self {
+    pub fn new(size: usize, guess_range: usize, f: fn(T, usize) -> usize) -> Self {
         Self {
             len_samples: size,
             guess_range,
-            sum_0: Array2::zeros((guess_range as usize, size)),
-            sum_1: Array2::zeros((guess_range as usize, size)),
-            count_0: Array1::zeros(guess_range as usize),
-            count_1: Array1::zeros(guess_range as usize),
-            corr: Array2::zeros((guess_range as usize, size)),
-            max_corr: Array1::zeros(guess_range as usize),
-            rank_slice: Array2::zeros((guess_range as usize, 1)),
+            sum_0: Array2::zeros((guess_range, size)),
+            sum_1: Array2::zeros((guess_range, size)),
+            count_0: Array1::zeros(guess_range),
+            count_1: Array1::zeros(guess_range),
+            corr: Array2::zeros((guess_range, size)),
+            max_corr: Array1::zeros(guess_range),
+            rank_slice: Array2::zeros((guess_range, 1)),
             rank_traces: 0,
             leakage_func: f,
             len_leakages: 0,
@@ -58,24 +58,24 @@ impl<T: Clone> Dpa<T> {
 
         /* This function updates the internal arrays of the DPA
         It accepts trace_batch and plaintext_batch to update them*/
-        for guess in 0..self.guess_range as i16 {
-            let index = (self.leakage_func)(metadata.clone(), guess as usize);
+        for guess in 0..self.guess_range {
+            let index = (self.leakage_func)(metadata.clone(), guess);
             if index & 1 == 1 {
                 // classification is performed based on the lsb
-                // let tmp_row: Array1<f32> = self.sum_1.row(guess as usize).to_owned() + tmp_trace.clone();
-                // self.sum_1.row_mut(guess as usize).assign(&tmp_row);
+                // let tmp_row: Array1<f32> = self.sum_1.row(guess).to_owned() + tmp_trace.clone();
+                // self.sum_1.row_mut(guess).assign(&tmp_row);
                 for i in 0..self.len_samples {
-                    self.sum_1[[guess as usize, i]] += trace[i].into();
+                    self.sum_1[[guess, i]] += trace[i].into();
                 }
-                self.count_1[guess as usize] += 1;
+                self.count_1[guess] += 1;
             } else {
-                // let tmp_row: Array1<f32> = self.sum_0.row(guess as usize).to_owned() + tmp_trace.clone();
-                // self.sum_0.row_mut(guess as usize).assign(&tmp_row);
+                // let tmp_row: Array1<f32> = self.sum_0.row(guess).to_owned() + tmp_trace.clone();
+                // self.sum_0.row_mut(guess).assign(&tmp_row);
 
                 for i in 0..self.len_samples {
-                    self.sum_0[[guess as usize, i]] += trace[i].into();
+                    self.sum_0[[guess, i]] += trace[i].into();
                 }
-                self.count_0[guess as usize] += 1;
+                self.count_0[guess] += 1;
             }
         }
         self.len_leakages += 1;
@@ -117,10 +117,10 @@ impl<T: Clone> Dpa<T> {
     pub fn finalize(&mut self) {
         /* This function finalizes the calculation after
         feeding all stored acc arrays */
-        let mut tmp_avg_0 = Array2::zeros((self.guess_range as usize, self.len_samples));
-        let mut tmp_avg_1 = Array2::zeros((self.guess_range as usize, self.len_samples));
+        let mut tmp_avg_0 = Array2::zeros((self.guess_range, self.len_samples));
+        let mut tmp_avg_1 = Array2::zeros((self.guess_range, self.len_samples));
 
-        for row in 0..self.guess_range as usize {
+        for row in 0..self.guess_range {
             let tmp_row_0 = self.sum_0.row(row).to_owned() / self.count_0[row] as f32;
             let tmp_row_1 = self.sum_1.row(row).to_owned() / self.count_1[row] as f32;
             tmp_avg_0.row_mut(row).assign(&tmp_row_0);
@@ -145,15 +145,17 @@ impl<T: Clone> Dpa<T> {
         self.corr.view()
     }
 
-    pub fn pass_guess(&self) -> i32 {
-        let mut init_value: f32 = 0.0;
-        let mut guess: i32 = 0;
+    pub fn pass_guess(&self) -> usize {
+        let mut init_value = 0.0;
+        let mut guess = 0;
+
         for i in 0..self.guess_range {
-            if self.max_corr[i as usize] > init_value {
-                init_value = self.max_corr[i as usize];
+            if self.max_corr[i] > init_value {
+                init_value = self.max_corr[i];
                 guess = i;
             }
         }
+
         guess
     }
 }