cleaner implementation: next_custom

src/lfsr.rs

@@ -6,14 +6,6 @@ pub struct Lfsr {
     taps: Vec<usize>,
 }
 
-pub type FeedbackFn = fn(&[u8]) -> u8;
-
-#[derive(Clone, Copy, Default)]
-pub struct StepConfig {
-    /// Computes feedback from the selected bytes.
-    pub feedback: Option<FeedbackFn>,
-}
-
 impl Lfsr {
     /// Create a new LFSR
     ///
@@ -53,24 +45,19 @@ impl Lfsr {
     }
 
     /// Generic step:
-    pub fn next_custom(&mut self, feedback_indices: &[usize], cfg: StepConfig) -> u8 {
+    pub fn next_custom<const N: usize, F>(&mut self, indices: [usize; N], f: F) -> u8
+    where
+        F: Fn([u8; N]) -> u8,
+    {
         let output = *self.state.last().unwrap();
 
-        for &idx in feedback_indices {
+        for &idx in &indices {
             assert!(idx < self.state.len(), "Feedback index out of bounds");
         }
 
-        let selected: Vec<u8> = feedback_indices
+        let selected: [u8; N] = std::array::from_fn(|i| self.state[indices[i]]);
-            .iter()
-            .map(|&idx| self.state[idx])
-            .collect();
 
-        let feedback = if let Some(f) = cfg.feedback {
+        let feedback = f(selected);
-            f(&selected)
-        } else {
-            // Default: XOR
-            selected.iter().copied().fold(0u8, |acc, b| acc ^ b)
-        };
 
         // Shift right
         for i in (1..self.state.len()).rev() {
@@ -135,12 +122,7 @@ mod tests {
         // feedback = k0 + S(k2 + k7)
         let mut lfsr = Lfsr::new(10, vec![], vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
 
-        let output = lfsr.next_custom(
+        let output = lfsr.next_custom([0, 2, 7], |[k0, k2, k7]| k0 ^ s(k2 ^ k7));
-            &[0, 2, 7],
-            StepConfig {
-                feedback: Some(|b| b[0] ^ s(b[1] ^ b[2])),
-            },
-        );
 
         // Expected:
         // k0 = 1