From 4b4f4fbf80e62d6aa49a2a950658ff4a3326a588 Mon Sep 17 00:00:00 2001
From: kageru <kageru@encode.moe>
Date: Thu, 8 Jul 2021 19:58:30 +0200
Subject: [PATCH] Add tests and clean up add/sub/mul for positions

---
 2020/src/grid.rs          | 13 +++++++++++++
 2020/src/grid/position.rs | 27 ++++++++++++---------------
 2 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/2020/src/grid.rs b/2020/src/grid.rs
index 621eb8a..347029b 100644
--- a/2020/src/grid.rs
+++ b/2020/src/grid.rs
@@ -61,9 +61,22 @@ pub fn draw_ascii<T: Display + Default, S: BuildHasher>(coordinates: &HashMap<Po
 #[cfg(test)]
 mod tests {
     use super::*;
+
     #[test]
     fn test_add() {
         assert_eq!(PositionND::from([0, 2]) + PositionND::from([-1, 0]), [-1, 2].into());
         assert_eq!(PositionND::from([0, -1]) + PositionND::from(Direction::Up), [0, 0].into());
     }
+
+    #[test]
+    fn test_sub() {
+        assert_eq!(PositionND::from([0, 2]) - PositionND::from([-1, 0]), [1, 2].into());
+        assert_eq!(PositionND::from([0, -1]) - PositionND::from([0, -1]), [0, 0].into());
+    }
+
+    #[test]
+    fn test_mul() {
+        assert_eq!(PositionND::from([0, 2]) * 5, [0, 10].into());
+        assert_eq!(PositionND::from([0, -1]) * -2, [0, 2].into());
+    }
 }
diff --git a/2020/src/grid/position.rs b/2020/src/grid/position.rs
index ef5e17f..45b2c86 100644
--- a/2020/src/grid/position.rs
+++ b/2020/src/grid/position.rs
@@ -117,36 +117,33 @@ impl<const DIMS: usize> PositionND<DIMS> {
 impl<const D: usize> Mul<i64> for PositionND<D> {
     type Output = PositionND<D>;
 
-    fn mul(self, rhs: i64) -> Self::Output {
-        let mut points = [0; D];
-        for i in 0..D {
-            points[i] = self.points[i] * rhs;
+    fn mul(mut self, rhs: i64) -> Self::Output {
+        for p in self.points.iter_mut() {
+            *p *= rhs;
         }
-        PositionND { points }
+        self
     }
 }
 
 impl<const D: usize> Add<PositionND<D>> for PositionND<D> {
     type Output = PositionND<D>;
 
-    fn add(self, rhs: PositionND<D>) -> Self::Output {
-        let mut points = [0; D];
-        for i in 0..D {
-            points[i] = self.points[i] + rhs.points[i];
+    fn add(mut self, rhs: PositionND<D>) -> Self::Output {
+        for (x, y) in self.points.iter_mut().zip(rhs.points) {
+            *x += y;
         }
-        PositionND { points }
+        self
     }
 }
 
 impl<const D: usize> Sub<PositionND<D>> for PositionND<D> {
     type Output = PositionND<D>;
 
-    fn sub(self, rhs: PositionND<D>) -> Self::Output {
-        let mut points = [0; D];
-        for i in 0..D {
-            points[i] = self.points[i] - rhs.points[i];
+    fn sub(mut self, rhs: PositionND<D>) -> Self::Output {
+        for (x, y) in self.points.iter_mut().zip(rhs.points) {
+            *x -= y;
         }
-        PositionND { points }
+        self
     }
 }