iterators

src/iterator.rs

@@ -0,0 +1,139 @@
+use crate::*;
+use crate::c_abi::xdrfile::exdrENDOFFILE;
+use failure::Error;
+use std::rc::Rc;
+
+impl IntoIterator for XTCTrajectory {
+    type Item = Result<Rc<Frame>, Error>;
+    type IntoIter = XTCTrajectoryIterator;
+
+    fn into_iter(mut self) -> Self::IntoIter {
+        // TODO this should be handled without the requirement of unwrap
+        let num_atoms = self.get_num_atoms().unwrap();
+        XTCTrajectoryIterator {
+            trajectory: self,
+            item: Rc::new(Frame::with_capacity(num_atoms)),
+            has_error: false
+        }
+    }
+}
+
+/* 
+Iterator for trajectories. This iterator yields a Result<Frame, Error>
+for each frame in the trajectory file and stops with yielding None once the 
+trajectory is finished. Also yields None after the first occurence of an error
+*/
+pub struct XTCTrajectoryIterator {
+    trajectory: XTCTrajectory,
+    item: Rc<Frame>,
+    has_error: bool,
+}
+
+impl Iterator for XTCTrajectoryIterator {
+    type Item = Result<Rc<Frame>, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> { // Reuse old frame
+        if self.has_error {
+            return None;
+        }
+
+        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));
+                Rc::get_mut(&mut self.item).unwrap()
+            }
+
+        };
+        match self.trajectory.read(item) {
+            Ok(()) => Some(Ok(Rc::clone(&self.item))),
+            Err(msg) => {
+                if msg.to_string().contains(&exdrENDOFFILE.to_string()) {
+                    None
+                } else {
+                    self.has_error = true;
+                    Some(Err(msg))
+                }
+            }
+        } 
+    }
+}
+
+impl IntoIterator for TRRTrajectory {
+    type Item = Result<Rc<Frame>, Error>;
+    type IntoIter = TRRTrajectoryIterator;
+
+    fn into_iter(mut self) -> Self::IntoIter {
+        // TODO this should be handled without the requirement of unwrap
+        let num_atoms = self.get_num_atoms().unwrap();
+        TRRTrajectoryIterator {
+            trajectory: self,
+            item: Rc::new(Frame::with_capacity(num_atoms)),
+            has_error: false
+        }
+    }
+}
+
+/* 
+Iterator for trajectories. This iterator yields a Result<Frame, Error>
+for each frame in the trajectory file and stops with yielding None once the 
+trajectory is finished. Also yields None after the first occurence of an error
+*/
+pub struct TRRTrajectoryIterator {
+    trajectory: TRRTrajectory,
+    item: Rc<Frame>,
+    has_error: bool,
+}
+
+impl Iterator for TRRTrajectoryIterator {
+    type Item = Result<Rc<Frame>, Error>;
+
+    fn next(&mut self) -> Option<Self::Item> { // Reuse old frame
+        if self.has_error {
+            return None;
+        }
+
+        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));
+                Rc::get_mut(&mut self.item).unwrap()
+            }
+
+        };
+        match self.trajectory.read(item) {
+            Ok(()) => Some(Ok(Rc::clone(&self.item))),
+            Err(msg) => {
+                if msg.to_string().contains(&exdrENDOFFILE.to_string()) {
+                    None
+                } else {
+                    self.has_error = true;
+                    Some(Err(msg))
+                }
+            }
+        } 
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    pub fn test_xtc_trajectory_iterator() {
+        let traj = XTCTrajectory::open(Path::new("tests/1l2y.xtc"), FileMode::Read).unwrap();
+        let frames: Vec<Rc<Frame>> = traj.into_iter().filter_map(Result::ok).collect();
+        assert!(frames.len() == 38);
+        assert!(frames[0].step == 1, frames[0].step);
+        assert!(frames[37].step == 38);
+    }
+
+    #[test]
+    pub fn test_trr_trajectory_iterator() {
+        let traj = TRRTrajectory::open(Path::new("tests/1l2y.trr"), FileMode::Read).unwrap();
+        let frames: Vec<Rc<Frame>> = traj.into_iter().filter_map(Result::ok).collect();
+        assert!(frames.len() == 38);
+        assert!(frames[0].step == 1, frames[0].step);
+        assert!(frames[37].step == 38);
+    }
+}

src/lib.rs

@@ -21,7 +21,7 @@
 //! // without instantiating data arrays for every step
 //! let mut frame = Frame::with_capacity(num_atoms);
 //! 
-//! // read the first trame of the trajectory
+//! // read the first frame of the trajectory
 //! let result = trj.read(&mut frame);
 //! match result {
 //!    Ok(_) => {
@@ -36,6 +36,27 @@
 //!    }
 //! }
 //! ```
+//!
+//! # Frame iteration
+//! For convenience, the trajectory implementations provide "into_iter" to
+//! be turned into an iterator that yields Rc<Frame>. If a frame is not kept
+//! during iteration, the Iterator reuses it for better performance (and hence,
+//! Rc is required)
+//!
+//! ```rust
+//! use xdrfile::*;
+//! use std::path::Path; 
+//!
+//! let mut path = Path::new("tests/1l2y.xtc");
+//! // get a handle to the file
+//! let trj = XTCTrajectory::open(path, FileMode::Read).unwrap();
+//! 
+//! // iterate over all frames
+//! for (idx, frame) in trj.into_iter().filter_map(Result::ok).enumerate() {
+//!     println!("{}", frame.time);
+//!     assert_eq!(idx+1, frame.step as usize);
+//! }
+//! ```
 
 
 #[cfg(test)]
@@ -43,8 +64,10 @@
 extern crate lazy_init;
 
 mod frame;
+mod iterator;
 pub mod c_abi; 
 pub use frame::Frame;
+pub use iterator::*;
 
 use c_abi::xdrfile;
 use c_abi::xdrfile::XDRFILE;

tests/integration.rs

@@ -0,0 +1,32 @@
+#[cfg(test)]
+mod integration {
+
+    use std::rc::Rc;
+    use xdrfile::*;
+    use std::path::Path;
+
+    #[test]
+    fn test_use_library() {
+        let path = Path::new("tests/1l2y.xtc");
+
+        let mut trj = XTCTrajectory::open(path, FileMode::Read).unwrap();
+        let num_atoms = trj.get_num_atoms().unwrap();
+        let mut frame = Frame::with_capacity(num_atoms);
+
+        trj.read(&mut frame).unwrap();
+        trj.read(&mut frame).unwrap();
+        assert_eq!(frame.step, 2);
+    }
+
+    #[test]
+    fn test_use_library_iterator() {
+        let path = Path::new("tests/1l2y.xtc");
+
+        let trj = XTCTrajectory::open(path, FileMode::Read).unwrap();
+        let frames: Vec<Rc<Frame>> = trj.into_iter().filter_map(Result::ok).collect();
+        for (idx, frame) in frames.iter().enumerate() {
+            assert_eq!(frame.step as usize, idx+1);
+        }
+    }
+
+}