Impl NdIndex<IxDyn> for &T and Vec<Ix>

rust-ndarray · Jan 23, 2022 · 712bfd6 · 712bfd6
1 parent a7d1fd6
commit 712bfd6
Showing 1 changed file with 20 additions and 8 deletions.
diff --git a/src/dimension/ndindex.rs b/src/dimension/ndindex.rs
@@ -219,22 +219,34 @@ impl<'a> IntoDimension for &'a [Ix] {
     }
 }
 
-unsafe impl<'a> NdIndex<IxDyn> for &'a IxDyn {
+unsafe impl<'a> NdIndex<IxDyn> for &'a [Ix] {
     fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
-        (**self).index_checked(dim, strides)
+        stride_offset_checked(dim.ix(), strides.ix(), *self)
     }
     fn index_unchecked(&self, strides: &IxDyn) -> isize {
-        (**self).index_unchecked(strides)
+        zip(strides.ix(), *self)
+            .map(|(&s, &i)| stride_offset(i, s))
+            .sum()
     }
 }
 
-unsafe impl<'a> NdIndex<IxDyn> for &'a [Ix] {
+unsafe impl<'a, T> NdIndex<IxDyn> for &T
+where
+    T: NdIndex<IxDyn>,
+{
     fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
-        stride_offset_checked(dim.ix(), strides.ix(), *self)
+        (*self).index_checked(dim, strides)
     }
     fn index_unchecked(&self, strides: &IxDyn) -> isize {
-        zip(strides.ix(), *self)
-            .map(|(&s, &i)| stride_offset(i, s))
-            .sum()
+        (*self).index_unchecked(strides)
+    }
+}
+
+unsafe impl<'a> NdIndex<IxDyn> for Vec<Ix> {
+    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
+        (&**self).index_checked(dim, strides)
+    }
+    fn index_unchecked(&self, strides: &IxDyn) -> isize {
+        (&**self).index_unchecked(strides)
     }
 }