diff --git a/src/errors.rs b/src/errors.rs index e874f48..482b036 100644 --- a/src/errors.rs +++ b/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"), } } } diff --git a/src/lib.rs b/src/lib.rs index ebaf74a..c4803cc 100644 --- a/src/lib.rs +++ b/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 { + 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 { + 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 { + 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> { + 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> { + 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()?;