mini_enigma/components/
wheel.rs

1use core::mem::MaybeUninit;
2
3use super::utils::{Letter, DEFAULT_ROTOR_NUM};
4use super::Rotor;
5
6#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
7/// Represents the configuration (i.e. ordering) of the `Rotor` set in an Enigma machine
8pub struct Wheel<const ROTOR_NUM: usize = DEFAULT_ROTOR_NUM> {
9    walzenlage: [Rotor; ROTOR_NUM],
10    double_step: bool,
11}
12
13impl<const ROTOR_NUM: usize> Default for Wheel<ROTOR_NUM> {
14    fn default() -> Self {
15        Wheel {
16            walzenlage: [Rotor::default(); ROTOR_NUM],
17            double_step: false,
18        }
19    }
20}
21
22impl<const ROTOR_NUM: usize> Wheel<ROTOR_NUM> {
23    /// Creates a new `Rotor` configuration in left-to-right ordering
24    #[must_use]
25    pub fn new(rotors: [Rotor; ROTOR_NUM]) -> Self {
26        Wheel {
27            walzenlage: rotors,
28            double_step: false,
29        }
30    }
31
32    /// Send a signal forward through the rotors
33    pub fn forward(&mut self, letter: Letter) -> Letter {
34        // If arbritrary wheels A, B, and C are arranged in the enigma left-to-right this means the output
35        // from the Eintrittswalze (or entry stator wheel) enters C first, then B, then A before hitting the
36        // reflector and flowing back through A, B and finally C, since walzenlage represents the left-to-right
37        // ordering we must traverse this in the reverse order.
38        self.walzenlage
39            .iter()
40            .rev()
41            .fold(letter, |a, e| e.wiring_map(a))
42    }
43
44    fn simple_step(&mut self, from: usize) {
45        self.walzenlage.iter_mut().take(from).rev().all(Rotor::step);
46    }
47
48    /// Step the rotors
49    pub fn step(&mut self) {
50        self.simple_step(ROTOR_NUM);
51        if self.double_step {
52            self.double_step = false;
53            self.simple_step(ROTOR_NUM - 1);
54        } else {
55            let mut idx = 1;
56            while idx < ROTOR_NUM - 1 && !self.walzenlage[idx].will_turn() {
57                idx += 1;
58            }
59            if idx < ROTOR_NUM - 1 {
60                self.double_step = true;
61            }
62        }
63    }
64
65    /// Send a signal backward through the rotors
66    #[must_use]
67    pub fn backward(&self, letter: Letter) -> Letter {
68        self.walzenlage
69            .iter()
70            .fold(letter, |a, e| e.rev_wiring_map(a))
71    }
72
73    /// Set the positions of each `Rotor` in left-to-right ordering
74    pub fn set_positions(&mut self, positions: [Letter; ROTOR_NUM]) {
75        positions
76            .iter()
77            .zip(self.walzenlage.iter_mut())
78            .for_each(|(pos, rotor)| rotor.set_grundstellung(*pos));
79    }
80
81    /// Set the wiring offset (aka ring setting) of each `Rotor` in left-to-right ordering
82    pub fn set_ring_wiring_offset(&mut self, positions: [u8; ROTOR_NUM]) {
83        positions
84            .iter()
85            .zip(self.walzenlage.iter_mut())
86            .for_each(|(pos, rotor)| rotor.set_ringstellung(*pos));
87    }
88
89    /// Get current `Rotor` positions
90    #[must_use]
91    pub fn get_positions(&self) -> [Letter; ROTOR_NUM] {
92        let mut ret: [MaybeUninit<Letter>; ROTOR_NUM] =
93            unsafe { MaybeUninit::uninit().assume_init() };
94        self.walzenlage
95            .iter()
96            .enumerate()
97            .for_each(|(i, x)| ret[i] = MaybeUninit::new(x.get_grundstellung()));
98        unsafe { core::mem::transmute_copy::<_, [Letter; ROTOR_NUM]>(&ret) }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use crate::components::{ROTOR_I, ROTOR_II, ROTOR_III};
105
106    use super::*;
107
108    #[test]
109    fn test_single_step() {
110        let mut wheel = Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]);
111        wheel.set_positions([Letter::A, Letter::A, Letter::U]);
112        assert_eq!([Letter::A, Letter::A, Letter::U], wheel.get_positions());
113        wheel.step();
114        assert_eq!([Letter::A, Letter::A, Letter::V], wheel.get_positions());
115        wheel.step();
116        assert_eq!([Letter::A, Letter::B, Letter::W], wheel.get_positions());
117        wheel.step();
118        assert_eq!([Letter::A, Letter::B, Letter::X], wheel.get_positions());
119    }
120
121    #[test]
122    fn test_double_step() {
123        let mut wheel = Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]);
124        wheel.set_positions([Letter::A, Letter::D, Letter::U]);
125        assert_eq!([Letter::A, Letter::D, Letter::U], wheel.get_positions());
126        wheel.step();
127        assert_eq!([Letter::A, Letter::D, Letter::V], wheel.get_positions());
128        wheel.step();
129        assert_eq!([Letter::A, Letter::E, Letter::W], wheel.get_positions());
130        wheel.step();
131        assert_eq!([Letter::B, Letter::F, Letter::X], wheel.get_positions());
132        wheel.step();
133        assert_eq!([Letter::B, Letter::F, Letter::Y], wheel.get_positions());
134    }
135}
mini_enigma/components/wheel.rs

mini_enigma/components/
wheel.rs
