diff --git a/common/src/lib.rs b/common/src/lib.rs index 3f47aa8d6..7e7e89151 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -13,6 +13,7 @@ mod group_by; pub mod json_path_writer; mod key_tracking; mod serialize; +mod versatile_buffer; mod vint; mod writer; pub use bitset::*; @@ -23,6 +24,7 @@ pub use json_path_writer::JsonPathWriter; pub use key_tracking::{KeyTracking, WithoutKeys}; pub use ownedbytes::{OwnedBytes, StableDeref}; pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize}; +pub use versatile_buffer::VersatileBuffer; pub use vint::{ VInt, VIntU128, read_u32_vint, read_u32_vint_no_advance, serialize_vint_u32, write_u32_vint, }; diff --git a/common/src/versatile_buffer.rs b/common/src/versatile_buffer.rs new file mode 100644 index 000000000..45bb4061b --- /dev/null +++ b/common/src/versatile_buffer.rs @@ -0,0 +1,188 @@ +use std::fmt::Debug; +use std::mem::ManuallyDrop; +use std::ops::{Deref, DerefMut}; + +/// A reusable allocation that can be lent as an empty `Vec` for any `T` with the same size +/// and alignment as `A`. +/// +/// Example: `VersatileBuffer` can lend `Vec`, `Vec`, `Vec`, ... +/// +/// Hidden contract: `buffer.len()` is always 0. No `A` value is ever built: `A` only sets the +/// size and alignment of the allocation. Only the allocation is reused. +pub struct VersatileBuffer { + buffer: Vec, +} + +impl Default for VersatileBuffer { + #[inline(always)] + fn default() -> Self { + VersatileBuffer { buffer: Vec::new() } + } +} + +impl VersatileBuffer { + #[inline(always)] + pub fn new() -> Self { + Self::default() + } + + #[inline(always)] + pub fn with_capacity(capacity: usize) -> Self { + VersatileBuffer { + buffer: Vec::with_capacity(capacity), + } + } + + /// Lends the allocation as an empty `Vec`. + /// + /// The allocation goes back to `self` when the returned guard is dropped. + /// If the guard is leaked (e.g. `mem::forget`), the allocation is leaked too and `self` + /// starts again from an empty `Vec`. + #[inline(always)] + pub fn borrow(&mut self) -> ClearOnDrop<'_, T, A> { + const { + assert!(std::mem::size_of::() == std::mem::size_of::()); + assert!(std::mem::align_of::() == std::mem::align_of::()); + } + let mut storage = ManuallyDrop::new(std::mem::take(&mut self.buffer)); + let capacity = storage.capacity(); + let ptr = storage.as_mut_ptr().cast::(); + // SAFETY: `ptr` and `capacity` come from a `Vec` we own. `T` and `A` have the same + // size and alignment (checked above), so the allocation is valid for a `Vec` with the + // same capacity. Length is 0, so nothing is read. + let values: Vec = unsafe { Vec::from_raw_parts(ptr, 0, capacity) }; + ClearOnDrop { + home: &mut self.buffer, + values, + } + } +} + +/// An empty `Vec` lent by a `VersatileBuffer`. +/// +/// On drop, the values are dropped and the allocation goes back to the `VersatileBuffer`. +/// This also happens if the code using the vec panics. +pub struct ClearOnDrop<'a, T, A> { + home: &'a mut Vec, + values: Vec, +} + +impl Debug for ClearOnDrop<'_, T, A> { + #[inline(always)] + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + self.values.fmt(f) + } +} + +impl Drop for ClearOnDrop<'_, T, A> { + #[inline(always)] + fn drop(&mut self) { + self.values.clear(); + let mut values = ManuallyDrop::new(std::mem::take(&mut self.values)); + let capacity = values.capacity(); + let ptr = values.as_mut_ptr().cast::(); + // SAFETY: same layout argument as in `VersatileBuffer::borrow`. `values` may have been + // reallocated by `push`, but always with `T`'s layout, which equals `A`'s. + let buffer: Vec = unsafe { Vec::from_raw_parts(ptr, 0, capacity) }; + // `home` still holds the empty `Vec` left by `mem::take` in `borrow`. Forgetting it + // instead of dropping it removes a dead dealloc branch, which the compiler cannot prove + // dead on its own. + std::mem::forget(std::mem::replace(self.home, buffer)); + } +} + +impl Deref for ClearOnDrop<'_, T, A> { + type Target = Vec; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + &self.values + } +} + +impl DerefMut for ClearOnDrop<'_, T, A> { + #[inline(always)] + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.values + } +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + + use crate::versatile_buffer::{ClearOnDrop, VersatileBuffer}; + + #[test] + fn test_versatile_borrowable_buffer() { + let mut my_buff: VersatileBuffer = VersatileBuffer::default(); + + { + let mut arr: ClearOnDrop = my_buff.borrow(); + arr.push(1u64); + arr.push(2u64); + assert_eq!(&arr[..], &[1u64, 2u64]); + } + { + let arr: ClearOnDrop = my_buff.borrow(); + assert!(arr.is_empty()); + } + } + + #[test] + fn test_versatile_buffer_keeps_capacity() { + let mut buffer: VersatileBuffer = VersatileBuffer::with_capacity(16); + { + let mut values = buffer.borrow::(); + assert!(values.capacity() >= 16); + values.extend(0..100u64); + } + let values = buffer.borrow::(); + assert!(values.is_empty()); + assert!(values.capacity() >= 100); + } + + #[test] + fn test_versatile_buffer_forgotten_guard() { + let mut buffer: VersatileBuffer = VersatileBuffer::new(); + let mut values = buffer.borrow::(); + values.extend([1u64, 2u64, 3u64]); + std::mem::forget(values); + let values = buffer.borrow::(); + assert!(values.is_empty()); + } + + #[test] + fn test_versatile_buffer_16_bytes() { + let mut buffer: VersatileBuffer<(u64, u64)> = VersatileBuffer::new(); + { + let mut values = buffer.borrow::<(u64, u32, u32)>(); + values.push((1, 2, 3)); + assert_eq!(&values[..], &[(1, 2, 3)]); + } + let values = buffer.borrow::<[u64; 2]>(); + assert!(values.is_empty()); + } + + struct DropCounter<'a>(&'a Cell); + + impl Drop for DropCounter<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } + } + + #[test] + fn test_versatile_buffer_drops_values() { + let drop_count = Cell::new(0u32); + let mut buffer: VersatileBuffer = VersatileBuffer::new(); + { + let mut values = buffer.borrow::(); + for _ in 0..5 { + values.push(DropCounter(&drop_count)); + } + assert_eq!(drop_count.get(), 0); + } + assert_eq!(drop_count.get(), 5); + } +} diff --git a/src/query/doc_predicate_query/jitexpr_predicate.rs b/src/query/doc_predicate_query/jitexpr_predicate.rs index 9698ac3fd..20e62cb19 100644 --- a/src/query/doc_predicate_query/jitexpr_predicate.rs +++ b/src/query/doc_predicate_query/jitexpr_predicate.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::io; use columnar::{ColumnIndex, ColumnType, DynamicColumn, StrColumn}; +use common::VersatileBuffer; use jitexpr::ast::{ infer_types_with_target, required_presence_for_true, InferredTypeSet, TypeError, UntypedExpr, VariablePresenceCondition, @@ -178,7 +179,7 @@ impl DocPredicate for JitExprPredicate { compiled: compiled_fn.context(), columns_opt, string_inputs: vec![String::new(); num_string_inputs], - input_values: Vec::with_capacity(num_inputs), + input_values: VersatileBuffer::with_capacity(num_inputs), }; Ok(ConstOrVariableSegmentPredicate::Variable { predicate, @@ -278,52 +279,27 @@ pub struct JitExprEvalState { columns_opt: Vec>, // One reusable buffer per string column. string_inputs: Vec, - // Reusable argument slots. - // - // Hidden contract: this vector is always empty between evaluations, so the - // `'static` lifetime is a placeholder for an unused element type rather - // than a claim about any stored string. Only its capacity carries over, - // which is what makes restoring the `'static` type after an evaluation - // sound. `eval` is responsible for upholding this on every return path. - input_values: Vec>, -} - -/// A wrapper to make sure the variable value buffer is cleared even if the evaluation -/// panicked. -struct ClearOnDrop<'a>(&'a mut Vec>); - -impl<'a> ClearOnDrop<'a> { - fn wrap(input_values: &'a mut Vec>) -> Self { - debug_assert!(input_values.is_empty()); - // Input_values is just a buffer we share to avoid allocations - let lower_lifetime_input_values: &mut Vec> = - unsafe { std::mem::transmute(input_values) }; - ClearOnDrop(lower_lifetime_input_values) - } -} - -impl<'a> Drop for ClearOnDrop<'a> { - fn drop(&mut self) { - self.0.clear(); - } + // Reusable allocation for the argument slots, lent as a `Vec` for each + // evaluation. + input_values: VersatileBuffer>, } impl SegmentDocPredicate for JitExprEvalState { fn eval(&mut self, doc_id: DocId) -> bool { - // Input_values is just a buffer we share to avoid allocations - let inputs_vec = ClearOnDrop::wrap(&mut self.input_values); + // The allocation goes back to `self.input_values` when `input_values` is dropped. + let mut input_values = self.input_values.borrow::(); fill_input_values( &self.columns_opt, &mut self.string_inputs, - inputs_vec.0, + &mut input_values, doc_id, ); // SAFETY: Columns follow compiled.inputs() and their types were checked // during setup. Each slot uses the matching union arm. String buffers // remain borrowed, and cannot be mutated, until this call finishes. - let eval_result: Option = unsafe { self.compiled.call(inputs_vec.0).as_bool() }; + let eval_result: Option = unsafe { self.compiled.call(&input_values).as_bool() }; eval_result == Some(true) }