From b02e7207ac33f0b765929ccd7b4bd9c508f240d9 Mon Sep 17 00:00:00 2001
From: Matej Focko <me@mfocko.xyz>
Date: Fri, 30 Dec 2022 12:11:41 +0100
Subject: [PATCH] vector2d: refactor indexing

---
 src/vector2d.rs | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/vector2d.rs b/src/vector2d.rs
index 28e7435..a90c1d9 100644
--- a/src/vector2d.rs
+++ b/src/vector2d.rs
@@ -1,7 +1,7 @@
 use std::cmp::Eq;
 use std::fmt::Debug;
 use std::hash::Hash;
-use std::ops::{Add, Mul, Sub};
+use std::ops::{Add, Index, IndexMut, Mul, Sub};
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct Vector2D<T> {
@@ -23,16 +23,28 @@ impl<T> Vector2D<T> {
     }
 }
 
-pub fn index<'a, T, U>(v: &'a [Vec<U>], idx: &Vector2D<T>) -> &'a U
+pub fn index<'a, C, I>(v: &'a [C], idx: &Vector2D<I>) -> &'a C::Output
 where
-    usize: TryFrom<T>,
-    <usize as TryFrom<T>>::Error: Debug,
-    T: Copy,
+    I: Copy,
+    C: Index<usize>,
+    usize: TryFrom<I>,
+    <usize as TryFrom<I>>::Error: Debug,
 {
     let (x, y): (usize, usize) = (idx.x.try_into().unwrap(), idx.y.try_into().unwrap());
     &v[y][x]
 }
 
+pub fn index_mut<'a, C, I>(v: &'a mut [C], idx: &Vector2D<I>) -> &'a mut C::Output
+where
+    I: Copy,
+    C: IndexMut<usize>,
+    usize: TryFrom<I>,
+    <usize as TryFrom<I>>::Error: Debug,
+{
+    let (x, y): (usize, usize) = (idx.x.try_into().unwrap(), idx.y.try_into().unwrap());
+    &mut v[y][x]
+}
+
 pub fn in_range<T, U>(v: &[Vec<U>], idx: &Vector2D<T>) -> bool
 where
     usize: TryInto<T>,