From 53dd997129f762048e45dd5fba3a46a7c2829a62 Mon Sep 17 00:00:00 2001
From: Josh Mitchell <yoshanuikabundi@gmail.com>
Date: Wed, 11 Nov 2020 17:11:55 +1100
Subject: [PATCH] Broke out next_inner

---
 src/errors.rs   | 28 ++++++++++++++++++++--------
 src/iterator.rs | 39 ++++++++++++++++++++++-----------------
 src/lib.rs      | 12 ++++++------
 3 files changed, 48 insertions(+), 31 deletions(-)

diff --git a/src/errors.rs b/src/errors.rs
index b875cce..1e148cd 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -1,8 +1,8 @@
 use crate::c_abi;
 use crate::FileMode;
 use crate::Frame;
-use std::path::{Path, PathBuf};
 use std::error::Error as StdError;
+use std::path::{Path, PathBuf};
 
 /// Error type for the xdrfile library
 #[derive(Debug, Clone, PartialEq)]
@@ -26,7 +26,8 @@ pub enum Error {
     InvalidOsStr,
     /// A path could not be converted to &CStr because it had a null byte
     NullInStr(std::ffi::NulError),
-    CouldNotCheckNAtoms(Box<Error>),
+    CheckNAtomsDuringRead(Box<Error>),
+    CheckNAtomsDuringIter(Box<Error>),
 }
 
 impl Error {
@@ -60,9 +61,10 @@ impl Error {
 
 impl std::error::Error for Error {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        use Error::*;
         match &self {
-            Error::NullInStr(err) => Some(err),
-            Error::CouldNotCheckNAtoms(err) => Some(err.as_ref()),
+            NullInStr(err) => Some(err),
+            CheckNAtomsDuringRead(err) | CheckNAtomsDuringIter(err) => Some(err.as_ref()),
             _ => None,
         }
     }
@@ -121,10 +123,14 @@ impl std::fmt::Display for Error {
             }
             InvalidOsStr => write!(f, "Paths must be valid unicode on this platform"),
             NullInStr(_err) => write!(f, "Paths cannot include null bytes"),
-            CouldNotCheckNAtoms(_err) => write!(
+            CheckNAtomsDuringRead(_err) => write!(
                 f,
                 "Failed to check number of atoms in trajectory while reading a frame"
             ),
+            CheckNAtomsDuringIter(_err) => write!(
+                f,
+                "Failed to check number of atoms in trajectory while creating iterator"
+            ),
         }
     }
 }
@@ -234,8 +240,8 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::ffi::NulError;
     use std::ffi::CString;
+    use std::ffi::NulError;
 
     #[test]
     fn test_is_eof() {
@@ -293,12 +299,18 @@ mod tests {
 
         let path = Path::new(".");
         let mode = FileMode::Read;
-        let expected = Error::CouldNotOpen{path: path.to_path_buf(), mode: mode.to_owned()};
+        let expected = Error::CouldNotOpen {
+            path: path.to_path_buf(),
+            mode: mode.to_owned(),
+        };
         let err = Error::from((path, mode));
         assert_eq!(expected, err);
 
         let frame = Frame::with_capacity(0);
-        let expected = Error::WrongSizeFrame{expected: 10, found: 0};
+        let expected = Error::WrongSizeFrame {
+            expected: 10,
+            found: 0,
+        };
         let err = Error::from((&frame, 10));
         assert_eq!(expected, err);
     }
diff --git a/src/iterator.rs b/src/iterator.rs
index 1bea8cc..71b3657 100644
--- a/src/iterator.rs
+++ b/src/iterator.rs
@@ -45,24 +45,13 @@ pub struct TrajectoryIterator<T> {
     num_atoms: Result<u32>,
 }
 
-impl<T> Iterator for TrajectoryIterator<T>
-where
-    T: Trajectory,
-{
-    type Item = Result<Rc<Frame>>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.has_error {
-            return None;
-        }
-
+impl<T: Trajectory> TrajectoryIterator<T> {
+    /// Inner function for `next()`  to seperate error handling from iteration logic
+    fn next_inner(&mut self) -> <Self as Iterator>::Item {
         // If we couldn't read the number of frames when we called into_iter, return that error now
         let num_atoms = match &self.num_atoms {
             &Ok(n) => n,
-            Err(e) => {
-                self.has_error = true;
-                return Some(Err(e.clone()));
-            }
+            Err(e) => Err(Error::CheckNAtomsDuringIter(Box::new(e.clone())))?,
         };
 
         // Reuse old frame
@@ -75,8 +64,24 @@ where
             }
         };
 
-        match self.trajectory.read(item) {
-            Ok(()) => Some(Ok(Rc::clone(&self.item))),
+        self.trajectory.read(item)?;
+        Ok(Rc::clone(&self.item))
+    }
+}
+
+impl<T> Iterator for TrajectoryIterator<T>
+where
+    T: Trajectory,
+{
+    type Item = Result<Rc<Frame>>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.has_error {
+            return None;
+        }
+
+        match self.next_inner() {
+            Ok(item) => Some(Ok(item)),
             Err(e) if e.is_eof() => None,
             Err(e) => {
                 self.has_error = true;
diff --git a/src/lib.rs b/src/lib.rs
index cbd1fe4..83fc2dc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -216,9 +216,9 @@ impl Trajectory for XTCTrajectory {
     fn read(&mut self, frame: &mut Frame) -> Result<()> {
         let mut step: i32 = 0;
 
-        let num_atoms = self
-            .get_num_atoms()
-            .map_err(|e| Error::CouldNotCheckNAtoms(Box::new(e)))? as usize;
+        let num_atoms =
+            self.get_num_atoms()
+                .map_err(|e| Error::CheckNAtomsDuringRead(Box::new(e)))? as usize;
         if num_atoms != frame.coords.len() {
             Err((&*frame, num_atoms))?;
         }
@@ -335,9 +335,9 @@ impl Trajectory for TRRTrajectory {
         let mut step: i32 = 0;
         let mut lambda: f32 = 0.0;
 
-        let num_atoms = self
-            .get_num_atoms()
-            .map_err(|e| Error::CouldNotCheckNAtoms(Box::new(e)))? as usize;
+        let num_atoms =
+            self.get_num_atoms()
+                .map_err(|e| Error::CheckNAtomsDuringRead(Box::new(e)))? as usize;
         if num_atoms != frame.coords.len() {
             Err((&*frame, num_atoms))?;
         }