Small buffer optimization for InvertedAxes

Since InvertedAxes is basically a set of small integers, we can avoid dynamic
allocations by using a word-sized bitmap in the common case of up to 32 axes.

This should also not unduly increase code slice as there is no dynamic spilling
into the large representation, but the decision is rather taken once based on
the dimensionality of the array.

Author: adamreichold

src/array.rs

@@ -358,12 +358,46 @@ impl<T, D> PyArray<T, D> {
     }
 }
 
-struct InvertedAxises(Vec<Axis>);
+enum InvertedAxes {
+    Short(u32),
+    Long(Vec<usize>),
+}
+
+impl InvertedAxes {
+    fn new(len: usize) -> Self {
+        if len <= 32 {
+            Self::Short(0)
+        } else {
+            Self::Long(Vec::new())
+        }
+    }
+
+    fn push(&mut self, axis: usize) {
+        match self {
+            Self::Short(axes) => {
+                debug_assert!(axis < 32);
+                *axes |= 1 << axis;
+            }
+            Self::Long(axes) => {
+                axes.push(axis);
+            }
+        }
+    }
 
-impl InvertedAxises {
     fn invert<S: RawData, D: Dimension>(self, array: &mut ArrayBase<S, D>) {
-        for axis in self.0 {
-            array.invert_axis(axis);
+        match self {
+            Self::Short(mut axes) => {
+                while axes != 0 {
+                    let axis = axes.trailing_zeros() as usize;
+                    axes &= !(1 << axis);
+                    array.invert_axis(Axis(axis));
+                }
+            }
+            Self::Long(axes) => {
+                for axis in axes {
+                    array.invert_axis(Axis(axis));
+                }
+            }
         }
     }
 }
@@ -372,36 +406,39 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     /// Same as [shape](#method.shape), but returns `D`
     #[inline(always)]
     pub fn dims(&self) -> D {
-        D::from_dimension(&Dim(self.shape())).expect("PyArray::dims different dimension")
+        D::from_dimension(&Dim(self.shape())).expect("mismatching dimensions")
     }
 
-    fn ndarray_shape_ptr(&self) -> (StrideShape<D>, *mut T, InvertedAxises) {
-        let shape_slice = self.shape();
-        let shape: Shape<_> = Dim(self.dims()).into();
-        let sizeof_t = mem::size_of::<T>();
+    fn ndarray_shape_ptr(&self) -> (StrideShape<D>, *mut T, InvertedAxes) {
+        let shape = self.shape();
         let strides = self.strides();
+
         let mut new_strides = D::zeros(strides.len());
         let mut data_ptr = unsafe { self.data() };
-        let mut inverted_axises = vec![];
+        let mut inverted_axes = InvertedAxes::new(strides.len());
+
         for i in 0..strides.len() {
             // TODO(kngwyu): Replace this hacky negative strides support with
             // a proper constructor, when it's implemented.
             // See https://github.com/rust-ndarray/ndarray/issues/842 for more.
             if strides[i] < 0 {
                 // Move the pointer to the start position
-                let offset = strides[i] * (shape_slice[i] as isize - 1) / sizeof_t as isize;
+                let offset = strides[i] * (shape[i] as isize - 1) / mem::size_of::<T>() as isize;
                 unsafe {
                     data_ptr = data_ptr.offset(offset);
                 }
-                new_strides[i] = (-strides[i]) as usize / sizeof_t;
-                inverted_axises.push(Axis(i));
+                new_strides[i] = (-strides[i]) as usize / mem::size_of::<T>();
+
+                inverted_axes.push(i);
             } else {
-                new_strides[i] = strides[i] as usize / sizeof_t;
+                new_strides[i] = strides[i] as usize / mem::size_of::<T>();
             }
         }
-        let st = D::from_dimension(&Dim(new_strides))
-            .expect("PyArray::ndarray_shape: dimension mismatching");
-        (shape.strides(st), data_ptr, InvertedAxises(inverted_axises))
+
+        let shape = Shape::from(D::from_dimension(&Dim(shape)).expect("mismatching dimensions"));
+        let new_strides = D::from_dimension(&Dim(new_strides)).expect("mismatching dimensions");
+
+        (shape.strides(new_strides), data_ptr, inverted_axes)
     }
 
     /// Creates a new uninitialized PyArray in python heap.
@@ -818,9 +855,9 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     /// If the internal array is not readonly and can be mutated from Python code,
     /// holding the `ArrayView` might cause undefined behavior.
     pub unsafe fn as_array(&self) -> ArrayView<'_, T, D> {
-        let (shape, ptr, inverted_axises) = self.ndarray_shape_ptr();
+        let (shape, ptr, inverted_axes) = self.ndarray_shape_ptr();
         let mut res = ArrayView::from_shape_ptr(shape, ptr);
-        inverted_axises.invert(&mut res);
+        inverted_axes.invert(&mut res);
         res
     }
 
@@ -830,25 +867,25 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     /// If another reference to the internal data exists(e.g., `&[T]` or `ArrayView`),
     /// it might cause undefined behavior.
     pub unsafe fn as_array_mut(&self) -> ArrayViewMut<'_, T, D> {
-        let (shape, ptr, inverted_axises) = self.ndarray_shape_ptr();
+        let (shape, ptr, inverted_axes) = self.ndarray_shape_ptr();
         let mut res = ArrayViewMut::from_shape_ptr(shape, ptr);
-        inverted_axises.invert(&mut res);
+        inverted_axes.invert(&mut res);
         res
     }
 
     /// Returns the internal array as [`RawArrayView`] enabling element access via raw pointers
     pub fn as_raw_array(&self) -> RawArrayView<T, D> {
-        let (shape, ptr, inverted_axises) = self.ndarray_shape_ptr();
+        let (shape, ptr, inverted_axes) = self.ndarray_shape_ptr();
         let mut res = unsafe { RawArrayView::from_shape_ptr(shape, ptr) };
-        inverted_axises.invert(&mut res);
+        inverted_axes.invert(&mut res);
         res
     }
 
     /// Returns the internal array as [`RawArrayViewMut`] enabling element access via raw pointers
     pub fn as_raw_array_mut(&self) -> RawArrayViewMut<T, D> {
-        let (shape, ptr, inverted_axises) = self.ndarray_shape_ptr();
+        let (shape, ptr, inverted_axes) = self.ndarray_shape_ptr();
         let mut res = unsafe { RawArrayViewMut::from_shape_ptr(shape, ptr) };
-        inverted_axises.invert(&mut res);
+        inverted_axes.invert(&mut res);
         res
     }