ARROW-12398: [Rust] remove redundant bound check in iterators

This PR removes the bound checks as discussed in #9994.

Furthermore I added `unsafe` versions of the `value` method to `PrimitiveArray` and `BooleanArray`. The `safe` marked methods are actually `unsafe`. This way we can slowly transition to explicitly using the `unsafe` variant and later make the "safe" one truly safe.

For the time being I also added a `debug_assert` bounds check in those "safe" methods that are `unsafe`. That way we at least get a panic in debug mode instead of UB in safe code.

Closes #10046 from ritchie46/iterator_bounds

Authored-by: Ritchie Vink <[email protected]>
Signed-off-by: Andrew Lamb <[email protected]>

Author: ritchie46

rust/arrow/src/array/array_boolean.rs

@@ -67,12 +67,21 @@ impl BooleanArray {
         &self.data.buffers()[0]
     }
 
+    /// Returns the boolean value at index `i`.
+    ///
+    /// # Safety
+    /// This doesn't check bounds, the caller must ensure that index < self.len()
+    pub unsafe fn value_unchecked(&self, i: usize) -> bool {
+        let offset = i + self.offset();
+        bit_util::get_bit_raw(self.raw_values.as_ptr(), offset)
+    }
+
     /// Returns the boolean value at index `i`.
     ///
     /// Note this doesn't do any bound checking, for performance reason.
     pub fn value(&self, i: usize) -> bool {
-        let offset = i + self.offset();
-        unsafe { bit_util::get_bit_raw(self.raw_values.as_ptr(), offset) }
+        debug_assert!(i < self.len());
+        unsafe { self.value_unchecked(i) }
     }
 }
 

rust/arrow/src/array/array_primitive.rs

@@ -88,14 +88,24 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
         PrimitiveBuilder::<T>::new(capacity)
     }
 
+    /// Returns the primitive value at index `i`.
+    ///
+    /// # Safety
+    ///
+    /// caller must ensure that the passed in offset is less than the array len()
+    pub unsafe fn value_unchecked(&self, i: usize) -> T::Native {
+        let offset = i + self.offset();
+        *self.raw_values.as_ptr().add(offset)
+    }
+
     /// Returns the primitive value at index `i`.
     ///
     /// Note this doesn't do any bound checking, for performance reason.
     /// # Safety
     /// caller must ensure that the passed in offset is less than the array len()
     pub fn value(&self, i: usize) -> T::Native {
-        let offset = i + self.offset();
-        unsafe { *self.raw_values.as_ptr().add(offset) }
+        debug_assert!(i < self.len());
+        unsafe { self.value_unchecked(i) }
     }
 
     /// Creates a PrimitiveArray based on an iterator of values without nulls

rust/arrow/src/array/iterator.rs

@@ -56,7 +56,12 @@ impl<'a, T: ArrowPrimitiveType> std::iter::Iterator for PrimitiveIter<'a, T> {
         } else {
             let old = self.current;
             self.current += 1;
-            Some(Some(self.array.value(old)))
+            // Safety:
+            // we just checked bounds in `self.current_end == self.current`
+            // this is safe on the premise that this struct is initialized with
+            // current = array.len()
+            // and that current_end is ever only decremented
+            unsafe { Some(Some(self.array.value_unchecked(old))) }
         }
     }
 
@@ -77,7 +82,12 @@ impl<'a, T: ArrowPrimitiveType> std::iter::DoubleEndedIterator for PrimitiveIter
             Some(if self.array.is_null(self.current_end) {
                 None
             } else {
-                Some(self.array.value(self.current_end))
+                // Safety:
+                // we just checked bounds in `self.current_end == self.current`
+                // this is safe on the premise that this struct is initialized with
+                // current = array.len()
+                // and that current_end is ever only decremented
+                unsafe { Some(self.array.value_unchecked(self.current_end)) }
             })
         }
     }
@@ -118,7 +128,12 @@ impl<'a> std::iter::Iterator for BooleanIter<'a> {
         } else {
             let old = self.current;
             self.current += 1;
-            Some(Some(self.array.value(old)))
+            // Safety:
+            // we just checked bounds in `self.current_end == self.current`
+            // this is safe on the premise that this struct is initialized with
+            // current = array.len()
+            // and that current_end is ever only decremented
+            unsafe { Some(Some(self.array.value_unchecked(old))) }
         }
     }
 
@@ -139,7 +154,12 @@ impl<'a> std::iter::DoubleEndedIterator for BooleanIter<'a> {
             Some(if self.array.is_null(self.current_end) {
                 None
             } else {
-                Some(self.array.value(self.current_end))
+                // Safety:
+                // we just checked bounds in `self.current_end == self.current`
+                // this is safe on the premise that this struct is initialized with
+                // current = array.len()
+                // and that current_end is ever only decremented
+                unsafe { Some(self.array.value_unchecked(self.current_end)) }
             })
         }
     }
@@ -182,7 +202,12 @@ impl<'a, T: StringOffsetSizeTrait> std::iter::Iterator for GenericStringIter<'a,
             Some(None)
         } else {
             self.current += 1;
-            Some(Some(self.array.value(i)))
+            // Safety:
+            // we just checked bounds in `self.current_end == self.current`
+            // this is safe on the premise that this struct is initialized with
+            // current = array.len()
+            // and that current_end is ever only decremented
+            unsafe { Some(Some(self.array.value_unchecked(i))) }
         }
     }
 
@@ -205,7 +230,12 @@ impl<'a, T: StringOffsetSizeTrait> std::iter::DoubleEndedIterator
             Some(if self.array.is_null(self.current_end) {
                 None
             } else {
-                Some(self.array.value(self.current_end))
+                // Safety:
+                // we just checked bounds in `self.current_end == self.current`
+                // this is safe on the premise that this struct is initialized with
+                // current = array.len()
+                // and that current_end is ever only decremented
+                unsafe { Some(self.array.value_unchecked(self.current_end)) }
             })
         }
     }
@@ -251,7 +281,12 @@ impl<'a, T: BinaryOffsetSizeTrait> std::iter::Iterator for GenericBinaryIter<'a,
             Some(None)
         } else {
             self.current += 1;
-            Some(Some(self.array.value(i)))
+            // Safety:
+            // we just checked bounds in `self.current_end == self.current`
+            // this is safe on the premise that this struct is initialized with
+            // current = array.len()
+            // and that current_end is ever only decremented
+            unsafe { Some(Some(self.array.value_unchecked(i))) }
         }
     }
 
@@ -274,7 +309,12 @@ impl<'a, T: BinaryOffsetSizeTrait> std::iter::DoubleEndedIterator
             Some(if self.array.is_null(self.current_end) {
                 None
             } else {
-                Some(self.array.value(self.current_end))
+                // Safety:
+                // we just checked bounds in `self.current_end == self.current`
+                // this is safe on the premise that this struct is initialized with
+                // current = array.len()
+                // and that current_end is ever only decremented
+                unsafe { Some(self.array.value_unchecked(self.current_end)) }
             })
         }
     }
@@ -318,7 +358,12 @@ impl<'a, S: OffsetSizeTrait> std::iter::Iterator for GenericListArrayIter<'a, S>
             Some(None)
         } else {
             self.current += 1;
-            Some(Some(self.array.value(i)))
+            // Safety:
+            // we just checked bounds in `self.current_end == self.current`
+            // this is safe on the premise that this struct is initialized with
+            // current = array.len()
+            // and that current_end is ever only decremented
+            unsafe { Some(Some(self.array.value_unchecked(i))) }
         }
     }
 
@@ -341,7 +386,12 @@ impl<'a, S: OffsetSizeTrait> std::iter::DoubleEndedIterator
             Some(if self.array.is_null(self.current_end) {
                 None
             } else {
-                Some(self.array.value(self.current_end))
+                // Safety:
+                // we just checked bounds in `self.current_end == self.current`
+                // this is safe on the premise that this struct is initialized with
+                // current = array.len()
+                // and that current_end is ever only decremented
+                unsafe { Some(self.array.value_unchecked(self.current_end)) }
             })
         }
     }