diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index fcb572c3159..28dc061972d 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -87,7 +87,7 @@ where let exponents = ::alp::encode_into( values_slice, exponents, - &mut encoded.spare_capacity_mut()[..values_slice.len()], + encoded.spare_capacity_mut(values_slice.len()), &mut exceptional_positions, &mut exceptional_values, &mut chunk_offsets, diff --git a/encodings/fastlanes/src/bit_transpose.rs b/encodings/fastlanes/src/bit_transpose.rs index b9761ee20cf..0abf6ca220c 100644 --- a/encodings/fastlanes/src/bit_transpose.rs +++ b/encodings/fastlanes/src/bit_transpose.rs @@ -65,13 +65,8 @@ fn bits_op_with_copy( let output_len = bytes.len().div_ceil(8).next_multiple_of(16); let mut output = BufferMut::::with_capacity(output_len); let (input_chunks, input_trailer) = bytes.as_chunks::<128>(); - // Bound to the requested `output_len`: `spare_capacity_mut` may expose extra over-aligned - // capacity, which would otherwise split into spurious trailing chunks and make `last_mut` - // below target a chunk past the data we actually initialize. let (output_chunks, _) = unsafe { - mem::transmute::<&mut [MaybeUninit], &mut [u64]>( - &mut output.spare_capacity_mut()[..output_len], - ) + mem::transmute::<&mut [MaybeUninit], &mut [u64]>(output.spare_capacity_mut(output_len)) } .as_chunks_mut::<16>(); diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index c2ef38ceb79..e111b6b1c87 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -86,7 +86,9 @@ where // Allocate result arrays. let mut bases = BufferMut::with_capacity(bases_len); let mut deltas = BufferMut::with_capacity(padded_len); - let (output_deltas, _) = deltas.spare_capacity_mut().as_chunks_mut::(); + let (output_deltas, _) = deltas + .spare_capacity_mut(padded_len) + .as_chunks_mut::(); // Loop over all full 1024-element chunks. let mut transposed: [T; FL_CHUNK_SIZE] = [T::default(); FL_CHUNK_SIZE]; diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index 3ccfede27c1..aa4adaac143 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -69,9 +69,9 @@ where // Allocate a result array. let mut output = BufferMut::with_capacity(deltas.len()); - // Bound to the requested length: `spare_capacity_mut` may expose extra over-aligned capacity - // beyond `deltas.len()`, which would desync the `zip_eq` with `chunks` below and panic. - let (output_chunks, _) = output.spare_capacity_mut()[..deltas.len()].as_chunks_mut::<1024>(); + let (output_chunks, _) = output + .spare_capacity_mut(deltas.len()) + .as_chunks_mut::<1024>(); // Loop over all the chunks let mut transposed: [T; 1024] = [T::default(); 1024]; diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..af46a2dd87d 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -157,7 +157,7 @@ pub(crate) fn fill_forward_nulls( .iter() .zip( to_fill_mut - .spare_capacity_mut() + .spare_capacity_mut(to_fill.len()) .iter_mut() .zip(bit_buffer.iter()), ) diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index f623b84d778..8c728323ceb 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -58,10 +58,9 @@ where // Pre-allocate for one offset per chunk. let mut values_idx_offsets = BufferMut::::with_capacity(len.div_ceil(FL_CHUNK_SIZE)); - let values_uninit = values_buf.spare_capacity_mut(); - // We don't care about the trailing chunk that exists due to overallocation by the underlying allocator. + let values_uninit = values_buf.spare_capacity_mut(padded_len); let (indices_uninit, _) = indices_buf - .spare_capacity_mut() + .spare_capacity_mut(padded_len) .as_chunks_mut::(); let mut value_count_acc = 0; // Chunk value count prefix sum. @@ -104,8 +103,6 @@ where // accounting for an additional value change. let mut padded_chunk = [values[len - 1]; FL_CHUNK_SIZE]; padded_chunk[..remainder.len()].copy_from_slice(remainder); - // There might be more entries in indices_uninit than necessary if the allocator gave us extra memory. - // Remainder has to go to the last chunk after full chunks have been processed. let last_idx_chunk = &mut indices_uninit[chunks.len()]; process_chunk(&padded_chunk, last_idx_chunk); } diff --git a/encodings/fastlanes/src/rle/array/rle_decompress.rs b/encodings/fastlanes/src/rle/array/rle_decompress.rs index 0c841b7cba6..08309f38df5 100644 --- a/encodings/fastlanes/src/rle/array/rle_decompress.rs +++ b/encodings/fastlanes/src/rle/array/rle_decompress.rs @@ -91,7 +91,9 @@ where let num_chunks = chunk_end_idx - chunk_start_idx; let mut buffer = BufferMut::::with_capacity(num_chunks * FL_CHUNK_SIZE); - let (out_buf, _) = buffer.spare_capacity_mut().as_chunks_mut::(); + let (out_buf, _) = buffer + .spare_capacity_mut(num_chunks * FL_CHUNK_SIZE) + .as_chunks_mut::(); for (chunk_idx, (chunk_indices, chunk_out)) in indices_sl.iter().zip(out_buf.iter_mut()).enumerate() diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 28942648aa4..c638d26028e 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -121,7 +121,7 @@ pub(crate) fn fsst_decode_bytes( let mut uncompressed_bytes = ByteBufferMut::with_capacity(plan.total_size + FSST_DECODE_SLACK); let len = plan.decode_into( &fsst_array.decompressor(), - uncompressed_bytes.spare_capacity_mut(), + uncompressed_bytes.spare_capacity_mut(plan.total_size + FSST_DECODE_SLACK), )?; // SAFETY: `decode_into` initialized the first `len` bytes. unsafe { uncompressed_bytes.set_len(len) }; diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs index 7ed5bfe6aa6..0509e4aedae 100644 --- a/encodings/onpair/src/canonical.rs +++ b/encodings/onpair/src/canonical.rs @@ -145,7 +145,7 @@ pub(crate) fn onpair_decode_bytes( ) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { let plan = OnPairDecodePlan::new(array, ctx)?; let mut out_bytes = ByteBufferMut::with_capacity(plan.total_size); - let written = plan.decode_into(out_bytes.spare_capacity_mut())?; + let written = plan.decode_into(out_bytes.spare_capacity_mut(plan.total_size))?; // SAFETY: `decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; Ok((out_bytes, plan.lengths)) diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index 81dda5940fd..9622c862fd5 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1401,8 +1401,8 @@ impl ZstdData { // the ones before it, bounded by the size the metadata declared, so a frame that // expands further than advertised is refused by zstd rather than overrunning. let mut destination = UninitDestination::new( - &mut decompressed.spare_capacity_mut() - [uncompressed_start..uncompressed_size_to_decompress], + &mut decompressed.spare_capacity_mut(uncompressed_size_to_decompress) + [uncompressed_start..], ); uncompressed_start += decompressor.decompress_to_buffer(frame.as_slice(), &mut destination)?; diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index f21deee6f64..2d817510c39 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -262,7 +262,7 @@ impl ZstdBuffersData { let compressed = buf.clone().try_to_host_sync()?; validate_frame_content_size(compressed.as_slice(), uncompressed_size, i)?; let mut output = ByteBufferMut::with_capacity_aligned(size, aligned); - let spare = output.spare_capacity_mut(); + let spare = output.spare_capacity_mut(size); // This is currently guaranteed, but still good to check because // of the unsafe calls below. diff --git a/vortex-array/benches/take_slices_to_buffer_matrix.rs b/vortex-array/benches/take_slices_to_buffer_matrix.rs index 21cb539ce32..b02d3e3b466 100644 --- a/vortex-array/benches/take_slices_to_buffer_matrix.rs +++ b/vortex-array/benches/take_slices_to_buffer_matrix.rs @@ -191,7 +191,10 @@ fn take_advancing_ptr_safe( ) -> Buffer { let mut result = BufferMut::::with_capacity(output_len); let mut cursor = 0usize; - let mut dst = result.spare_capacity_mut().as_mut_ptr().cast::(); + let mut dst = result + .spare_capacity_mut(output_len) + .as_mut_ptr() + .cast::(); for (&start, &length) in starts.iter().zip(lengths) { let end = cursor.checked_add(length).unwrap(); assert!(end <= output_len); @@ -261,7 +264,10 @@ fn take_preverify_advancing_ptr_unchecked( preverify(values.len(), starts, lengths, output_len); let mut result = BufferMut::::with_capacity(output_len); - let mut dst = result.spare_capacity_mut().as_mut_ptr().cast::(); + let mut dst = result + .spare_capacity_mut(output_len) + .as_mut_ptr() + .cast::(); for (&start, &length) in starts.iter().zip(lengths) { // SAFETY: `preverify` checked every source range and the summed output length. unsafe { @@ -289,7 +295,7 @@ fn preverify(source_len: usize, starts: &[usize], lengths: &[usize], output_len: } fn copy_to_spare(result: &mut BufferMut, cursor: usize, source: &[u16]) { - let dst = &mut result.spare_capacity_mut()[cursor..][..source.len()]; + let dst = &mut result.spare_capacity_mut(cursor + source.len())[cursor..]; // SAFETY: `dst` has exactly `source.len()` spare slots and does not overlap with source. unsafe { copy_to_uninit(dst.as_mut_ptr().cast(), source) }; } @@ -298,7 +304,7 @@ unsafe fn copy_to_spare_unchecked(result: &mut BufferMut, cursor: usize, so // SAFETY: callers ensure `cursor..cursor + source.len()` is within spare capacity. let dst = unsafe { result - .spare_capacity_mut() + .spare_capacity_mut(result.capacity() - result.len()) .get_unchecked_mut(cursor..cursor + source.len()) }; // SAFETY: `dst` has exactly `source.len()` spare slots and does not overlap with source. diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 8770a9bc8f4..8822525347a 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -171,7 +171,7 @@ fn cast_to_f64( let values = array.buffer::(); let values = values.as_slice(); let mut out = BufferMut::::with_capacity(n); - values.map_into(&mut out.spare_capacity_mut()[..n], |v: F| { + values.map_into(out.spare_capacity_mut(n), |v: F| { to_f64_lossy::(v) * inv_factor }); // SAFETY: map_into wrote every lane before returning. @@ -184,7 +184,7 @@ fn cast_to_f64( let mut out = BufferMut::::with_capacity(n); let write_result = values.try_map_masked_into( mask_values.bit_buffer(), - &mut out.spare_capacity_mut()[..n], + out.spare_capacity_mut(n), |v: F| Some(to_f64_lossy::(v) * inv_factor), ); debug_assert!(write_result.is_ok()); @@ -244,13 +244,13 @@ where let mut buffer = BufferMut::::with_capacity(values.len()); match valid_values { Mask::AllTrue(_) => { - values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; + values.try_map_into(buffer.spare_capacity_mut(values.len()), &cast)?; } Mask::AllFalse(_) => return Ok(BufferMut::::zeroed(values.len()).freeze()), Mask::Values(mask) => { values.try_map_masked_into( mask.bit_buffer(), - &mut buffer.spare_capacity_mut()[..values.len()], + buffer.spare_capacity_mut(values.len()), &cast, )?; } diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 032b11a56b3..a29cf8f76ae 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -121,7 +121,9 @@ fn filter_chunk_into( return; } - let out_ptr = out.spare_capacity_mut().as_mut_ptr(); + let out_ptr = out + .spare_capacity_mut(out.capacity() - out.len()) + .as_mut_ptr(); if chunk.len() == 8 && mask_byte == 0xFF { // All 8 selected, so bulk copy. // SAFETY: write_pos + 8 <= capacity. diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs index 807259cd5e7..8b0954473f8 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs +++ b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs @@ -62,7 +62,9 @@ pub(super) fn filter_slice_by_bitmap( let written = unsafe { kernel( values.as_ptr().cast(), - out.spare_capacity_mut().as_mut_ptr().cast(), + out.spare_capacity_mut(true_count + SLACK_BYTES / size_of::()) + .as_mut_ptr() + .cast(), mask, ) }; diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 52d1328c92e..7100b058a64 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -64,7 +64,7 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> let output_len = mask.true_count(); let mut out = BufferMut::::with_capacity(output_len); let src_ptr = slice.as_ptr(); - let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let out_ptr = out.spare_capacity_mut(output_len).as_mut_ptr().cast::(); let mut write_pos = 0; for_each_mask_word(mask, |word, word_start, word_len| { @@ -102,7 +102,10 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { let mut out = BufferMut::::with_capacity(indices.len()); let src_ptr = slice.as_ptr(); - let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let out_ptr = out + .spare_capacity_mut(indices.len()) + .as_mut_ptr() + .cast::(); for (write_pos, &index) in indices.iter().enumerate() { // SAFETY: mask indices are validated when the mask is constructed and the output has one diff --git a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs index 6518293db41..6a289c96a90 100644 --- a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs +++ b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs @@ -207,7 +207,7 @@ where L: Fn(usize) -> usize, { let mut out = BufferMut::::with_capacity(ranks.len()); - let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let out_ptr = out.spare_capacity_mut(ranks.len()).as_mut_ptr().cast::(); for (idx, rank) in ranks.iter().enumerate() { let value = if ranks_validity.value(idx) { let rank = validate_rank(*rank, translated_len)?; @@ -240,7 +240,7 @@ where L: Fn(usize) -> usize, { let mut out = BufferMut::::with_capacity(ranks.len()); - let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let out_ptr = out.spare_capacity_mut(ranks.len()).as_mut_ptr().cast::(); for (idx, rank) in ranks.iter().enumerate() { let rank = validate_rank(*rank, translated_len)?; let child_idx = translate(rank); diff --git a/vortex-array/src/arrays/filter/execute/take/rank.rs b/vortex-array/src/arrays/filter/execute/take/rank.rs index 08170ab85d3..6f96480af19 100644 --- a/vortex-array/src/arrays/filter/execute/take/rank.rs +++ b/vortex-array/src/arrays/filter/execute/take/rank.rs @@ -103,7 +103,10 @@ where L: Fn(usize) -> usize, { let mut translated = BufferMut::::with_capacity(ranks.len()); - let translated_ptr = translated.spare_capacity_mut().as_mut_ptr().cast::(); + let translated_ptr = translated + .spare_capacity_mut(ranks.len()) + .as_mut_ptr() + .cast::(); for (idx, rank) in ranks.iter().enumerate() { let translated_rank = match ranks_validity { diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs index cba089f1011..10a7aa886bb 100644 --- a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs @@ -114,7 +114,7 @@ where let max_index = Idx::from(values.len()); let mut buffer = BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); - let buf_uninit = buffer.spare_capacity_mut(); + let buf_uninit = buffer.spare_capacity_mut(indices_len); let mut offset = 0; // SAFETY: `exec_take` is only called by `take_avx2`, whose caller guarantees AVX2 support. diff --git a/vortex-array/src/arrays/fixed_width/take/scalar.rs b/vortex-array/src/arrays/fixed_width/take/scalar.rs index 5a3510efc2d..986ac417ce4 100644 --- a/vortex-array/src/arrays/fixed_width/take/scalar.rs +++ b/vortex-array/src/arrays/fixed_width/take/scalar.rs @@ -15,7 +15,10 @@ pub(crate) fn take_values_scalar( // The explicit pointer loop keeps the source length in a register and avoids a capacity check // for every output value. let mut result = BufferMut::with_capacity(indices.len()); - let result_ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); + let result_ptr = result + .spare_capacity_mut(indices.len()) + .as_mut_ptr() + .cast::(); for (output_index, index) in indices.iter().enumerate() { // SAFETY: `indices.len()` elements were reserved and each output position is written once. diff --git a/vortex-array/src/arrays/fixed_width/take/slices.rs b/vortex-array/src/arrays/fixed_width/take/slices.rs index a4d9ceda84a..181d03dc949 100644 --- a/vortex-array/src/arrays/fixed_width/take/slices.rs +++ b/vortex-array/src/arrays/fixed_width/take/slices.rs @@ -71,7 +71,7 @@ fn copy_slices( .checked_mul(byte_width) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; let mut result = BufferMut::::with_capacity_aligned(output_byte_len, values.alignment()); - let spare = &mut result.spare_capacity_mut()[..output_byte_len]; + let spare = result.spare_capacity_mut(output_byte_len); let mut cursor = 0usize; for (start, length) in slices { diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index d5e2a6bf406..126a816db5e 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -93,8 +93,8 @@ impl ZipKernel for ListView { let false_offsets = false_offsets.as_slice(); let false_sizes = false_sizes.as_slice(); - let offsets_out = offsets.spare_capacity_mut(); - let sizes_out = sizes.spare_capacity_mut(); + let offsets_out = offsets.spare_capacity_mut(len); + let sizes_out = sizes.spare_capacity_mut(len); // We matched `Mask::Values` above, so the bit buffer is materialized. `unaligned_chunks` // iterates faster than `chunks`: it exposes the byte-aligned body as a plain `&[u64]` diff --git a/vortex-array/src/arrays/patched/array.rs b/vortex-array/src/arrays/patched/array.rs index 4522c8223b8..0ead683f394 100644 --- a/vortex-array/src/arrays/patched/array.rs +++ b/vortex-array/src/arrays/patched/array.rs @@ -331,8 +331,8 @@ fn transpose( } // Loop over patches, writing them to final positions. - let indices_out = indices_buffer.spare_capacity_mut(); - let values_out = values_buffer.spare_capacity_mut(); + let indices_out = indices_buffer.spare_capacity_mut(indices_in.len()); + let values_out = values_buffer.spare_capacity_mut(values_in.len()); for (index, &value) in std::iter::zip(indices_in, values_in) { let index = index.as_() - offset; let chunk = index / 1024; diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 93f10bdc5cf..43e6b09498c 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -415,12 +415,12 @@ where let mut buffer = BufferMut::::with_capacity(values.len()); match valid_values { Mask::AllTrue(_) => { - values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?; + values.try_map_into(buffer.spare_capacity_mut(values.len()), &cast)?; } Mask::Values(mask) => { values.try_map_masked_into( mask.bit_buffer(), - &mut buffer.spare_capacity_mut()[..values.len()], + buffer.spare_capacity_mut(values.len()), &cast, )?; } @@ -526,7 +526,7 @@ where } None => { let mut buffer = BufferMut::::with_capacity(len); - values.map_into(&mut buffer.spare_capacity_mut()[..len], |v| v.as_()); + values.map_into(buffer.spare_capacity_mut(len), |v| v.as_()); // SAFETY: map_into initializes every lane. unsafe { buffer.set_len(len) }; Ok(PrimitiveArray::new(buffer.freeze(), new_validity).into_array()) @@ -548,9 +548,7 @@ where (Mask::AllTrue(_), None) => { let mut buffer = BufferMut::::with_capacity(len); values - .try_map_into(&mut buffer.spare_capacity_mut()[..len], |v| { - ::from(v) - }) + .try_map_into(buffer.spare_capacity_mut(len), |v| ::from(v)) .map_err(|_| overflow())?; // SAFETY: initialized every lane. unsafe { buffer.set_len(len) }; @@ -568,11 +566,9 @@ where (Mask::Values(m), None) => { let mut buffer = BufferMut::::with_capacity(len); values - .try_map_masked_into( - m.bit_buffer(), - &mut buffer.spare_capacity_mut()[..len], - |v| ::from(v), - ) + .try_map_masked_into(m.bit_buffer(), buffer.spare_capacity_mut(len), |v| { + ::from(v) + }) .map_err(|_| overflow())?; // SAFETY: initialized every lane. unsafe { buffer.set_len(len) }; diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 35e49831acc..425d2ef3907 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -78,7 +78,7 @@ fn select_values( let len = true_values.len(); let mut out = BufferMut::::with_capacity(len); { - let out_slice = out.spare_capacity_mut(); + let out_slice = out.spare_capacity_mut(len); let mask_bits = mask .values() diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 5990ecee371..132aa631af4 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -265,7 +265,7 @@ impl VarBinBuilder { self.data.reserve(capacity); let data_len = self.data.len(); - let written = decode(self.data.spare_capacity_mut())?; + let written = decode(self.data.spare_capacity_mut(capacity))?; vortex_ensure!( written == num_bytes, "Decoded {written} bytes, expected {num_bytes}" @@ -478,7 +478,7 @@ impl VarBinBuilder { // Writing into the spare capacity keeps the output cursor in a register: `push` rewrites // the buffer length on every value, which the optimizer cannot hoist out of the loop. - let spare = &mut self.offsets.spare_capacity_mut()[..count]; + let spare = self.offsets.spare_capacity_mut(count); let mut end_offsets = end_offsets; let mut previous = 0usize; for slot in spare.iter_mut() { @@ -531,7 +531,7 @@ impl VarBinBuilder { // Disjoint field borrows: the offsets spare capacity stays valid while the byte buffer // grows, since the two are separate allocations. let Self { offsets, data, .. } = self; - let spare = &mut offsets.spare_capacity_mut()[..count]; + let spare = offsets.spare_capacity_mut(count); match validity.bit_buffer() { AllOr::All => { diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 59689d9116c..5dd0ba8255a 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -630,7 +630,7 @@ where } let mut new_data = ByteBufferMut::with_capacity(output_bytes); - let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let spare = new_data.spare_capacity_mut(output_bytes); let mut cursor = 0usize; for start in starts { let start = start.as_(); @@ -722,7 +722,7 @@ where ); let mut new_data = ByteBufferMut::with_capacity(output_bytes); - let spare = &mut new_data.spare_capacity_mut()[..output_bytes]; + let spare = new_data.spare_capacity_mut(output_bytes); let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); diff --git a/vortex-array/src/arrays/varbinview/build_views.rs b/vortex-array/src/arrays/varbinview/build_views.rs index 8a6a5a0530d..eda80004bf5 100644 --- a/vortex-array/src/arrays/varbinview/build_views.rs +++ b/vortex-array/src/arrays/varbinview/build_views.rs @@ -152,7 +152,7 @@ fn extend_views_single_buffer( // loop-invariant, so it reloads and rewrites the output cursor through the stack each // iteration. Writing into the spare slice keeps the cursor in a register and the length is // set once after the loop. - let spare = &mut views.spare_capacity_mut()[..count]; + let spare = views.spare_capacity_mut(count); for (i, slot) in spare.iter_mut().enumerate() { let len = len_at(i); let value = &data[offset..offset + len]; diff --git a/vortex-array/src/arrays/varbinview/compute/take.rs b/vortex-array/src/arrays/varbinview/compute/take.rs index c10b16cd419..2c662f950a3 100644 --- a/vortex-array/src/arrays/varbinview/compute/take.rs +++ b/vortex-array/src/arrays/varbinview/compute/take.rs @@ -175,7 +175,7 @@ where ); let mut views = BufferMut::::with_capacity(output_len); - let spare = &mut views.spare_capacity_mut()[..output_len]; + let spare = views.spare_capacity_mut(output_len); let mut cursor = 0usize; for &start in starts { let start = start.as_(); @@ -213,7 +213,7 @@ where L: UnsignedPType, { let mut views = BufferMut::::with_capacity(output_len); - let spare = &mut views.spare_capacity_mut()[..output_len]; + let spare = views.spare_capacity_mut(output_len); let mut cursor = 0usize; for (&start, &length) in starts.iter().zip_eq(lengths) { let start = start.as_(); diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index aca2db36286..1501dd304d9 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -245,7 +245,7 @@ impl UninitRange<'_, T> { #[inline] pub fn set_value(&mut self, index: usize, value: T) { assert!(index < self.len, "index out of bounds"); - let spare = self.builder.values.spare_capacity_mut(); + let spare = self.builder.values.spare_capacity_mut(self.len); spare[index] = MaybeUninit::new(value); } @@ -304,10 +304,10 @@ impl UninitRange<'_, T> { // SAFETY: &[T] and &[MaybeUninit] have the same layout. let uninit_src: &[MaybeUninit] = unsafe { std::mem::transmute(src) }; - // Note: spare_capacity_mut() returns the spare capacity starting from the current length, + // Note: spare_capacity_mut returns the spare capacity starting from the current length, // so we just use local_offset directly. - let dst = - &mut self.builder.values.spare_capacity_mut()[local_offset..local_offset + src.len()]; + let dst = &mut self.builder.values.spare_capacity_mut(self.len) + [local_offset..local_offset + src.len()]; dst.copy_from_slice(uninit_src); } @@ -332,7 +332,7 @@ impl UninitRange<'_, T> { len, self.len ); - &mut self.builder.values.spare_capacity_mut()[offset..offset + len] + &mut self.builder.values.spare_capacity_mut(self.len)[offset..offset + len] } /// Finish building this range, marking it as initialized and advancing the length of the diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index 1901f0260c0..65f21430b58 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -48,7 +48,7 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; + let out = values.spare_capacity_mut(len); match valid_bits { None => source.try_map_into(out, apply)?, Some(valid_bits) => source.try_map_masked_into(valid_bits, out, apply)?, diff --git a/vortex-buffer/src/bit/ops.rs b/vortex-buffer/src/bit/ops.rs index ec922b54bc2..5779ae576ea 100644 --- a/vortex-buffer/src/bit/ops.rs +++ b/vortex-buffer/src/bit/ops.rs @@ -114,7 +114,7 @@ pub(super) fn bitwise_unary_op_copy u64>(buffer: &BitBuffer, op let src = buffer.inner().as_slice(); let mut bytes = ByteBufferMut::with_capacity(src.len()); map_u64_words( - OutOfPlaceBitWordTarget::new(src, bytes.spare_capacity_mut()), + OutOfPlaceBitWordTarget::new(src, bytes.spare_capacity_mut(src.len())), op, ); // SAFETY: `map_u64_words` initializes every byte in `0..src.len()` for diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 3b48889fa8d..7dfafd70cbc 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -351,7 +351,7 @@ impl Buffer { let allocator = buf.allocator().clone(); let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf - .spare_capacity_mut() + .spare_capacity_mut(len) .iter_mut() .zip(buf) .for_each(|(out, in_)| { diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 27c3dad3532..28b18bfa3ae 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -538,20 +538,25 @@ impl BufferMut { self.capacity = logical_size / size_of::(); } - /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. - /// Has identical semantics to [`Vec::spare_capacity_mut`]. + /// Returns the first `len` elements of the buffer's spare capacity as a slice of + /// `MaybeUninit`. /// /// The returned slice can be used to fill the buffer with data (e.g. by /// reading from a file) before marking the data as initialized using the /// [`set_len`] method. /// - /// Note that the returned slice may be larger than the capacity requested at - /// construction, since the underlying allocation can be rounded up (e.g. to - /// satisfy alignment requirements). + /// `len` is the number of additional elements, not the buffer's final length. + /// Pass `self.capacity() - self.len()` to access the full spare capacity, as with + /// [`Vec::spare_capacity_mut`]. The allocation can be rounded up to satisfy alignment + /// requirements, so the full spare capacity may exceed the capacity requested at construction. /// /// [`set_len`]: BufferMut::set_len /// [`Vec::spare_capacity_mut`]: Vec::spare_capacity_mut /// + /// # Panics + /// + /// Panics if `len` exceeds `self.capacity() - self.len()`. + /// /// # Examples /// /// ``` @@ -561,7 +566,7 @@ impl BufferMut { /// let mut b = BufferMut::::with_capacity(10); /// /// // Fill in the first 3 elements. - /// let uninit = b.spare_capacity_mut(); + /// let uninit = b.spare_capacity_mut(3); /// uninit[0].write(0); /// uninit[1].write(1); /// uninit[2].write(2); @@ -574,10 +579,15 @@ impl BufferMut { /// assert_eq!(b.as_slice(), &[0u64, 1, 2]); /// ``` #[inline] - pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { - // SAFETY: offset + length is within the allocation and points at spare capacity. + pub fn spare_capacity_mut(&mut self, len: usize) -> &mut [MaybeUninit] { + assert!( + len <= self.capacity() - self.length, + "requested spare capacity exceeds available capacity" + ); + // SAFETY: self.length is within the allocation and points at spare capacity. let dst = unsafe { self.as_mut_ptr().add(self.length) }.cast::>(); - unsafe { std::slice::from_raw_parts_mut(dst, self.capacity() - self.length) } + // SAFETY: the check above ensures that all `len` elements are within spare capacity. + unsafe { std::slice::from_raw_parts_mut(dst, len) } } /// Sets the length of the buffer. @@ -848,7 +858,7 @@ impl BufferMut { let unwritten = self.capacity() - self.len(); // We store `begin` in the case that the lower bound hint is incorrect. - let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut(unwritten).as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); // As a first step, we manually iterate the iterator up to the known capacity. @@ -893,7 +903,10 @@ impl BufferMut { .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self + .spare_capacity_mut(self.capacity() - self.len()) + .as_mut_ptr() + .cast(); let mut dst: *mut T = begin.cast_mut(); iter.for_each(|item| { @@ -982,11 +995,54 @@ impl FromIterator for BufferMut { } #[cfg(test)] -mod test { +mod tests { + use rstest::rstest; + use crate::Alignment; use crate::BufferMut; use crate::buffer_mut; + #[test] + fn spare_capacity_mut_prefix() { + let mut buffer = BufferMut::::with_capacity_aligned(4, Alignment::new(64)); + buffer.push(10); + assert!(buffer.spare_capacity_mut(0).is_empty()); + let slots = buffer.spare_capacity_mut(2); + assert_eq!(slots.len(), 2); + slots[0].write(20); + slots[1].write(30); + assert_eq!(buffer.len(), 1); + // SAFETY: the existing element and both new elements are initialized. + unsafe { buffer.set_len(3) }; + assert_eq!(buffer.as_slice(), &[10, 20, 30]); + let spare = buffer.capacity() - buffer.len(); + assert_eq!(buffer.spare_capacity_mut(spare).len(), spare); + } + + #[test] + fn spare_capacity_mut_without_spare_capacity() { + let mut buffer = BufferMut::::with_capacity(0); + assert!(buffer.spare_capacity_mut(0).is_empty()); + let mut buffer = BufferMut::::with_capacity(1); + buffer.push_n(10, buffer.capacity()); + assert!(buffer.spare_capacity_mut(0).is_empty()); + } + + #[rstest] + #[case(false)] + #[case(true)] + #[should_panic(expected = "requested spare capacity exceeds available capacity")] + fn spare_capacity_mut_exceeds_spare_capacity(#[case] overflowing: bool) { + let mut buffer = BufferMut::::with_capacity(4); + buffer.push(10); + let len = if overflowing { + usize::MAX + } else { + buffer.capacity() + }; + let _ = buffer.spare_capacity_mut(len); + } + #[test] fn capacity() { let mut n = 57; diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index 869f5c4dc95..57995a75f55 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -405,7 +405,7 @@ impl DeviceBuffer for CudaDeviceBuffer { // `cuMemcpyDtoHAsync_v2` fully initializes the memory. unsafe { sys::cuMemcpyDtoHAsync_v2( - host_buffer.spare_capacity_mut().as_mut_ptr().cast(), + host_buffer.spare_capacity_mut(len).as_mut_ptr().cast(), src_ptr, len, stream.cu_stream(), diff --git a/vortex-jni/src/io/read_at.rs b/vortex-jni/src/io/read_at.rs index 2c45fc347cd..f20f3664d20 100644 --- a/vortex-jni/src/io/read_at.rs +++ b/vortex-jni/src/io/read_at.rs @@ -154,7 +154,7 @@ impl VortexReadAt for JavaReadable { // the buffer after it returns. let dst = unsafe { env.new_direct_byte_buffer( - buffer.spare_capacity_mut().as_mut_ptr().cast(), + buffer.spare_capacity_mut(length).as_mut_ptr().cast(), length, )? };