Move to streaming style encoders and decoders

src/baudot.rs

@@ -50,63 +50,140 @@ enum ShiftState {
     Figures,
 }
 
-pub fn encode(msg: &str) -> Result<Vec<u8>, u8> {
-    let mut res = Vec::<u8>::new();
-
-    let mut shift = ShiftState::None;
-
-    for &byte in msg.to_ascii_uppercase().as_bytes().iter() {
-        let letter = code_from_letter(&byte);
-        let figure = code_from_figure(&byte);
-        match (letter, figure, shift) {
-            (Some(l), Some(_), ShiftState::None) => {
-                res.push(SHIFT_LETTERS);
-                shift = ShiftState::Letters;
-                res.push(l);
-            }
-            (Some(l), Some(_), _) => res.push(l),
-            (Some(l), None, ShiftState::Letters) => res.push(l),
-            (Some(l), None, _) => {
-                res.push(SHIFT_LETTERS);
-                shift = ShiftState::Letters;
-                res.push(l);
-            }
-            (None, Some(f), ShiftState::Figures) => res.push(f),
-            (None, Some(f), _) => {
-                res.push(SHIFT_FIGURES);
-                shift = ShiftState::Figures;
-                res.push(f);
-            }
-            (None, None, _) => return Err(byte),
-        }
-    }
-
-    Ok(res)
+struct Encoder<'a> {
+    shift_state: ShiftState,
+    buffer: Option<u8>,
+    iterator: Box<dyn Iterator<Item = u8> + 'a>,
 }
 
-pub fn decode(msg: &Vec<u8>) -> Result<String, u8> {
-    let mut res = String::new();
+impl<'a> Encoder<'a> {
+    pub fn from<I>(iterator: I) -> Encoder<'a>
+    where
+        I: Iterator<Item = u8> + 'a,
+    {
+        Encoder {
+            shift_state: ShiftState::None,
+            buffer: None,
+            iterator: Box::new(iterator),
+        }
+    }
+}
 
-    // Lets assume this is a sane starting point.
-    let mut shift = ShiftState::Letters;
+impl<'a> Iterator for Encoder<'a> {
+    type Item = Result<u8, u8>;
 
-    for &code in msg.iter() {
-        match (code, shift) {
-            (SHIFT_LETTERS, _) => shift = ShiftState::Letters,
-            (SHIFT_FIGURES, _) => shift = ShiftState::Figures,
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some(c) = self.buffer {
+            self.buffer = None;
+            return Some(Ok(c));
+        }
+
+        let byte = match self.iterator.next() {
+            Some(b) => b,
+            None => return None,
+        };
+
+        let letter = code_from_letter(&byte);
+        let figure = code_from_figure(&byte);
+
+        match (letter, figure, self.shift_state) {
+            (Some(l), Some(_), ShiftState::None) => {
+                self.shift_state = ShiftState::Letters;
+                self.buffer = Some(l);
+
+                Some(Ok(SHIFT_LETTERS))
+            }
+            (Some(l), Some(_), _) => Some(Ok(l)),
+            (Some(l), None, ShiftState::Letters) => Some(Ok(l)),
+            (Some(l), None, _) => {
+                self.shift_state = ShiftState::Letters;
+                self.buffer = Some(l);
+
+                Some(Ok(SHIFT_LETTERS))
+            }
+            (None, Some(f), ShiftState::Figures) => Some(Ok(f)),
+            (None, Some(f), _) => {
+                self.shift_state = ShiftState::Figures;
+                self.buffer = Some(f);
+
+                Some(Ok(SHIFT_FIGURES))
+            }
+            (None, None, _) => Some(Err(byte)),
+        }
+    }
+}
+
+pub fn encode(msg: &str) -> Result<Vec<u8>, u8> {
+    let encoder = Encoder::from(msg.to_ascii_uppercase().into_bytes().into_iter());
+    encoder.collect()
+}
+
+struct Decoder<'a> {
+    shift_state: ShiftState,
+    iterator: Box<dyn Iterator<Item = u8> + 'a>,
+}
+
+impl<'a> Decoder<'a> {
+    pub fn from<I>(iterator: I) -> Decoder<'a>
+    where
+        I: Iterator<Item = u8> + 'a,
+    {
+        Decoder {
+            shift_state: ShiftState::Letters,
+            iterator: Box::new(iterator),
+        }
+    }
+}
+
+impl<'a> Iterator for Decoder<'a> {
+    // Option<Result<u8, u8>> for nice error handling if somebody wants to use collect
+    // See https://doc.rust-lang.org/stable/rust-by-example/error/iter_result.html
+    type Item = Result<u8, u8>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut byte = match self.iterator.next() {
+            Some(b) => b,
+            None => return None,
+        };
+
+        match byte {
+            SHIFT_LETTERS => {
+                self.shift_state = ShiftState::Letters;
+                byte = match self.iterator.next() {
+                    Some(b) => b,
+                    None => return None,
+                };
+            }
+            SHIFT_FIGURES => {
+                self.shift_state = ShiftState::Figures;
+                byte = match self.iterator.next() {
+                    Some(b) => b,
+                    None => return None,
+                };
+            }
+            _ => {}
+        };
+
+        match (byte, self.shift_state) {
             (c, ShiftState::Letters) => match letter_from_code(c) {
-                Some(l) => res.push(l as char),
-                None => return Err(c),
+                Some(l) => Some(Ok(l)),
+                None => Some(Err(c)),
             },
             (c, ShiftState::Figures) => match figure_from_code(c) {
-                Some(l) => res.push(l as char),
-                None => return Err(c),
+                Some(f) => Some(Ok(f)),
+                None => Some(Err(c)),
             },
             (_, ShiftState::None) => panic!("Shift state is expectedly None"),
         }
     }
+}
 
-    Ok(res)
+pub fn decode(msg: &Vec<u8>) -> Result<String, u8> {
+    let decoder = Decoder::from(msg.iter().copied());
+    match decoder.collect() {
+        Ok(v) => Ok(String::from_utf8(v).unwrap()),
+        Err(c) => Err(c),
+    }
 }
 
 #[cfg(test)]