diff --git a/src/main.rs b/src/main.rs
index 3751241..890d06f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -7,7 +7,7 @@ use std::{
 
 use aes::{
     Aes128,
-    cipher::{BlockEncrypt, KeyInit, generic_array::GenericArray},
+    cipher::{self, BlockEncrypt, KeyInit, generic_array::GenericArray},
 };
 use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
@@ -30,21 +30,21 @@ static RSBOX: [u8; 256] = [
     0x17, 0x2b, 0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55, 0x21, 0x0c, 0x7d,
 ];
 
-fn calculate_models(ciphertexts: &[[u8; 16]]) -> Vec<[[u8; 16]; 256]> {
-    ciphertexts
-        .iter()
-        .map(|c| {
-            let mut row = [[0; 16]; 256];
+fn calculate_models(ciphertexts: &[[u8; 16]]) -> Vec<Vec<[u8; 16]>> {
+    let mut r = vec![Vec::with_capacity(ciphertexts.len()); 256];
 
-            for i in 0..256 {
-                for j in 0..16 {
-                    row[i][j] = RSBOX[(c[j] ^ (i as u8)) as usize];
-                }
+    for (i, c) in ciphertexts.iter().enumerate() {
+        for j in 0..256 {
+            let mut row = [0; 16];
+
+            for k in 0..16 {
+                row[k] = RSBOX[(c[k] ^ (j as u8)) as usize];
             }
+            r[j].push(row);
+        }
+    }
 
-            row
-        })
-        .collect()
+    r
 }
 
 fn read_msgs(path: impl AsRef<Path>) -> Vec<[u8; 16]> {
@@ -78,38 +78,35 @@ fn read_msgs(path: impl AsRef<Path>) -> Vec<[u8; 16]> {
 
 const TRACES: usize = 87;
 
-fn read_traces(path: impl AsRef<Path>) -> Vec<[u8; TRACES]> {
+fn read_traces(path: impl AsRef<Path>, capacity: usize) -> Vec<Vec<u8>> {
     let file = std::fs::File::open(path).unwrap();
 
     let bufreader = std::io::BufReader::new(file);
 
-    bufreader
-        .lines()
-        .map(|l| {
-            let l = l.unwrap();
+    let mut r: Vec<Vec<u8>> = (0..TRACES).map(|_| vec![0; capacity]).collect();
 
-            let mut trace = [0; TRACES];
+    for (i, l) in bufreader.lines().enumerate() {
+        let l = l.unwrap();
 
-            for (i, t) in l
-                .trim()
-                .split(',')
-                .map(|t| t.parse::<u8>().unwrap())
-                .enumerate()
-            {
-                trace[i] = t;
-            }
+        for (j, t) in l
+            .trim()
+            .split(',')
+            .map(|t| t.parse::<u8>().unwrap())
+            .enumerate()
+        {
+            r[j][i] = t;
+        }
+    }
 
-            trace
-        })
-        .collect()
+    r
 }
 
 fn correlation(
     bit: usize,
     key_hypothesis: usize,
     trace_index: usize,
-    cyphtertext: &[[[u8; 16]; 256]],
-    traces: &[[u8; TRACES]],
+    cyphtertext: &[Vec<[u8; 16]>],
+    traces: &[Vec<u8>],
 ) -> f64 {
     let mut x = 0i64;
     let mut y = 0i64;
@@ -117,9 +114,9 @@ fn correlation(
     let mut xsqr = 0i64;
     let mut ysqr = 0i64;
 
-    for i in 0..traces.len() {
-        let xi = (cyphtertext[i][key_hypothesis][bit / 8] & (1 << (bit % 8))) as i64;
-        let yi = traces[i][trace_index] as i64;
+    for i in 0..traces[0].len() {
+        let xi = (cyphtertext[key_hypothesis][i][bit / 8] & (1 << (bit % 8))) as i64;
+        let yi = traces[trace_index][i] as i64;
 
         x += xi;
         y += yi;
@@ -128,7 +125,7 @@ fn correlation(
         ysqr += yi * yi;
     }
 
-    let n = traces.len() as i64;
+    let n = traces[0].len() as i64;
     let num = (n * xy - x * y) as f64;
     let denom = f64::sqrt((n * xsqr - x * x) as f64) * f64::sqrt((n * ysqr - y * y) as f64);
 
@@ -145,7 +142,10 @@ fn main() {
     println!("calculate models: {:?}", start.elapsed());
 
     let start = std::time::Instant::now();
-    let traces = read_traces("./alan/Task-3-example_traces/test_traces.csv");
+    let traces = read_traces(
+        "./alan/Task-3-example_traces/test_traces.csv",
+        cyphertext.len(),
+    );
     println!("read traces: {:?}", start.elapsed());
 
     let start = std::time::Instant::now();
@@ -153,7 +153,7 @@ fn main() {
         let (max_index, max) = (0..256)
             .into_par_iter()
             .map(|key_hypothesis| {
-                let m = (0..traces[0].len())
+                let m = (0..TRACES)
                     .map(|trace_index| {
                         correlation(bit, key_hypothesis, trace_index, &models, &traces).abs()
                     })