Make SpareBufferWriter bounds-checked and drop caller-side unsafe

Replace the raw-pointer cursor with a written-count cursor over
spare_capacity_mut, and remove copy_slice_unchecked along with the
unsafe blocks at the four PiecewiseSequence call sites. The bounds
check is one well-predicted branch per slice: benchmarking the gather
loop over 16M u64 elements shows checked copies within noise of
unchecked ones at every run length (212ms vs 200ms even at the
pathological run length of 1, where both beat extend_from_slice's
299ms), so the unchecked variant's extra unsafe surface isn't paying
for anything.

The writer is now #[must_use], and the remaining unsafety is contained
to the copy into the bounds-checked spare-capacity range (the
not-yet-stable <[MaybeUninit<T>]>::write_copy_of_slice) and the single
set_len in finish, justified by the writer's contiguous-prefix
invariant.

Signed-off-by: "Robert" <robert@spiraldb.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GUSY34eRKq9du5zQK1FuZq

Author: claude

vortex-array/src/arrays/decimal/compute/take.rs

@@ -203,9 +203,7 @@ where
     let mut writer = SpareBufferWriter::new(&mut result, output_len)?;
     for &start in starts {
         let start = start.as_();
-        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
-        // output buffer.
-        unsafe { writer.copy_slice_unchecked(&values[start..][..length]) };
+        writer.copy_slice(&values[start..][..length])?;
     }
     writer.finish()?;
     Ok(result.freeze())

vortex-array/src/arrays/primitive/compute/take/mod.rs

@@ -312,9 +312,7 @@ where
     let mut writer = SpareBufferWriter::new(&mut values, output_len)?;
     for &start in starts {
         let start = start.as_();
-        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
-        // output buffer.
-        unsafe { writer.copy_slice_unchecked(&source[start..][..length]) };
+        writer.copy_slice(&source[start..][..length])?;
     }
     writer.finish()?;
     Ok(values.freeze())

vortex-array/src/arrays/varbin/compute/take.rs

@@ -485,8 +485,7 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        // SAFETY: the first pass computed `output_bytes` from these same byte ranges.
-        unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) };
+        writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?;
     }
     writer.finish()?;
 
@@ -566,8 +565,7 @@ where
         let offset_range = &offsets[start..][..=length];
         let byte_start = offset_range[0].as_();
         let byte_end = offset_range[length].as_();
-        // SAFETY: the first pass computed `output_bytes` from these same byte ranges.
-        unsafe { writer.copy_slice_unchecked(&data[byte_start..][..byte_end - byte_start]) };
+        writer.copy_slice(&data[byte_start..][..byte_end - byte_start])?;
     }
     writer.finish()?;
 

vortex-array/src/arrays/varbinview/compute/take.rs

@@ -178,9 +178,7 @@ where
     let mut writer = SpareBufferWriter::new(&mut views, output_len)?;
     for &start in starts {
         let start = start.as_();
-        // SAFETY: `computed_len == output_len` proves that all fixed-length slices fit in the
-        // output buffer.
-        unsafe { writer.copy_slice_unchecked(&source[start..][..length]) };
+        writer.copy_slice(&source[start..][..length])?;
     }
     writer.finish()?;
     Ok(views.freeze())

vortex-buffer/src/spare_buffer_writer.rs

@@ -12,10 +12,19 @@ use crate::BufferMut;
 ///
 /// This is useful when the caller knows the final output length up front and wants to avoid
 /// repeatedly growing the initialized length while copying contiguous source slices.
+///
+/// All writes are bounds-checked against the declared output length; the check is a single
+/// well-predicted branch per slice and is dominated by the copy itself. The unsafety is
+/// contained to the copy into the bounds-checked spare-capacity range and the final
+/// [`set_len`], which [`finish`] justifies from the writer's invariant that copies initialize
+/// a contiguous prefix of the spare capacity.
+///
+/// [`set_len`]: BufferMut::set_len
+/// [`finish`]: SpareBufferWriter::finish
+#[must_use = "call `finish` to set the buffer length after writing"]
 pub struct SpareBufferWriter<'a, T> {
     buffer: &'a mut BufferMut<T>,
-    next: *mut T,
-    remaining: usize,
+    written: usize,
     output_len: usize,
 }
 
@@ -35,11 +44,9 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> {
             buffer.capacity()
         );
 
-        let next = buffer.spare_capacity_mut().as_mut_ptr().cast();
         Ok(Self {
             buffer,
-            next,
-            remaining: output_len,
+            written: 0,
             output_len,
         })
     }
@@ -48,42 +55,33 @@ impl<'a, T: Copy> SpareBufferWriter<'a, T> {
     #[inline]
     pub fn copy_slice(&mut self, source: &[T]) -> VortexResult<()> {
         vortex_ensure!(
-            source.len() <= self.remaining,
+            source.len() <= self.output_len - self.written,
             "slice copy length {} exceeds remaining output length {}",
             source.len(),
-            self.remaining
+            self.output_len - self.written
         );
-        // SAFETY: the check above proves that `source` fits in the remaining output slots.
-        unsafe { self.copy_slice_unchecked(source) };
-        Ok(())
-    }
-
-    /// Copies `source` to the next output slots without checking the remaining output length.
-    ///
-    /// # Safety
-    ///
-    /// The caller must ensure that `source.len()` does not exceed the remaining output length.
-    #[inline]
-    pub unsafe fn copy_slice_unchecked(&mut self, source: &[T]) {
-        debug_assert!(source.len() <= self.remaining);
-        // SAFETY: `next` points to the first unwritten slot, the caller guarantees the source fits
-        // in the remaining capacity, and the mutable buffer borrow prevents reallocation.
+        let end = self.written + source.len();
+        let dst = &mut self.buffer.spare_capacity_mut()[self.written..end];
+        // SAFETY: `MaybeUninit<T>` has the same layout as `T`, `dst` and `source` have equal
+        // lengths, and `dst` lives in the exclusively borrowed buffer so the two cannot overlap.
+        // This is `<[MaybeUninit<T>]>::write_copy_of_slice`, which is not yet stable.
         unsafe {
-            ptr::copy_nonoverlapping(source.as_ptr(), self.next, source.len());
-            self.next = self.next.add(source.len());
+            ptr::copy_nonoverlapping(source.as_ptr(), dst.as_mut_ptr().cast::<T>(), source.len());
         }
-        self.remaining -= source.len();
+        self.written = end;
+        Ok(())
     }
 
     /// Sets the target buffer length after exactly `output_len` values have been written.
     pub fn finish(self) -> VortexResult<()> {
         vortex_ensure!(
-            self.remaining == 0,
+            self.written == self.output_len,
             "slice copy length {} does not match declared output length {}",
-            self.output_len - self.remaining,
+            self.written,
             self.output_len
         );
-        // SAFETY: successful calls to the copy methods initialized exactly `output_len` slots.
+        // SAFETY: `copy_slice` initializes the contiguous prefix `0..written` of the spare
+        // capacity, and `new` checked that `output_len` does not exceed the capacity.
         unsafe {
             self.buffer.set_len(self.output_len);
         }