Add VersatileBuffer, a reusable allocation lent as a Vec (#3166)

`VersatileBuffer<A>` (tantivy-common) keeps an allocation and lends it
as an empty `Vec<T>`, for any `T` with the size and alignment of `A`.
The allocation goes back to the buffer when the guard is dropped,
including on panic. If the guard is leaked, the buffer starts again
from an empty `Vec`.

This makes it possible to reuse a buffer of values borrowing data with
a short lifetime, without self-referential structs.

The jitexpr predicate uses it for its argument slots. It replaces the
`ClearOnDrop` guard, which transmuted a `Vec<VariableValue<'static>>`
to a shorter lifetime.

Co-authored-by: Paul Masurel <paul.masurel@datadoghq.com>

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,
 };

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<T>` for any `T` with the same size
+/// and alignment as `A`.
+///
+/// Example: `VersatileBuffer<u64>` can lend `Vec<u64>`, `Vec<i64>`, `Vec<f64>`, ...
+///
+/// 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<A> {
+    buffer: Vec<A>,
+}
+
+impl<A> Default for VersatileBuffer<A> {
+    #[inline(always)]
+    fn default() -> Self {
+        VersatileBuffer { buffer: Vec::new() }
+    }
+}
+
+impl<A> VersatileBuffer<A> {
+    #[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<T>`.
+    ///
+    /// 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<T>(&mut self) -> ClearOnDrop<'_, T, A> {
+        const {
+            assert!(std::mem::size_of::<T>() == std::mem::size_of::<A>());
+            assert!(std::mem::align_of::<T>() == std::mem::align_of::<A>());
+        }
+        let mut storage = ManuallyDrop::new(std::mem::take(&mut self.buffer));
+        let capacity = storage.capacity();
+        let ptr = storage.as_mut_ptr().cast::<T>();
+        // SAFETY: `ptr` and `capacity` come from a `Vec<A>` we own. `T` and `A` have the same
+        // size and alignment (checked above), so the allocation is valid for a `Vec<T>` with the
+        // same capacity. Length is 0, so nothing is read.
+        let values: Vec<T> = unsafe { Vec::from_raw_parts(ptr, 0, capacity) };
+        ClearOnDrop {
+            home: &mut self.buffer,
+            values,
+        }
+    }
+}
+
+/// An empty `Vec<T>` 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<A>,
+    values: Vec<T>,
+}
+
+impl<T: Debug, A> Debug for ClearOnDrop<'_, T, A> {
+    #[inline(always)]
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        self.values.fmt(f)
+    }
+}
+
+impl<T, A> 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::<A>();
+        // 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<A> = 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<T, A> Deref for ClearOnDrop<'_, T, A> {
+    type Target = Vec<T>;
+
+    #[inline(always)]
+    fn deref(&self) -> &Self::Target {
+        &self.values
+    }
+}
+
+impl<T, A> 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<u64> = VersatileBuffer::default();
+
+        {
+            let mut arr: ClearOnDrop<u64, u64> = my_buff.borrow();
+            arr.push(1u64);
+            arr.push(2u64);
+            assert_eq!(&arr[..], &[1u64, 2u64]);
+        }
+        {
+            let arr: ClearOnDrop<i64, u64> = my_buff.borrow();
+            assert!(arr.is_empty());
+        }
+    }
+
+    #[test]
+    fn test_versatile_buffer_keeps_capacity() {
+        let mut buffer: VersatileBuffer<u64> = VersatileBuffer::with_capacity(16);
+        {
+            let mut values = buffer.borrow::<u64>();
+            assert!(values.capacity() >= 16);
+            values.extend(0..100u64);
+        }
+        let values = buffer.borrow::<f64>();
+        assert!(values.is_empty());
+        assert!(values.capacity() >= 100);
+    }
+
+    #[test]
+    fn test_versatile_buffer_forgotten_guard() {
+        let mut buffer: VersatileBuffer<u64> = VersatileBuffer::new();
+        let mut values = buffer.borrow::<u64>();
+        values.extend([1u64, 2u64, 3u64]);
+        std::mem::forget(values);
+        let values = buffer.borrow::<f64>();
+        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<u32>);
+
+    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<usize> = VersatileBuffer::new();
+        {
+            let mut values = buffer.borrow::<DropCounter>();
+            for _ in 0..5 {
+                values.push(DropCounter(&drop_count));
+            }
+            assert_eq!(drop_count.get(), 0);
+        }
+        assert_eq!(drop_count.get(), 5);
+    }
+}

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<Option<DynamicColumn>>,
     // One reusable buffer per string column.
     string_inputs: Vec<String>,
-    // 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<VariableValue<'static>>,
-}
-
-/// A wrapper to make sure the variable value buffer is cleared even if the evaluation
-/// panicked.
-struct ClearOnDrop<'a>(&'a mut Vec<VariableValue<'a>>);
-
-impl<'a> ClearOnDrop<'a> {
-    fn wrap(input_values: &'a mut Vec<VariableValue<'static>>) -> 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<VariableValue<'_>> =
-            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<VariableValue>` for each
+    // evaluation.
+    input_values: VersatileBuffer<VariableValue<'static>>,
 }
 
 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::<VariableValue>();
 
         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<bool> = unsafe { self.compiled.call(inputs_vec.0).as_bool() };
+        let eval_result: Option<bool> = unsafe { self.compiled.call(&input_values).as_bool() };
 
         eval_result == Some(true)
     }