mz_repr/
datum_vec.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A re-useable vector of `Datum` with varying lifetimes.
11//!
12//! This type is meant to allow us to recycle an underlying allocation with
13//! a specific lifetime, under the condition that the vector is emptied before
14//! this happens (to prevent leaking of invalid references).
15//!
16//! It uses `ore::vec::repurpose_allocation` to accomplish this, which contains
17//! unsafe code.
18
19use crate::{Datum, RowRef};
20
21/// A re-useable vector of `Datum` with no particular lifetime.
22#[derive(Debug, Default, Clone)]
23pub struct DatumVec {
24    outer: Vec<Datum<'static>>,
25}
26
27impl DatumVec {
28    /// Allocate a new instance.
29    pub fn new() -> Self {
30        Self { outer: Vec::new() }
31    }
32    /// Borrow an instance with a specific lifetime.
33    ///
34    /// When the result is dropped, its allocation will be returned to `self`.
35    pub fn borrow<'a>(&'a mut self) -> DatumVecBorrow<'a> {
36        let inner = std::mem::take(&mut self.outer);
37        DatumVecBorrow {
38            outer: &mut self.outer,
39            inner,
40        }
41    }
42
43    /// Borrow an instance with a specific lifetime, and pre-populate with a `Row`.
44    pub fn borrow_with<'a>(&'a mut self, row: &'a RowRef) -> DatumVecBorrow<'a> {
45        let mut borrow = self.borrow();
46        borrow.extend(row.iter());
47        borrow
48    }
49
50    /// Borrow an instance with a specific lifetime, and pre-populate with a `Row` with up to
51    /// `limit` elements. If `limit` is greater than the number of elements in `row`, the borrow
52    /// will contain all elements of `row`.
53    pub fn borrow_with_limit<'a>(
54        &'a mut self,
55        row: &'a RowRef,
56        limit: usize,
57    ) -> DatumVecBorrow<'a> {
58        let mut borrow = self.borrow();
59        borrow.extend(row.iter().take(limit));
60        borrow
61    }
62}
63
64/// A borrowed allocation of `Datum` with a specific lifetime.
65///
66/// When an instance is dropped, its allocation is returned to the vector from
67/// which it was extracted.
68#[derive(Debug)]
69pub struct DatumVecBorrow<'outer> {
70    outer: &'outer mut Vec<Datum<'static>>,
71    inner: Vec<Datum<'outer>>,
72}
73
74impl<'outer> Drop for DatumVecBorrow<'outer> {
75    fn drop(&mut self) {
76        *self.outer = mz_ore::vec::repurpose_allocation(std::mem::take(&mut self.inner));
77    }
78}
79
80impl<'outer> std::ops::Deref for DatumVecBorrow<'outer> {
81    type Target = Vec<Datum<'outer>>;
82    fn deref(&self) -> &Self::Target {
83        &self.inner
84    }
85}
86
87impl<'outer> std::ops::DerefMut for DatumVecBorrow<'outer> {
88    fn deref_mut(&mut self) -> &mut Self::Target {
89        &mut self.inner
90    }
91}
92
93#[cfg(test)]
94mod test {
95    use super::*;
96    use crate::Row;
97
98    #[mz_ore::test]
99    fn miri_test_datum_vec() {
100        let mut d = DatumVec::new();
101
102        assert_eq!(d.borrow().len(), 0);
103
104        let r = Row::pack_slice(&[Datum::String("first"), Datum::Dummy]);
105
106        {
107            let borrow = d.borrow_with(&r);
108            assert_eq!(borrow.len(), 2);
109            assert_eq!(borrow[0], Datum::String("first"));
110        }
111
112        {
113            // different lifetime, so that rust is happy with the reference lifetimes
114            let r2 = Row::pack_slice(&[Datum::String("second")]);
115            let mut borrow = d.borrow();
116            borrow.extend(&r);
117            borrow.extend(&r2);
118            assert_eq!(borrow.len(), 3);
119            assert_eq!(borrow[2], Datum::String("second"));
120        }
121    }
122}