From 6d911fae0b3eacc6c7b6ba2777de19c9eabe8ead Mon Sep 17 00:00:00 2001
From: Josh Mitchell <yoshanuikabundi@gmail.com>
Date: Wed, 11 Nov 2020 16:40:37 +1100
Subject: [PATCH] Clearer errors if get_num_atoms fails during into_iter

---
 src/iterator.rs | 53 +++++++++++++++++++++++++++++--------------------
 src/lib.rs      |  9 ++++-----
 2 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/src/iterator.rs b/src/iterator.rs
index 5f69fab..1bea8cc 100644
--- a/src/iterator.rs
+++ b/src/iterator.rs
@@ -1,20 +1,26 @@
 use crate::*;
 use std::rc::Rc;
 
+fn into_iter_inner<T: Trajectory>(mut traj: T) -> TrajectoryIterator<T> {
+    let num_atoms = traj.get_num_atoms();
+    let frame = match &num_atoms {
+        Ok(num_atoms) => Frame::with_capacity(*num_atoms),
+        Err(_) => Frame::new(),
+    };
+    TrajectoryIterator {
+        trajectory: traj,
+        item: Rc::new(frame),
+        has_error: false,
+        num_atoms,
+    }
+}
+
 impl IntoIterator for XTCTrajectory {
     type Item = Result<Rc<Frame>>;
     type IntoIter = TrajectoryIterator<XTCTrajectory>;
 
-    fn into_iter(mut self) -> Self::IntoIter {
-        let frame = match self.get_num_atoms() {
-            Ok(num_atoms) => Frame::with_capacity(num_atoms),
-            Err(_) => Frame::new(),
-        };
-        TrajectoryIterator {
-            trajectory: self,
-            item: Rc::new(frame),
-            has_error: false,
-        }
+    fn into_iter(self) -> Self::IntoIter {
+        into_iter_inner(self)
     }
 }
 
@@ -22,16 +28,8 @@ impl IntoIterator for TRRTrajectory {
     type Item = Result<Rc<Frame>>;
     type IntoIter = TrajectoryIterator<TRRTrajectory>;
 
-    fn into_iter(mut self) -> Self::IntoIter {
-        let frame = match self.get_num_atoms() {
-            Ok(num_atoms) => Frame::with_capacity(num_atoms),
-            Err(_) => Frame::new(),
-        };
-        TrajectoryIterator {
-            trajectory: self,
-            item: Rc::new(frame),
-            has_error: false,
-        }
+    fn into_iter(self) -> Self::IntoIter {
+        into_iter_inner(self)
     }
 }
 
@@ -44,6 +42,7 @@ pub struct TrajectoryIterator<T> {
     trajectory: T,
     item: Rc<Frame>,
     has_error: bool,
+    num_atoms: Result<u32>,
 }
 
 impl<T> Iterator for TrajectoryIterator<T>
@@ -53,19 +52,29 @@ where
     type Item = Result<Rc<Frame>>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        // Reuse old frame
         if self.has_error {
             return None;
         }
 
+        // 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()));
+            }
+        };
+
+        // Reuse old frame
         let item: &mut Frame = match Rc::get_mut(&mut self.item) {
             Some(item) => item,
             None => {
                 // caller kept frame. Create new one
-                self.item = Rc::new(Frame::with_capacity(self.item.num_atoms));
+                self.item = Rc::new(Frame::with_capacity(num_atoms));
                 Rc::get_mut(&mut self.item).expect("Could not get mutable access to new Rc")
             }
         };
+
         match self.trajectory.read(item) {
             Ok(()) => Some(Ok(Rc::clone(&self.item))),
             Err(e) if e.is_eof() => None,
diff --git a/src/lib.rs b/src/lib.rs
index 46e5fad..cbd1fe4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -102,7 +102,7 @@ impl FileMode {
 
 fn path_to_cstring(path: impl AsRef<Path>) -> Result<CString> {
     let s = path.as_ref().to_str().ok_or(Error::InvalidOsStr)?;
-    CString::new(s).map_err(Error::from)
+    Ok(CString::new(s)?)
 }
 
 /// Convert an error code from a C call to an Error
@@ -119,7 +119,6 @@ pub fn check_code(code: impl Into<ErrorCode>, task: ErrorTask) -> Option<Error>
     }
 }
 
-
 /// A safe wrapper around the c implementation of an XDRFile
 struct XDRFile {
     xdrfile: *mut XDRFILE,
@@ -150,7 +149,7 @@ impl XDRFile {
                 })
             } else {
                 // Something went wrong. But the C api does not tell us what
-                Err(Error::from((path, filemode)))
+                Err((path, filemode))?
             }
         }
     }
@@ -221,7 +220,7 @@ impl Trajectory for XTCTrajectory {
             .get_num_atoms()
             .map_err(|e| Error::CouldNotCheckNAtoms(Box::new(e)))? as usize;
         if num_atoms != frame.coords.len() {
-            return Err(Error::from((&*frame, num_atoms)));
+            Err((&*frame, num_atoms))?;
         }
 
         unsafe {
@@ -340,7 +339,7 @@ impl Trajectory for TRRTrajectory {
             .get_num_atoms()
             .map_err(|e| Error::CouldNotCheckNAtoms(Box::new(e)))? as usize;
         if num_atoms != frame.coords.len() {
-            return Err(Error::from((&*frame, num_atoms)));
+            Err((&*frame, num_atoms))?;
         }
 
         unsafe {