From ee9158555e920f392cba33c7c0e10ba3786a558c Mon Sep 17 00:00:00 2001
From: daniel <bauer@cbs.tu-darmstadt.de>
Date: Tue, 17 Nov 2020 09:00:11 +0100
Subject: [PATCH] array-like frame indexing

---
 README.md    |  2 +-
 src/frame.rs | 69 +++++++++++++++++++++++++++++++++++++++++++++++-----
 src/lib.rs   |  3 ++-
 3 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 6b0f29b..2addf56 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ fn main() -> Result<()> {
     assert_eq!(frame.step, 1);
     assert_eq!(frame.len(), num_atoms);
 
-    let first_atom_coords = frame.coords[0];
+    let first_atom_coords = frame[0];  // shorthand for frame.coords[0]
     assert_eq!(first_atom_coords, [-0.8901, 0.4127, -0.055499997]);
 
     Ok(())
diff --git a/src/frame.rs b/src/frame.rs
index 08e81c5..464b30f 100644
--- a/src/frame.rs
+++ b/src/frame.rs
@@ -1,3 +1,5 @@
+use std::ops::{Index, IndexMut};
+
 /// A frame represents a single step in a trajectory.
 #[derive(Clone, Debug)]
 pub struct Frame {
@@ -66,6 +68,20 @@ impl Frame {
     }
 }
 
+impl Index<usize> for Frame {
+    type Output = [f32; 3];
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.coords[index]
+    }
+}
+
+impl IndexMut<usize> for Frame {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output { {
+        &mut self.coords[index]
+    }}
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -80,15 +96,15 @@ mod tests {
     #[test]
     fn test_frame_filter_atoms() {
         let mut frame = Frame::with_len(3);
-        frame.coords[0] = [1.0, 2.0, 3.0];
-        frame.coords[1] = [4.0, 5.0, 6.0];
-        frame.coords[2] = [7.0, 8.0, 9.0];
+        frame[0] = [1.0, 2.0, 3.0];
+        frame[1] = [4.0, 5.0, 6.0];
+        frame[2] = [7.0, 8.0, 9.0];
         let filter: Vec<usize> = vec![1, 2];
         let mut frame_new = frame.clone();
         frame_new.filter_coords(&filter);
         assert!(frame_new.len() == filter.len());
-        assert!(frame_new.coords[0] == frame.coords[1]);
-        assert!(frame_new.coords[1] == frame.coords[2]);
+        assert!(frame_new.coords[0] == frame[1]);
+        assert!(frame_new.coords[1] == frame[2]);
     }
 
     #[test]
@@ -108,7 +124,48 @@ mod tests {
 
         frame.filter_coords(&[1]);
         for i in 0..3 {
-            assert_approx_eq!(frame.coords[0][i], 1.0);
+            assert_approx_eq!(frame[0][i], 1.0);
         }
     }
+
+    #[test]
+    #[allow(unused_mut)]
+    fn test_index() {
+        // test Index
+        let frame = Frame {
+            step: 0,
+            time: 0.0,
+            box_vector: [[0.0; 3]; 3],
+            coords: vec![[0.0; 3], [1.0; 3], [2.0; 3]]
+        };
+        for i in 0..frame.len() {
+            for j in 0..3 {
+                assert_approx_eq!(frame[i][j], frame[i][j]);
+            }
+        }
+
+        // test IndexMut
+        let mut frame = Frame {
+            step: 0,
+            time: 0.0,
+            box_vector: [[0.0; 3]; 3],
+            coords: vec![[0.0; 3], [1.0; 3], [2.0; 3]]
+        };
+        for i in 0..frame.len() {
+            for j in 0..3 {
+                assert_approx_eq!(frame[i][j], frame.coords[i][j]);
+            }
+        }
+
+        frame[0] = [123.0; 3];
+        for i in 0..frame.len() {
+            for j in 0..3 {
+                assert_approx_eq!(frame[i][j], frame.coords[i][j]);
+                if i == 0 {
+                    assert_approx_eq!(frame[i][j], 123.0);    
+                }
+            }
+        }
+
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 6a916f5..9783561 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -26,7 +26,7 @@
 //!     assert_eq!(frame.step, 1);
 //!     assert_eq!(frame.len(), num_atoms);
 //!
-//!     let first_atom_coords = frame.coords[0];
+//!     let first_atom_coords = frame[0];
 //!     assert_eq!(first_atom_coords, [-0.8901, 0.4127, -0.055499997]);
 //!
 //!     Ok(())
@@ -232,6 +232,7 @@ pub trait Trajectory {
 
     /// Get the number of atoms from the give trajectory
     fn get_num_atoms(&mut self) -> Result<usize>;
+
 }
 
 /// Handle to Read/Write XTC Trajectories