Merge pull request #5 from Yoshanuikabundi/xdr-seek

XDR seek

src/errors.rs

@@ -139,8 +139,10 @@ pub enum ErrorTask {
     Read,
     /// A frame was being written to a file
     Write,
-    /// An file was being flushed to disk
+    /// A file was being flushed to disk
     Flush,
+    /// A seek operation was being run on a file
+    Seek,
 }
 
 impl std::fmt::Display for ErrorTask {
@@ -151,6 +153,7 @@ impl std::fmt::Display for ErrorTask {
             Read => write!(f, "reading trajectory"),
             Write => write!(f, "writing trajectory"),
             Flush => write!(f, "flushing trajectory"),
+            Seek => write!(f, "seeking in trajectory"),
         }
     }
 }

src/lib.rs

@@ -78,7 +78,10 @@ use c_abi::xdrfile_xtc;
 use lazy_init::Lazy;
 use std::cell::Cell;
 use std::ffi::CString;
+use std::io;
 use std::path::{Path, PathBuf};
+use std::convert::TryInto;
+use std::io::SeekFrom;
 
 #[derive(Debug, Clone, PartialEq)]
 pub enum FileMode {
@@ -153,6 +156,32 @@ impl XDRFile {
             }
         }
     }
+
+    /// Get the current position in the file
+    pub fn tell(&self) -> u64 {
+        unsafe {
+            xdr_seek::xdr_tell(self.xdrfile)
+                .try_into()
+                .expect("i64 could not be converted to u64")
+        }
+    }
+}
+
+impl io::Seek for XDRFile {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        let (whence, pos) = match pos {
+            SeekFrom::Start(u) => (0, u as i64),
+            SeekFrom::Current(i) => (1, i),
+            SeekFrom::End(i) => (2, i),
+        };
+        unsafe {
+            let code = xdr_seek::xdr_seek(self.xdrfile, pos, whence) as u32;
+            match check_code(code, ErrorTask::Seek) {
+                None => Ok(self.tell()),
+                Some(err) => Err(io::Error::new(io::ErrorKind::Other, err)),
+            }
+        }
+    }
 }
 
 impl Drop for XDRFile {
@@ -299,6 +328,19 @@ impl Trajectory for XTCTrajectory {
     }
 }
 
+impl XTCTrajectory {
+    /// Get the current position in the file
+    pub fn tell(&self) -> u64 {
+        self.handle.tell()
+    }
+}
+
+impl io::Seek for XTCTrajectory {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        self.handle.seek(pos)
+    }
+}
+
 /// Read/Write TRR Trajectories
 pub struct TRRTrajectory {
     handle: XDRFile,
@@ -422,11 +464,26 @@ impl Trajectory for TRRTrajectory {
     }
 }
 
+impl TRRTrajectory {
+    /// Get the current position in the file
+    pub fn tell(&self) -> u64 {
+        self.handle.tell()
+    }
+}
+
+impl io::Seek for TRRTrajectory {
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        self.handle.seek(pos)
+    }
+}
+
 #[cfg(test)]
 mod tests {
 
     use super::*;
     use tempfile::NamedTempFile;
+    use std::io::Seek;
+    use std::io::Write;
 
     #[test]
     fn test_read_write_xtc() -> Result<()> {
@@ -573,6 +630,74 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_tell() -> std::result::Result<(), Box<dyn std::error::Error>> {
+        let tempfile = NamedTempFile::new()?;
+        let tmp_path = tempfile.path();
+
+        let natoms: u32 = 2;
+        let frame = Frame {
+            num_atoms: natoms,
+            step: 5,
+            time: 2.0,
+            box_vector: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+            coords: vec![[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]],
+        };
+        let mut f = TRRTrajectory::open_write(tmp_path)?;
+        assert_eq!(f.tell(), 0);
+        f.write(&frame)?;
+        assert_eq!(f.tell(), 144);
+        f.flush()?;
+
+        let mut new_frame = Frame::with_capacity(natoms);
+        let mut f = TRRTrajectory::open_read(tmp_path)?;
+        assert_eq!(f.tell(), 0);
+
+        f.read(&mut new_frame)?;
+        assert_eq!(f.tell(), 144);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_seek() -> std::result::Result<(), Box<dyn std::error::Error>> {
+        let tempfile = NamedTempFile::new()?;
+        let tmp_path = tempfile.path();
+
+        let natoms: u32 = 2;
+        let mut frame = Frame {
+            num_atoms: natoms,
+            step: 0,
+            time: 0.0,
+            box_vector: [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
+            coords: vec![[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]],
+        };
+        let mut f = TRRTrajectory::open_write(tmp_path)?;
+        f.write(&frame)?;
+        let after_first_frame = f.tell();
+        frame.step += 1;
+        frame.time += 10.0;
+        f.write(&frame)?;
+        let after_second_frame = f.tell();
+        f.flush()?;
+
+        let mut new_frame = Frame::with_capacity(natoms);
+        let mut f = TRRTrajectory::open_read(tmp_path)?;
+        let pos = f.seek(std::io::SeekFrom::Current(144))?;
+        assert_eq!(pos, after_first_frame);
+
+        f.read(&mut new_frame)?;
+        assert_eq!(f.tell(), after_second_frame);
+
+        assert_eq!(new_frame.num_atoms, frame.num_atoms);
+        assert_eq!(new_frame.step, frame.step);
+        assert_eq!(new_frame.time, frame.time);
+        assert_eq!(new_frame.box_vector, frame.box_vector);
+        assert_eq!(new_frame.coords, frame.coords);
+
+        Ok(())
+    }
+
     #[test]
     fn test_err_could_not_open() {
         let file_name = "non-existent.xtc";
@@ -647,7 +772,6 @@ mod tests {
         }
 
         let mut file = std::fs::File::create(tmp_path)?;
-        use std::io::Write as _;
         file.write_all(&[0; 999])?;
         file.flush()?;