From 3989aea5b3bc13fc4ea19ceecd6227b44946943f Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Thu, 2 Apr 2026 11:11:21 +0200
Subject: [PATCH] custom feedback

---
 Cargo.lock  |  2 +-
 Cargo.toml  |  2 +-
 src/lfsr.rs | 68 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b4c3f55..07b92cf 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3,5 +3,5 @@
 version = 3
 
 [[package]]
-name = "lfsr"
+name = "tea-3"
 version = "0.1.0"
diff --git a/Cargo.toml b/Cargo.toml
index 55a3683..d415757 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "lfsr"
+name = "tea-3"
 version = "0.1.0"
 edition = "2021"
 
diff --git a/src/lfsr.rs b/src/lfsr.rs
index a6b62a5..4c02fe0 100644
--- a/src/lfsr.rs
+++ b/src/lfsr.rs
@@ -6,6 +6,14 @@ 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
     ///
@@ -44,6 +52,36 @@ impl Lfsr {
         output
     }
 
+    /// Generic step:
+    pub fn next_custom(&mut self, feedback_indices: &[usize], cfg: StepConfig) -> u8 {
+        let output = *self.state.last().unwrap();
+
+        for &idx in feedback_indices {
+            assert!(idx < self.state.len(), "Feedback index out of bounds");
+        }
+
+        let selected: Vec<u8> = feedback_indices
+            .iter()
+            .map(|&idx| self.state[idx])
+            .collect();
+
+        let feedback = if let Some(f) = cfg.feedback {
+            f(&selected)
+        } else {
+            // Default: XOR
+            selected.iter().copied().fold(0u8, |acc, b| acc ^ b)
+        };
+
+        // Shift right
+        for i in (1..self.state.len()).rev() {
+            self.state[i] = self.state[i - 1];
+        }
+
+        self.state[0] = feedback;
+
+        output
+    }
+
     /// Get current state
     pub fn state(&self) -> &[u8] {
         &self.state
@@ -86,6 +124,36 @@ mod tests {
         assert_eq!(lfsr.state(), &[3, 1, 2]);
     }
 
+    #[test]
+    fn test_next_custom() {
+        // S-box placeholder: S(n) = n + 15
+        // + are XOR in this context
+        fn s(x: u8) -> u8 {
+            x ^ 15
+        }
+
+        // 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(
+            &[0, 2, 7],
+            StepConfig {
+                feedback: Some(|b| b[0] ^ s(b[1] ^ b[2])),
+            },
+        );
+
+        // Expected:
+        // k0 = 1
+        // k2 = 3
+        // k7 = 8
+        // k2 + k7 = 3 + 8 = 11
+        // s(11) = 11 + 15 = 4
+        // feedback = 1 + 4 = 5
+
+        assert_eq!(output, 10);
+        assert_eq!(lfsr.state(), &[5, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    }
+
     #[test]
     fn test_empty_taps() {
         let mut lfsr = Lfsr::new(3, vec![], vec![1, 2, 3]);