Skip to main content

mz_repr/
row.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//! In-memory `Tag`-based encoding of a tuple of `Datum`s.
11//!
12//! See `doc/developer/row-encoding.md` for the size limits this encoding and
13//! its datum types impose.
14
15use std::borrow::Borrow;
16use std::cell::{Cell, RefCell};
17use std::cmp::Ordering;
18use std::convert::{TryFrom, TryInto};
19use std::fmt::{self, Debug};
20use std::hash::{Hash, Hasher};
21use std::marker::PhantomData;
22use std::mem::{size_of, transmute};
23use std::ops::Deref;
24use std::str;
25
26use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, Timelike, Utc};
27use compact_bytes::CompactBytes;
28use mz_ore::cast::{CastFrom, ReinterpretCast};
29use mz_ore::soft_assert_no_log;
30use mz_ore::vec::Vector;
31use mz_persist_types::Codec64;
32use num_enum::{IntoPrimitive, TryFromPrimitive};
33use ordered_float::OrderedFloat;
34#[cfg(any(test, feature = "proptest"))]
35use proptest::prelude::*;
36#[cfg(any(test, feature = "proptest"))]
37use proptest::strategy::{BoxedStrategy, Strategy};
38use serde::{Deserialize, Serialize};
39use uuid::Uuid;
40
41use crate::adt::array::{
42    Array, ArrayDimension, ArrayDimensions, InvalidArrayError, MAX_ARRAY_DIMENSIONS,
43};
44use crate::adt::date::Date;
45use crate::adt::interval::Interval;
46use crate::adt::mz_acl_item::{AclItem, MzAclItem};
47use crate::adt::numeric;
48use crate::adt::numeric::Numeric;
49use crate::adt::range::{
50    self, InvalidRangeError, Range, RangeBound, RangeInner, RangeLowerBound, RangeUpperBound,
51};
52use crate::adt::timestamp::CheckedTimestamp;
53#[cfg(any(test, feature = "proptest"))]
54use crate::scalar::arb_datum;
55use crate::scalar::{DatumKind, SqlScalarType};
56use crate::{Datum, RelationDesc, Timestamp};
57
58pub(crate) mod encode;
59pub mod iter;
60
61include!(concat!(env!("OUT_DIR"), "/mz_repr.row.rs"));
62
63/// A packed representation for `Datum`s.
64///
65/// `Datum` is easy to work with but very space inefficient. A `Datum::Int32(42)`
66/// is laid out in memory like this:
67///
68///   tag: 3
69///   padding: 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
70///   data: 0 0 0 42
71///   padding: 0 0 0 0 0 0 0 0 0 0 0 0
72///
73/// For a total of 32 bytes! The second set of padding is needed in case we were
74/// to write a 16-byte datum into this location. The first set of padding is
75/// needed to align that hypothetical decimal to a 16 bytes boundary.
76///
77/// A `Row` stores zero or more `Datum`s without any padding. We avoid the need
78/// for the first set of padding by only providing access to the `Datum`s via
79/// calls to `ptr::read_unaligned`, which on modern x86 is barely penalized. We
80/// avoid the need for the second set of padding by not providing mutable access
81/// to the `Datum`. Instead, `Row` is append-only.
82///
83/// A `Row` can be built from a collection of `Datum`s using `Row::pack`, but it
84/// is more efficient to use `Row::pack_slice` so that a right-sized allocation
85/// can be created. If that is not possible, consider using the row buffer
86/// pattern: allocate one row, pack into it, and then call [`Row::clone`] to
87/// receive a copy of that row, leaving behind the original allocation to pack
88/// future rows.
89///
90/// Creating a row via [`Row::pack_slice`]:
91///
92/// ```
93/// # use mz_repr::{Row, Datum};
94/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
95/// assert_eq!(row.unpack(), vec![Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)])
96/// ```
97///
98/// `Row`s can be unpacked by iterating over them:
99///
100/// ```
101/// # use mz_repr::{Row, Datum};
102/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
103/// assert_eq!(row.iter().nth(1).unwrap(), Datum::Int32(1));
104/// ```
105///
106/// If you want random access to the `Datum`s in a `Row`, use `Row::unpack` to create a `Vec<Datum>`
107/// ```
108/// # use mz_repr::{Row, Datum};
109/// let row = Row::pack_slice(&[Datum::Int32(0), Datum::Int32(1), Datum::Int32(2)]);
110/// let datums = row.unpack();
111/// assert_eq!(datums[1], Datum::Int32(1));
112/// ```
113///
114/// # Performance
115///
116/// Rows are dynamically sized, but up to a fixed size their data is stored in-line.
117/// It is best to re-use a `Row` across multiple `Row` creation calls, as this
118/// avoids the allocations involved in `Row::new()`.
119#[derive(Default, Eq, PartialEq, Serialize, Deserialize)]
120pub struct Row {
121    data: CompactBytes,
122}
123
124impl Row {
125    const SIZE: usize = CompactBytes::MAX_INLINE;
126
127    /// A variant of `Row::from_proto` that allows for reuse of internal allocs
128    /// and validates the decoding against a provided [`RelationDesc`].
129    pub fn decode_from_proto(
130        &mut self,
131        proto: &ProtoRow,
132        desc: &RelationDesc,
133    ) -> Result<(), String> {
134        let mut packer = self.packer();
135        for (col_idx, _, _) in desc.iter_all() {
136            let d = match proto.datums.get(col_idx.to_raw()) {
137                Some(x) => x,
138                None => {
139                    packer.push(Datum::Null);
140                    continue;
141                }
142            };
143            packer.try_push_proto(d)?;
144        }
145
146        Ok(())
147    }
148
149    /// Allocate an empty `Row` with a pre-allocated capacity.
150    #[inline]
151    pub fn with_capacity(cap: usize) -> Self {
152        Self {
153            data: CompactBytes::with_capacity(cap),
154        }
155    }
156
157    /// Create an empty `Row`.
158    #[inline]
159    pub const fn empty() -> Self {
160        Self {
161            data: CompactBytes::empty(),
162        }
163    }
164
165    /// Creates a new row from supplied bytes.
166    ///
167    /// # Safety
168    ///
169    /// This method relies on `data` being an appropriate row encoding, and can
170    /// result in unsafety if this is not the case.
171    pub unsafe fn from_bytes_unchecked(data: &[u8]) -> Self {
172        Row {
173            data: CompactBytes::new(data),
174        }
175    }
176
177    /// Constructs a [`RowPacker`] that will pack datums into this row's
178    /// allocation.
179    ///
180    /// This method clears the existing contents of the row, but retains the
181    /// allocation.
182    pub fn packer(&mut self) -> RowPacker<'_> {
183        self.clear();
184        RowPacker { row: self }
185    }
186
187    /// Take some `Datum`s and pack them into a `Row`.
188    ///
189    /// This method builds a `Row` by repeatedly increasing the backing
190    /// allocation. If the contents of the iterator are known ahead of
191    /// time, consider [`Row::with_capacity`] to right-size the allocation
192    /// first, and then [`RowPacker::extend`] to populate it with `Datum`s.
193    /// This avoids the repeated allocation resizing and copying.
194    pub fn pack<'a, I, D>(iter: I) -> Row
195    where
196        I: IntoIterator<Item = D>,
197        D: Borrow<Datum<'a>>,
198    {
199        let mut row = Row::default();
200        row.packer().extend(iter);
201        row
202    }
203
204    /// Use `self` to pack `iter`, and then clone the result.
205    ///
206    /// This is a convenience method meant to reduce boilerplate around row
207    /// formation.
208    pub fn pack_using<'a, I, D>(&mut self, iter: I) -> Row
209    where
210        I: IntoIterator<Item = D>,
211        D: Borrow<Datum<'a>>,
212    {
213        self.packer().extend(iter);
214        self.clone()
215    }
216
217    /// Like [`Row::pack`], but the provided iterator is allowed to produce an
218    /// error, in which case the packing operation is aborted and the error
219    /// returned.
220    pub fn try_pack<'a, I, D, E>(iter: I) -> Result<Row, E>
221    where
222        I: IntoIterator<Item = Result<D, E>>,
223        D: Borrow<Datum<'a>>,
224    {
225        let mut row = Row::default();
226        row.packer().try_extend(iter)?;
227        Ok(row)
228    }
229
230    /// Pack a slice of `Datum`s into a `Row`.
231    ///
232    /// This method has the advantage over `pack` that it can determine the required
233    /// allocation before packing the elements, ensuring only one allocation and no
234    /// redundant copies required.
235    pub fn pack_slice<'a>(slice: &[Datum<'a>]) -> Row {
236        // Pre-allocate the needed number of bytes.
237        let mut row = Row::with_capacity(datums_size(slice.iter()));
238        row.packer().extend(slice.iter());
239        row
240    }
241
242    /// Returns the total amount of bytes used by this row.
243    pub fn byte_len(&self) -> usize {
244        let heap_size = if self.data.spilled() {
245            self.data.len()
246        } else {
247            0
248        };
249        let inline_size = std::mem::size_of::<Self>();
250        inline_size.saturating_add(heap_size)
251    }
252
253    /// The length of the encoded row in bytes. Does not include the size of the `Row` struct itself.
254    pub fn data_len(&self) -> usize {
255        self.data.len()
256    }
257
258    /// Returns the total capacity in bytes used by this row.
259    pub fn byte_capacity(&self) -> usize {
260        self.data.capacity()
261    }
262
263    /// Extracts a Row slice containing the entire [`Row`].
264    #[inline]
265    pub fn as_row_ref(&self) -> &RowRef {
266        // SAFETY: `Row` contains valid row data, by construction.
267        unsafe { RowRef::from_slice(self.data.as_slice()) }
268    }
269
270    /// Clear the contents of the [`Row`], leaving any allocation in place.
271    #[inline]
272    fn clear(&mut self) {
273        self.data.clear();
274    }
275}
276
277impl Borrow<RowRef> for Row {
278    #[inline]
279    fn borrow(&self) -> &RowRef {
280        self.as_row_ref()
281    }
282}
283
284impl AsRef<RowRef> for Row {
285    #[inline]
286    fn as_ref(&self) -> &RowRef {
287        self.as_row_ref()
288    }
289}
290
291impl Deref for Row {
292    type Target = RowRef;
293
294    #[inline]
295    fn deref(&self) -> &Self::Target {
296        self.as_row_ref()
297    }
298}
299
300// Nothing depends on Row being exactly 24, we just want to add visibility to the size.
301static_assertions::const_assert_eq!(std::mem::size_of::<Row>(), 24);
302
303impl Clone for Row {
304    fn clone(&self) -> Self {
305        Row {
306            data: self.data.clone(),
307        }
308    }
309
310    fn clone_from(&mut self, source: &Self) {
311        self.data.clone_from(&source.data);
312    }
313}
314
315// Row's `Hash` implementation defers to `RowRef` to ensure they hash equivalently.
316impl std::hash::Hash for Row {
317    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
318        self.as_row_ref().hash(state)
319    }
320}
321
322#[cfg(any(test, feature = "proptest"))]
323impl Arbitrary for Row {
324    type Parameters = prop::collection::SizeRange;
325    type Strategy = BoxedStrategy<Row>;
326
327    fn arbitrary_with(size: Self::Parameters) -> Self::Strategy {
328        prop::collection::vec(arb_datum(true), size)
329            .prop_map(|items| {
330                let mut row = Row::default();
331                let mut packer = row.packer();
332                for item in items.iter() {
333                    let datum: Datum<'_> = item.into();
334                    packer.push(datum);
335                }
336                row
337            })
338            .boxed()
339    }
340}
341
342impl PartialOrd for Row {
343    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
344        Some(self.cmp(other))
345    }
346}
347
348impl Ord for Row {
349    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
350        self.as_ref().cmp(other.as_ref())
351    }
352}
353
354#[allow(missing_debug_implementations)]
355mod columnation {
356    use columnation::{Columnation, Region};
357    use mz_ore::region::LgAllocRegion;
358
359    use crate::Row;
360
361    /// Region allocation for `Row` data.
362    ///
363    /// Content bytes are stored in stable contiguous memory locations,
364    /// and then a `Row` referencing them is falsified.
365    pub struct RowStack {
366        region: LgAllocRegion<u8>,
367    }
368
369    impl RowStack {
370        const LIMIT: usize = 2 << 20;
371    }
372
373    // Implement `Default` manually to specify a region allocation limit.
374    impl Default for RowStack {
375        fn default() -> Self {
376            Self {
377                // Limit the region size to 2MiB.
378                region: LgAllocRegion::with_limit(Self::LIMIT),
379            }
380        }
381    }
382
383    impl Columnation for Row {
384        type InnerRegion = RowStack;
385    }
386
387    impl Region for RowStack {
388        type Item = Row;
389        #[inline]
390        fn clear(&mut self) {
391            self.region.clear();
392        }
393        #[inline(always)]
394        unsafe fn copy(&mut self, item: &Row) -> Row {
395            if item.data.spilled() {
396                let bytes = self.region.copy_slice(&item.data[..]);
397                Row {
398                    data: compact_bytes::CompactBytes::from_raw_parts(
399                        bytes.as_mut_ptr(),
400                        item.data.len(),
401                        item.data.capacity(),
402                    ),
403                }
404            } else {
405                item.clone()
406            }
407        }
408
409        fn reserve_items<'a, I>(&mut self, items: I)
410        where
411            Self: 'a,
412            I: Iterator<Item = &'a Self::Item> + Clone,
413        {
414            let size = items
415                .filter(|row| row.data.spilled())
416                .map(|row| row.data.len())
417                .sum();
418            let size = std::cmp::min(size, Self::LIMIT);
419            self.region.reserve(size);
420        }
421
422        fn reserve_regions<'a, I>(&mut self, regions: I)
423        where
424            Self: 'a,
425            I: Iterator<Item = &'a Self> + Clone,
426        {
427            let size = regions.map(|r| r.region.len()).sum();
428            let size = std::cmp::min(size, Self::LIMIT);
429            self.region.reserve(size);
430        }
431
432        fn heap_size(&self, callback: impl FnMut(usize, usize)) {
433            self.region.heap_size(callback)
434        }
435    }
436}
437
438mod columnar {
439    use columnar::common::PushIndexAs;
440    use columnar::{
441        AsBytes, Borrow, Clear, Columnar, Container, FromBytes, Index, IndexAs, Len, Push,
442    };
443    use mz_ore::cast::CastFrom;
444    use std::ops::Range;
445
446    use crate::{Row, RowRef};
447
448    #[derive(
449        Copy,
450        Clone,
451        Debug,
452        Default,
453        PartialEq,
454        serde::Serialize,
455        serde::Deserialize
456    )]
457    pub struct Rows<BC = Vec<u64>, VC = Vec<u8>> {
458        /// Bounds container; provides indexed access to offsets.
459        bounds: BC,
460        /// Values container; provides slice access to bytes.
461        values: VC,
462    }
463
464    impl Columnar for Row {
465        #[inline(always)]
466        fn copy_from(&mut self, other: columnar::Ref<'_, Self>) {
467            self.clear();
468            self.data.extend_from_slice(other.data());
469        }
470        #[inline(always)]
471        fn into_owned(other: columnar::Ref<'_, Self>) -> Self {
472            other.to_owned()
473        }
474        type Container = Rows;
475        #[inline(always)]
476        fn reborrow<'b, 'a: 'b>(thing: columnar::Ref<'a, Self>) -> columnar::Ref<'b, Self>
477        where
478            Self: 'a,
479        {
480            thing
481        }
482    }
483
484    impl<BC: PushIndexAs<u64>> Borrow for Rows<BC, Vec<u8>> {
485        type Ref<'a> = &'a RowRef;
486        type Borrowed<'a>
487            = Rows<BC::Borrowed<'a>, &'a [u8]>
488        where
489            Self: 'a;
490        #[inline(always)]
491        fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
492            Rows {
493                bounds: self.bounds.borrow(),
494                values: self.values.borrow(),
495            }
496        }
497        #[inline(always)]
498        fn reborrow<'c, 'a: 'c>(item: Self::Borrowed<'a>) -> Self::Borrowed<'c>
499        where
500            Self: 'a,
501        {
502            Rows {
503                bounds: BC::reborrow(item.bounds),
504                values: item.values,
505            }
506        }
507
508        fn reborrow_ref<'b, 'a: 'b>(item: Self::Ref<'a>) -> Self::Ref<'b>
509        where
510            Self: 'a,
511        {
512            item
513        }
514    }
515
516    impl<BC: PushIndexAs<u64>> Container for Rows<BC, Vec<u8>> {
517        fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: Range<usize>) {
518            if !range.is_empty() {
519                // Imported bounds will be relative to this starting offset.
520                let values_len: u64 = self.values.len().try_into().expect("must fit");
521
522                // Push all bytes that we can, all at once.
523                let other_lower = if range.start == 0 {
524                    0
525                } else {
526                    other.bounds.index_as(range.start - 1)
527                };
528                let other_upper = other.bounds.index_as(range.end - 1);
529                self.values.extend_from_self(
530                    other.values,
531                    usize::try_from(other_lower).expect("must fit")
532                        ..usize::try_from(other_upper).expect("must fit"),
533                );
534
535                // Each bound needs to be shifted by `values_len - other_lower`.
536                if values_len == other_lower {
537                    self.bounds.extend_from_self(other.bounds, range);
538                } else {
539                    for index in range {
540                        let shifted = other.bounds.index_as(index) - other_lower + values_len;
541                        self.bounds.push(&shifted)
542                    }
543                }
544            }
545        }
546        fn reserve_for<'a, I>(&mut self, selves: I)
547        where
548            Self: 'a,
549            I: Iterator<Item = Self::Borrowed<'a>> + Clone,
550        {
551            self.bounds.reserve_for(selves.clone().map(|r| r.bounds));
552            self.values.reserve_for(selves.map(|r| r.values));
553        }
554    }
555
556    impl<'a, BC: AsBytes<'a>, VC: AsBytes<'a>> AsBytes<'a> for Rows<BC, VC> {
557        const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
558        #[inline(always)]
559        fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
560            debug_assert!(index < Self::SLICE_COUNT);
561            if index < BC::SLICE_COUNT {
562                self.bounds.get_byte_slice(index)
563            } else {
564                self.values.get_byte_slice(index - BC::SLICE_COUNT)
565            }
566        }
567    }
568    impl<'a, BC: FromBytes<'a>, VC: FromBytes<'a>> FromBytes<'a> for Rows<BC, VC> {
569        const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
570        #[inline(always)]
571        fn from_bytes(bytes: &mut impl Iterator<Item = &'a [u8]>) -> Self {
572            Self {
573                bounds: FromBytes::from_bytes(bytes),
574                values: FromBytes::from_bytes(bytes),
575            }
576        }
577    }
578
579    impl<BC: Len, VC> Len for Rows<BC, VC> {
580        #[inline(always)]
581        fn len(&self) -> usize {
582            self.bounds.len()
583        }
584    }
585
586    impl<'a, BC: Len + IndexAs<u64>> Index for Rows<BC, &'a [u8]> {
587        type Ref = &'a RowRef;
588        #[inline(always)]
589        fn get(&self, index: usize) -> Self::Ref {
590            let lower = if index == 0 {
591                0
592            } else {
593                self.bounds.index_as(index - 1)
594            };
595            let upper = self.bounds.index_as(index);
596            let lower = usize::cast_from(lower);
597            let upper = usize::cast_from(upper);
598            // SAFETY: self.values contains only valid row data, and self.metadata delimits only ranges
599            // that correspond to the original rows.
600            unsafe { RowRef::from_slice(&self.values[lower..upper]) }
601        }
602    }
603    impl<'a, BC: Len + IndexAs<u64>> Index for &'a Rows<BC, Vec<u8>> {
604        type Ref = &'a RowRef;
605        #[inline(always)]
606        fn get(&self, index: usize) -> Self::Ref {
607            let lower = if index == 0 {
608                0
609            } else {
610                self.bounds.index_as(index - 1)
611            };
612            let upper = self.bounds.index_as(index);
613            let lower = usize::cast_from(lower);
614            let upper = usize::cast_from(upper);
615            // SAFETY: self.values contains only valid row data, and self.metadata delimits only ranges
616            // that correspond to the original rows.
617            unsafe { RowRef::from_slice(&self.values[lower..upper]) }
618        }
619    }
620
621    impl<BC: Push<u64>> Push<&Row> for Rows<BC> {
622        #[inline(always)]
623        fn push(&mut self, item: &Row) {
624            self.values.extend_from_slice(item.data.as_slice());
625            self.bounds.push(u64::cast_from(self.values.len()));
626        }
627    }
628    impl<BC: for<'a> Push<&'a u64>> Push<&RowRef> for Rows<BC> {
629        #[inline(always)]
630        fn push(&mut self, item: &RowRef) {
631            self.values.extend_from_slice(item.data());
632            self.bounds.push(&u64::cast_from(self.values.len()));
633        }
634    }
635    impl<BC: Clear, VC: Clear> Clear for Rows<BC, VC> {
636        #[inline(always)]
637        fn clear(&mut self) {
638            self.bounds.clear();
639            self.values.clear();
640        }
641    }
642}
643
644/// A contiguous slice of bytes that are row data.
645///
646/// A [`RowRef`] is to [`Row`] as [`prim@str`] is to [`String`].
647#[derive(PartialEq, Eq, Hash)]
648#[repr(transparent)]
649pub struct RowRef([u8]);
650
651impl RowRef {
652    /// Create a [`RowRef`] from a slice of data.
653    ///
654    /// # Safety
655    ///
656    /// We do not check that the provided slice is valid [`Row`] data; the caller is required to
657    /// ensure this.
658    pub unsafe fn from_slice(row: &[u8]) -> &RowRef {
659        #[allow(clippy::as_conversions)]
660        let ptr = row as *const [u8] as *const RowRef;
661        // SAFETY: We know `ptr` is non-null and aligned because it came from a &[u8].
662        unsafe { &*ptr }
663    }
664
665    /// Unpack `self` into a `Vec<Datum>` for efficient random access.
666    pub fn unpack(&self) -> Vec<Datum<'_>> {
667        // It's usually cheaper to unpack twice to figure out the right length than it is to grow the vec as we go
668        let len = self.iter().count();
669        let mut vec = Vec::with_capacity(len);
670        vec.extend(self.iter());
671        vec
672    }
673
674    /// Return the first [`Datum`] in `self`
675    ///
676    /// Panics if the [`RowRef`] is empty.
677    pub fn unpack_first(&self) -> Datum<'_> {
678        self.iter().next().unwrap()
679    }
680
681    /// Iterate the [`Datum`] elements of the [`RowRef`].
682    pub fn iter(&self) -> DatumListIter<'_> {
683        DatumListIter { data: &self.0 }
684    }
685
686    /// Return the byte length of this [`RowRef`].
687    pub fn byte_len(&self) -> usize {
688        self.0.len()
689    }
690
691    /// For debugging only.
692    pub fn data(&self) -> &[u8] {
693        &self.0
694    }
695
696    /// True iff there is no data in this [`RowRef`].
697    pub fn is_empty(&self) -> bool {
698        self.0.is_empty()
699    }
700}
701
702impl ToOwned for RowRef {
703    type Owned = Row;
704
705    fn to_owned(&self) -> Self::Owned {
706        // SAFETY: RowRef has the invariant that the wrapped data must be a valid Row encoding.
707        unsafe { Row::from_bytes_unchecked(&self.0) }
708    }
709}
710
711impl<'a> IntoIterator for &'a RowRef {
712    type Item = Datum<'a>;
713    type IntoIter = DatumListIter<'a>;
714
715    fn into_iter(self) -> DatumListIter<'a> {
716        DatumListIter { data: &self.0 }
717    }
718}
719
720/// These implementations order first by length, and then by slice contents.
721/// This allows many comparisons to complete without dereferencing memory.
722/// Warning: These order by the u8 array representation, and NOT by Datum::cmp.
723impl PartialOrd for RowRef {
724    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
725        Some(self.cmp(other))
726    }
727}
728
729impl Ord for RowRef {
730    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
731        match self.0.len().cmp(&other.0.len()) {
732            std::cmp::Ordering::Less => std::cmp::Ordering::Less,
733            std::cmp::Ordering::Greater => std::cmp::Ordering::Greater,
734            std::cmp::Ordering::Equal => self.0.cmp(&other.0),
735        }
736    }
737}
738
739impl fmt::Debug for RowRef {
740    /// Debug representation using the internal datums
741    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
742        f.write_str("RowRef{")?;
743        f.debug_list().entries(&*self).finish()?;
744        f.write_str("}")
745    }
746}
747
748/// Packs datums into a [`Row`].
749///
750/// Creating a `RowPacker` via [`Row::packer`] starts a packing operation on the
751/// row. A packing operation always starts from scratch: the existing contents
752/// of the underlying row are cleared.
753///
754/// To complete a packing operation, drop the `RowPacker`.
755#[derive(Debug)]
756pub struct RowPacker<'a> {
757    row: &'a mut Row,
758}
759
760/// Infallible conversion from a [`Datum`] to a typed value.
761///
762/// Used by [`DatumList::typed_iter`] to yield elements as `T` rather than
763/// raw `Datum`s. At runtime, `T` is always `Datum<'a>`, so the conversion
764/// is identity.
765///
766/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
767/// behind the generic type parameter and type erasure.
768///
769/// This trait is sealed and cannot be implemented outside of this crate.
770pub trait FromDatum<'a>:
771    Sized + PartialEq + std::borrow::Borrow<Datum<'a>> + sealed::Sealed
772{
773    fn from_datum(datum: Datum<'a>) -> Self;
774}
775
776mod sealed {
777    use crate::Datum;
778
779    pub trait Sealed {}
780    impl<'a> Sealed for Datum<'a> {}
781}
782
783impl<'a> FromDatum<'a> for Datum<'a> {
784    #[inline]
785    fn from_datum(datum: Datum<'a>) -> Self {
786        datum
787    }
788}
789
790#[derive(Debug, Clone)]
791pub struct DatumListIter<'a> {
792    data: &'a [u8],
793}
794
795#[derive(Debug, Clone)]
796pub struct DatumListTypedIter<'a, T> {
797    inner: DatumListIter<'a>,
798    _phantom: PhantomData<fn() -> T>,
799}
800
801#[derive(Debug, Clone)]
802pub struct DatumDictIter<'a> {
803    data: &'a [u8],
804    prev_key: Option<&'a str>,
805}
806
807#[derive(Debug, Clone)]
808pub struct DatumDictTypedIter<'a, T> {
809    inner: DatumDictIter<'a>,
810    _phantom: PhantomData<fn() -> T>,
811}
812
813/// `RowArena` is used to hold on to temporary `Row`s for functions like `eval` that need to create complex `Datum`s but don't have a `Row` to put them in yet.
814#[derive(Debug)]
815pub struct RowArena {
816    // A stack of byte regions, used as a bump allocator. Bytes handed to
817    // `push_bytes` are *copied* into the active (last) region and a reference
818    // into that region is returned.
819    //
820    // The invariant that keeps returned references valid for the arena's
821    // lifetime is that a region is never reallocated once it holds data: when
822    // the active region lacks spare capacity for a push we allocate a *new*,
823    // larger region rather than growing the current one (which would move its
824    // bytes and dangle outstanding references). The outer `Vec` may itself
825    // reallocate as regions are added, but that only moves the `Vec<u8>`
826    // headers, not the heap buffers they own, so references remain valid.
827    //
828    // `clear` retains only the largest region (emptied) to right-size the arena
829    // for reuse; reusing one region across `clear` cycles makes a steady-state
830    // workload (e.g. decoding rows one at a time) allocation-free.
831    inner: RefCell<Vec<Vec<u8>>>,
832    // A single recycled scratch buffer backing `RowArena::writer`. A writer takes ownership of this
833    // buffer (or allocates a fresh one if absent), builds into it, and on drop returns it here for
834    // the next writer to reuse — so building values incrementally does not allocate per use once the
835    // buffer reaches its high-water mark. Holding `Option` (rather than the buffer directly) means
836    // `writer` borrows this cell only transiently, to take and return the buffer, never across the
837    // writer's lifetime. That keeps nested writers sound: a writer obtained while another is live
838    // finds the slot empty and allocates its own buffer instead of double-borrowing.
839    scratch: RefCell<Option<Vec<u8>>>,
840}
841
842// DatumList and DatumDict defined here rather than near Datum because we need private access to the unsafe data field
843
844/// A sequence of Datums
845///
846/// The type parameter `T` represents the element type of the list. It is a
847/// phantom parameter that carries no runtime data — the actual elements are
848/// stored as serialized bytes and `T` is not enforced at runtime. It is up
849/// to the caller to ensure `T` matches the actual element type. The default
850/// `T = Datum<'a>` means existing code that writes `DatumList<'a>` continues
851/// to work unchanged.
852///
853/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
854/// behind the generic type parameter.
855pub struct DatumList<'a, T = Datum<'a>> {
856    /// Points at the serialized datums
857    data: &'a [u8],
858    _phantom: PhantomData<fn() -> T>,
859}
860
861impl<'a, T> DatumList<'a, T> {
862    /// Private constructor. All `DatumList` values should be created through
863    /// this function to keep the `PhantomData` bookkeeping in one place.
864    pub(crate) fn new(data: &'a [u8]) -> Self {
865        DatumList {
866            data,
867            _phantom: PhantomData,
868        }
869    }
870}
871
872impl<'a, T> Clone for DatumList<'a, T> {
873    fn clone(&self) -> Self {
874        *self
875    }
876}
877
878impl<'a, T> Copy for DatumList<'a, T> {}
879
880impl<'a, T> Debug for DatumList<'a, T> {
881    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
882        f.debug_list().entries(self.iter()).finish()
883    }
884}
885
886impl<'a, T> PartialEq for DatumList<'a, T> {
887    #[inline(always)]
888    fn eq(&self, other: &DatumList<'a, T>) -> bool {
889        self.iter().eq(other.iter())
890    }
891}
892
893impl<'a, T> Eq for DatumList<'a, T> {}
894
895impl<'a, T> Hash for DatumList<'a, T> {
896    #[inline(always)]
897    fn hash<H: Hasher>(&self, state: &mut H) {
898        for d in self.iter() {
899            d.hash(state);
900        }
901    }
902}
903
904impl<T> Ord for DatumList<'_, T> {
905    #[inline(always)]
906    fn cmp(&self, other: &DatumList<'_, T>) -> Ordering {
907        // Grow the stack: lists can be arbitrarily deeply nested (e.g. jsonb).
908        mz_ore::stack::maybe_grow(|| self.iter().cmp(other.iter()))
909    }
910}
911
912impl<T> PartialOrd for DatumList<'_, T> {
913    #[inline(always)]
914    fn partial_cmp(&self, other: &DatumList<'_, T>) -> Option<Ordering> {
915        Some(self.cmp(other))
916    }
917}
918
919/// A mapping from string keys to Datums
920///
921/// The type parameter `T` represents the value type of the map. It is a
922/// phantom parameter — the actual values are stored as serialized bytes and
923/// `T` is not enforced at runtime. It is up to the caller to ensure `T`
924/// matches the actual value type. The default `T = Datum<'a>` means existing
925/// code that writes `DatumMap<'a>` continues to work unchanged.
926///
927/// See `doc/developer/design/20260311_sqlfunc_generic.md` for the design
928/// behind the generic type parameter.
929pub struct DatumMap<'a, T = Datum<'a>> {
930    /// Points at the serialized datums, which should be sorted in key order
931    data: &'a [u8],
932    _phantom: PhantomData<fn() -> T>,
933}
934
935impl<'a, T> DatumMap<'a, T> {
936    /// Private constructor. All `DatumMap` values should be created through
937    /// this function to keep the `PhantomData` bookkeeping in one place.
938    pub(crate) fn new(data: &'a [u8]) -> Self {
939        DatumMap {
940            data,
941            _phantom: PhantomData,
942        }
943    }
944}
945
946impl<'a, T> Clone for DatumMap<'a, T> {
947    fn clone(&self) -> Self {
948        *self
949    }
950}
951
952impl<'a, T> Copy for DatumMap<'a, T> {}
953
954impl<'a, T> PartialEq for DatumMap<'a, T> {
955    #[inline(always)]
956    fn eq(&self, other: &DatumMap<'a, T>) -> bool {
957        self.iter().eq(other.iter())
958    }
959}
960
961impl<'a, T> Eq for DatumMap<'a, T> {}
962
963impl<'a, T> Hash for DatumMap<'a, T> {
964    #[inline(always)]
965    fn hash<H: Hasher>(&self, state: &mut H) {
966        for (k, v) in self.iter() {
967            k.hash(state);
968            v.hash(state);
969        }
970    }
971}
972
973impl<'a, T> Ord for DatumMap<'a, T> {
974    #[inline(always)]
975    fn cmp(&self, other: &DatumMap<'a, T>) -> Ordering {
976        // Grow the stack: maps can be arbitrarily deeply nested (e.g. jsonb).
977        mz_ore::stack::maybe_grow(|| self.iter().cmp(other.iter()))
978    }
979}
980
981impl<'a, T> PartialOrd for DatumMap<'a, T> {
982    #[inline(always)]
983    fn partial_cmp(&self, other: &DatumMap<'a, T>) -> Option<Ordering> {
984        Some(self.cmp(other))
985    }
986}
987
988impl<'a> crate::scalar::SqlContainerType for DatumList<'a, Datum<'a>> {
989    fn unwrap_element_type(container: &SqlScalarType) -> &SqlScalarType {
990        container.unwrap_list_element_type()
991    }
992    fn wrap_element_type(element: SqlScalarType) -> SqlScalarType {
993        SqlScalarType::List {
994            element_type: Box::new(element),
995            custom_id: None,
996        }
997    }
998}
999
1000impl<'a> crate::scalar::SqlContainerType for DatumMap<'a, Datum<'a>> {
1001    fn unwrap_element_type(container: &SqlScalarType) -> &SqlScalarType {
1002        container.unwrap_map_value_type()
1003    }
1004    fn wrap_element_type(element: SqlScalarType) -> SqlScalarType {
1005        SqlScalarType::Map {
1006            value_type: Box::new(element),
1007            custom_id: None,
1008        }
1009    }
1010}
1011
1012/// Represents a single `Datum`, appropriate to be nested inside other
1013/// `Datum`s.
1014#[derive(Clone, Copy, Eq, PartialEq, Hash)]
1015pub struct DatumNested<'a> {
1016    val: &'a [u8],
1017}
1018
1019impl<'a> std::fmt::Display for DatumNested<'a> {
1020    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1021        std::fmt::Display::fmt(&self.datum(), f)
1022    }
1023}
1024
1025impl<'a> std::fmt::Debug for DatumNested<'a> {
1026    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1027        f.debug_struct("DatumNested")
1028            .field("val", &self.datum())
1029            .finish()
1030    }
1031}
1032
1033impl<'a> DatumNested<'a> {
1034    // Figure out which bytes `read_datum` returns (e.g. including the tag),
1035    // and then store a reference to those bytes, so we can "replay" this same
1036    // call later on without storing the datum itself.
1037    pub fn extract(data: &mut &'a [u8]) -> DatumNested<'a> {
1038        let prev = *data;
1039        let _ = unsafe { read_datum(data) };
1040        DatumNested {
1041            val: &prev[..(prev.len() - data.len())],
1042        }
1043    }
1044
1045    /// Returns the datum `self` contains.
1046    pub fn datum(&self) -> Datum<'a> {
1047        let mut temp = self.val;
1048        unsafe { read_datum(&mut temp) }
1049    }
1050}
1051
1052impl<'a> Ord for DatumNested<'a> {
1053    fn cmp(&self, other: &Self) -> Ordering {
1054        // Grow the stack: this recurses once per level of nested list/map values.
1055        mz_ore::stack::maybe_grow(|| self.datum().cmp(&other.datum()))
1056    }
1057}
1058
1059impl<'a> PartialOrd for DatumNested<'a> {
1060    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1061        Some(self.cmp(other))
1062    }
1063}
1064
1065// Prefer adding new tags to the end of the enum. Certain behavior, like row ordering and EXPLAIN
1066// PHYSICAL PLAN, rely on the ordering of this enum. Neither of these are breaking changes, but
1067// it's annoying when they change.
1068#[derive(Debug, Clone, Copy, PartialEq, Eq, IntoPrimitive, TryFromPrimitive)]
1069#[repr(u8)]
1070enum Tag {
1071    Null,
1072    False,
1073    True,
1074    Int16,
1075    Int32,
1076    Int64,
1077    UInt8,
1078    UInt32,
1079    Float32,
1080    Float64,
1081    Date,
1082    Time,
1083    Timestamp,
1084    TimestampTz,
1085    Interval,
1086    BytesTiny,
1087    BytesShort,
1088    BytesLong,
1089    BytesHuge,
1090    StringTiny,
1091    StringShort,
1092    StringLong,
1093    StringHuge,
1094    Uuid,
1095    Array,
1096    ListTiny,
1097    ListShort,
1098    ListLong,
1099    ListHuge,
1100    Dict,
1101    JsonNull,
1102    Dummy,
1103    Numeric,
1104    UInt16,
1105    UInt64,
1106    MzTimestamp,
1107    Range,
1108    MzAclItem,
1109    AclItem,
1110    // Everything except leap seconds and times beyond the range of
1111    // i64 nanoseconds. (Note that Materialize does not support leap
1112    // seconds, but this module does).
1113    CheapTimestamp,
1114    // Everything except leap seconds and times beyond the range of
1115    // i64 nanoseconds. (Note that Materialize does not support leap
1116    // seconds, but this module does).
1117    CheapTimestampTz,
1118    // The next several tags are for variable-length signed integer encoding.
1119    // The basic idea is that `NonNegativeIntN_K` is used to encode a datum of type
1120    // IntN whose actual value is positive or zero and fits in K bits, and similarly for
1121    // NegativeIntN_K with negative values.
1122    //
1123    // The order of these tags matters, because we want to be able to choose the
1124    // tag for a given datum quickly, with arithmetic, rather than slowly, with a
1125    // stack of `if` statements.
1126    //
1127    // Separate tags for non-negative and negative numbers are used to avoid having to
1128    // waste one bit in the actual data space to encode the sign.
1129    NonNegativeInt16_0, // i.e., 0
1130    NonNegativeInt16_8,
1131    NonNegativeInt16_16,
1132
1133    NonNegativeInt32_0,
1134    NonNegativeInt32_8,
1135    NonNegativeInt32_16,
1136    NonNegativeInt32_24,
1137    NonNegativeInt32_32,
1138
1139    NonNegativeInt64_0,
1140    NonNegativeInt64_8,
1141    NonNegativeInt64_16,
1142    NonNegativeInt64_24,
1143    NonNegativeInt64_32,
1144    NonNegativeInt64_40,
1145    NonNegativeInt64_48,
1146    NonNegativeInt64_56,
1147    NonNegativeInt64_64,
1148
1149    NegativeInt16_0, // i.e., -1
1150    NegativeInt16_8,
1151    NegativeInt16_16,
1152
1153    NegativeInt32_0,
1154    NegativeInt32_8,
1155    NegativeInt32_16,
1156    NegativeInt32_24,
1157    NegativeInt32_32,
1158
1159    NegativeInt64_0,
1160    NegativeInt64_8,
1161    NegativeInt64_16,
1162    NegativeInt64_24,
1163    NegativeInt64_32,
1164    NegativeInt64_40,
1165    NegativeInt64_48,
1166    NegativeInt64_56,
1167    NegativeInt64_64,
1168
1169    // These are like the ones above, but for unsigned types. The
1170    // situation is slightly simpler as we don't have negatives.
1171    UInt8_0, // i.e., 0
1172    UInt8_8,
1173
1174    UInt16_0,
1175    UInt16_8,
1176    UInt16_16,
1177
1178    UInt32_0,
1179    UInt32_8,
1180    UInt32_16,
1181    UInt32_24,
1182    UInt32_32,
1183
1184    UInt64_0,
1185    UInt64_8,
1186    UInt64_16,
1187    UInt64_24,
1188    UInt64_32,
1189    UInt64_40,
1190    UInt64_48,
1191    UInt64_56,
1192    UInt64_64,
1193}
1194
1195impl Tag {
1196    fn actual_int_length(self) -> Option<usize> {
1197        use Tag::*;
1198        let val = match self {
1199            NonNegativeInt16_0 | NonNegativeInt32_0 | NonNegativeInt64_0 | UInt8_0 | UInt16_0
1200            | UInt32_0 | UInt64_0 => 0,
1201            NonNegativeInt16_8 | NonNegativeInt32_8 | NonNegativeInt64_8 | UInt8_8 | UInt16_8
1202            | UInt32_8 | UInt64_8 => 1,
1203            NonNegativeInt16_16 | NonNegativeInt32_16 | NonNegativeInt64_16 | UInt16_16
1204            | UInt32_16 | UInt64_16 => 2,
1205            NonNegativeInt32_24 | NonNegativeInt64_24 | UInt32_24 | UInt64_24 => 3,
1206            NonNegativeInt32_32 | NonNegativeInt64_32 | UInt32_32 | UInt64_32 => 4,
1207            NonNegativeInt64_40 | UInt64_40 => 5,
1208            NonNegativeInt64_48 | UInt64_48 => 6,
1209            NonNegativeInt64_56 | UInt64_56 => 7,
1210            NonNegativeInt64_64 | UInt64_64 => 8,
1211            NegativeInt16_0 | NegativeInt32_0 | NegativeInt64_0 => 0,
1212            NegativeInt16_8 | NegativeInt32_8 | NegativeInt64_8 => 1,
1213            NegativeInt16_16 | NegativeInt32_16 | NegativeInt64_16 => 2,
1214            NegativeInt32_24 | NegativeInt64_24 => 3,
1215            NegativeInt32_32 | NegativeInt64_32 => 4,
1216            NegativeInt64_40 => 5,
1217            NegativeInt64_48 => 6,
1218            NegativeInt64_56 => 7,
1219            NegativeInt64_64 => 8,
1220
1221            _ => return None,
1222        };
1223        Some(val)
1224    }
1225}
1226
1227// --------------------------------------------------------------------------------
1228// reading data
1229
1230/// Read a byte slice starting at byte `offset`.
1231///
1232/// Updates `offset` to point to the first byte after the end of the read region.
1233fn read_untagged_bytes<'a>(data: &mut &'a [u8]) -> &'a [u8] {
1234    let len = u64::from_le_bytes(read_byte_array(data));
1235    let len = usize::cast_from(len);
1236    let (bytes, next) = data.split_at(len);
1237    *data = next;
1238    bytes
1239}
1240
1241/// Read a data whose length is encoded in the row before its contents.
1242///
1243/// Updates `offset` to point to the first byte after the end of the read region.
1244///
1245/// # Safety
1246///
1247/// This function is safe if the datum's length and contents were previously written by `push_lengthed_bytes`,
1248/// and it was only written with a `String` tag if it was indeed UTF-8.
1249unsafe fn read_lengthed_datum<'a>(data: &mut &'a [u8], tag: Tag) -> Datum<'a> {
1250    let len = match tag {
1251        Tag::BytesTiny | Tag::StringTiny | Tag::ListTiny => usize::from(read_byte(data)),
1252        Tag::BytesShort | Tag::StringShort | Tag::ListShort => {
1253            usize::from(u16::from_le_bytes(read_byte_array(data)))
1254        }
1255        Tag::BytesLong | Tag::StringLong | Tag::ListLong => {
1256            usize::cast_from(u32::from_le_bytes(read_byte_array(data)))
1257        }
1258        Tag::BytesHuge | Tag::StringHuge | Tag::ListHuge => {
1259            usize::cast_from(u64::from_le_bytes(read_byte_array(data)))
1260        }
1261        _ => unreachable!(),
1262    };
1263    let (bytes, next) = data.split_at(len);
1264    *data = next;
1265    match tag {
1266        Tag::BytesTiny | Tag::BytesShort | Tag::BytesLong | Tag::BytesHuge => Datum::Bytes(bytes),
1267        Tag::StringTiny | Tag::StringShort | Tag::StringLong | Tag::StringHuge => {
1268            Datum::String(str::from_utf8_unchecked(bytes))
1269        }
1270        Tag::ListTiny | Tag::ListShort | Tag::ListLong | Tag::ListHuge => {
1271            Datum::List(DatumList::new(bytes))
1272        }
1273        _ => unreachable!(),
1274    }
1275}
1276
1277fn read_byte(data: &mut &[u8]) -> u8 {
1278    let byte = data[0];
1279    *data = &data[1..];
1280    byte
1281}
1282
1283/// Read `length` bytes from `data` at `offset`, updating the
1284/// latter. Extend the resulting buffer to an array of `N` bytes by
1285/// inserting `FILL` in the k most significant bytes, where k = N - length.
1286///
1287/// SAFETY:
1288///   * length <= N
1289///   * offset + length <= data.len()
1290fn read_byte_array_sign_extending<const N: usize, const FILL: u8>(
1291    data: &mut &[u8],
1292    length: usize,
1293) -> [u8; N] {
1294    let mut raw = [FILL; N];
1295    let (prev, next) = data.split_at(length);
1296    (raw[..prev.len()]).copy_from_slice(prev);
1297    *data = next;
1298    raw
1299}
1300/// Read `length` bytes from `data` at `offset`, updating the
1301/// latter. Extend the resulting buffer to a negative `N`-byte
1302/// twos complement integer by filling the remaining bits with 1.
1303///
1304/// SAFETY:
1305///   * length <= N
1306///   * offset + length <= data.len()
1307fn read_byte_array_extending_negative<const N: usize>(data: &mut &[u8], length: usize) -> [u8; N] {
1308    read_byte_array_sign_extending::<N, 255>(data, length)
1309}
1310
1311/// Read `length` bytes from `data` at `offset`, updating the
1312/// latter. Extend the resulting buffer to a positive or zero `N`-byte
1313/// twos complement integer by filling the remaining bits with 0.
1314///
1315/// SAFETY:
1316///   * length <= N
1317///   * offset + length <= data.len()
1318fn read_byte_array_extending_nonnegative<const N: usize>(
1319    data: &mut &[u8],
1320    length: usize,
1321) -> [u8; N] {
1322    read_byte_array_sign_extending::<N, 0>(data, length)
1323}
1324
1325pub(super) fn read_byte_array<const N: usize>(data: &mut &[u8]) -> [u8; N] {
1326    let (prev, next) = data.split_first_chunk().unwrap();
1327    *data = next;
1328    *prev
1329}
1330
1331pub(super) fn read_date(data: &mut &[u8]) -> Date {
1332    let days = i32::from_le_bytes(read_byte_array(data));
1333    Date::from_pg_epoch(days).expect("unexpected date")
1334}
1335
1336pub(super) fn read_naive_date(data: &mut &[u8]) -> NaiveDate {
1337    let year = i32::from_le_bytes(read_byte_array(data));
1338    let ordinal = u32::from_le_bytes(read_byte_array(data));
1339    NaiveDate::from_yo_opt(year, ordinal).unwrap()
1340}
1341
1342pub(super) fn read_time(data: &mut &[u8]) -> NaiveTime {
1343    let secs = u32::from_le_bytes(read_byte_array(data));
1344    let nanos = u32::from_le_bytes(read_byte_array(data));
1345    NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos).unwrap()
1346}
1347
1348/// Read a datum starting at byte `offset`.
1349///
1350/// Updates `offset` to point to the first byte after the end of the read region.
1351///
1352/// # Safety
1353///
1354/// This function is safe if a `Datum` was previously written at this offset by `push_datum`.
1355/// Otherwise it could return invalid values, which is Undefined Behavior.
1356pub unsafe fn read_datum<'a>(data: &mut &'a [u8]) -> Datum<'a> {
1357    let tag = Tag::try_from_primitive(read_byte(data)).expect("unknown row tag");
1358    match tag {
1359        Tag::Null => Datum::Null,
1360        Tag::False => Datum::False,
1361        Tag::True => Datum::True,
1362        Tag::UInt8_0 | Tag::UInt8_8 => {
1363            let i = u8::from_le_bytes(read_byte_array_extending_nonnegative(
1364                data,
1365                tag.actual_int_length()
1366                    .expect("returns a value for variable-length-encoded integer tags"),
1367            ));
1368            Datum::UInt8(i)
1369        }
1370        Tag::Int16 => {
1371            let i = i16::from_le_bytes(read_byte_array(data));
1372            Datum::Int16(i)
1373        }
1374        Tag::NonNegativeInt16_0 | Tag::NonNegativeInt16_16 | Tag::NonNegativeInt16_8 => {
1375            // SAFETY:`tag.actual_int_length()` is <= 16 for these tags,
1376            // and `data` is big enough because it was encoded validly. These assumptions
1377            // are checked in debug asserts.
1378            let i = i16::from_le_bytes(read_byte_array_extending_nonnegative(
1379                data,
1380                tag.actual_int_length()
1381                    .expect("returns a value for variable-length-encoded integer tags"),
1382            ));
1383            Datum::Int16(i)
1384        }
1385        Tag::UInt16_0 | Tag::UInt16_8 | Tag::UInt16_16 => {
1386            let i = u16::from_le_bytes(read_byte_array_extending_nonnegative(
1387                data,
1388                tag.actual_int_length()
1389                    .expect("returns a value for variable-length-encoded integer tags"),
1390            ));
1391            Datum::UInt16(i)
1392        }
1393        Tag::Int32 => {
1394            let i = i32::from_le_bytes(read_byte_array(data));
1395            Datum::Int32(i)
1396        }
1397        Tag::NonNegativeInt32_0
1398        | Tag::NonNegativeInt32_32
1399        | Tag::NonNegativeInt32_8
1400        | Tag::NonNegativeInt32_16
1401        | Tag::NonNegativeInt32_24 => {
1402            // SAFETY:`tag.actual_int_length()` is <= 32 for these tags,
1403            // and `data` is big enough because it was encoded validly. These assumptions
1404            // are checked in debug asserts.
1405            let i = i32::from_le_bytes(read_byte_array_extending_nonnegative(
1406                data,
1407                tag.actual_int_length()
1408                    .expect("returns a value for variable-length-encoded integer tags"),
1409            ));
1410            Datum::Int32(i)
1411        }
1412        Tag::UInt32_0 | Tag::UInt32_8 | Tag::UInt32_16 | Tag::UInt32_24 | Tag::UInt32_32 => {
1413            let i = u32::from_le_bytes(read_byte_array_extending_nonnegative(
1414                data,
1415                tag.actual_int_length()
1416                    .expect("returns a value for variable-length-encoded integer tags"),
1417            ));
1418            Datum::UInt32(i)
1419        }
1420        Tag::Int64 => {
1421            let i = i64::from_le_bytes(read_byte_array(data));
1422            Datum::Int64(i)
1423        }
1424        Tag::NonNegativeInt64_0
1425        | Tag::NonNegativeInt64_64
1426        | Tag::NonNegativeInt64_8
1427        | Tag::NonNegativeInt64_16
1428        | Tag::NonNegativeInt64_24
1429        | Tag::NonNegativeInt64_32
1430        | Tag::NonNegativeInt64_40
1431        | Tag::NonNegativeInt64_48
1432        | Tag::NonNegativeInt64_56 => {
1433            // SAFETY:`tag.actual_int_length()` is <= 64 for these tags,
1434            // and `data` is big enough because it was encoded validly. These assumptions
1435            // are checked in debug asserts.
1436
1437            let i = i64::from_le_bytes(read_byte_array_extending_nonnegative(
1438                data,
1439                tag.actual_int_length()
1440                    .expect("returns a value for variable-length-encoded integer tags"),
1441            ));
1442            Datum::Int64(i)
1443        }
1444        Tag::UInt64_0
1445        | Tag::UInt64_8
1446        | Tag::UInt64_16
1447        | Tag::UInt64_24
1448        | Tag::UInt64_32
1449        | Tag::UInt64_40
1450        | Tag::UInt64_48
1451        | Tag::UInt64_56
1452        | Tag::UInt64_64 => {
1453            let i = u64::from_le_bytes(read_byte_array_extending_nonnegative(
1454                data,
1455                tag.actual_int_length()
1456                    .expect("returns a value for variable-length-encoded integer tags"),
1457            ));
1458            Datum::UInt64(i)
1459        }
1460        Tag::NegativeInt16_0 | Tag::NegativeInt16_16 | Tag::NegativeInt16_8 => {
1461            // SAFETY:`tag.actual_int_length()` is <= 16 for these tags,
1462            // and `data` is big enough because it was encoded validly. These assumptions
1463            // are checked in debug asserts.
1464            let i = i16::from_le_bytes(read_byte_array_extending_negative(
1465                data,
1466                tag.actual_int_length()
1467                    .expect("returns a value for variable-length-encoded integer tags"),
1468            ));
1469            Datum::Int16(i)
1470        }
1471        Tag::NegativeInt32_0
1472        | Tag::NegativeInt32_32
1473        | Tag::NegativeInt32_8
1474        | Tag::NegativeInt32_16
1475        | Tag::NegativeInt32_24 => {
1476            // SAFETY:`tag.actual_int_length()` is <= 32 for these tags,
1477            // and `data` is big enough because it was encoded validly. These assumptions
1478            // are checked in debug asserts.
1479            let i = i32::from_le_bytes(read_byte_array_extending_negative(
1480                data,
1481                tag.actual_int_length()
1482                    .expect("returns a value for variable-length-encoded integer tags"),
1483            ));
1484            Datum::Int32(i)
1485        }
1486        Tag::NegativeInt64_0
1487        | Tag::NegativeInt64_64
1488        | Tag::NegativeInt64_8
1489        | Tag::NegativeInt64_16
1490        | Tag::NegativeInt64_24
1491        | Tag::NegativeInt64_32
1492        | Tag::NegativeInt64_40
1493        | Tag::NegativeInt64_48
1494        | Tag::NegativeInt64_56 => {
1495            // SAFETY:`tag.actual_int_length()` is <= 64 for these tags,
1496            // and `data` is big enough because the row was encoded validly. These assumptions
1497            // are checked in debug asserts.
1498            let i = i64::from_le_bytes(read_byte_array_extending_negative(
1499                data,
1500                tag.actual_int_length()
1501                    .expect("returns a value for variable-length-encoded integer tags"),
1502            ));
1503            Datum::Int64(i)
1504        }
1505
1506        Tag::UInt8 => {
1507            let i = u8::from_le_bytes(read_byte_array(data));
1508            Datum::UInt8(i)
1509        }
1510        Tag::UInt16 => {
1511            let i = u16::from_le_bytes(read_byte_array(data));
1512            Datum::UInt16(i)
1513        }
1514        Tag::UInt32 => {
1515            let i = u32::from_le_bytes(read_byte_array(data));
1516            Datum::UInt32(i)
1517        }
1518        Tag::UInt64 => {
1519            let i = u64::from_le_bytes(read_byte_array(data));
1520            Datum::UInt64(i)
1521        }
1522        Tag::Float32 => {
1523            let f = f32::from_bits(u32::from_le_bytes(read_byte_array(data)));
1524            Datum::Float32(OrderedFloat::from(f))
1525        }
1526        Tag::Float64 => {
1527            let f = f64::from_bits(u64::from_le_bytes(read_byte_array(data)));
1528            Datum::Float64(OrderedFloat::from(f))
1529        }
1530        Tag::Date => Datum::Date(read_date(data)),
1531        Tag::Time => Datum::Time(read_time(data)),
1532        Tag::CheapTimestamp => {
1533            let ts = i64::from_le_bytes(read_byte_array(data));
1534            let secs = ts.div_euclid(1_000_000_000);
1535            let nsecs: u32 = ts.rem_euclid(1_000_000_000).try_into().unwrap();
1536            let ndt = DateTime::from_timestamp(secs, nsecs)
1537                .expect("We only write round-trippable timestamps")
1538                .naive_utc();
1539            Datum::Timestamp(
1540                CheckedTimestamp::from_timestamplike(ndt).expect("unexpected timestamp"),
1541            )
1542        }
1543        Tag::CheapTimestampTz => {
1544            let ts = i64::from_le_bytes(read_byte_array(data));
1545            let secs = ts.div_euclid(1_000_000_000);
1546            let nsecs: u32 = ts.rem_euclid(1_000_000_000).try_into().unwrap();
1547            let dt = DateTime::from_timestamp(secs, nsecs)
1548                .expect("We only write round-trippable timestamps");
1549            Datum::TimestampTz(
1550                CheckedTimestamp::from_timestamplike(dt).expect("unexpected timestamp"),
1551            )
1552        }
1553        Tag::Timestamp => {
1554            let date = read_naive_date(data);
1555            let time = read_time(data);
1556            Datum::Timestamp(
1557                CheckedTimestamp::from_timestamplike(date.and_time(time))
1558                    .expect("unexpected timestamp"),
1559            )
1560        }
1561        Tag::TimestampTz => {
1562            let date = read_naive_date(data);
1563            let time = read_time(data);
1564            Datum::TimestampTz(
1565                CheckedTimestamp::from_timestamplike(DateTime::from_naive_utc_and_offset(
1566                    date.and_time(time),
1567                    Utc,
1568                ))
1569                .expect("unexpected timestamptz"),
1570            )
1571        }
1572        Tag::Interval => {
1573            let months = i32::from_le_bytes(read_byte_array(data));
1574            let days = i32::from_le_bytes(read_byte_array(data));
1575            let micros = i64::from_le_bytes(read_byte_array(data));
1576            Datum::Interval(Interval {
1577                months,
1578                days,
1579                micros,
1580            })
1581        }
1582        Tag::BytesTiny
1583        | Tag::BytesShort
1584        | Tag::BytesLong
1585        | Tag::BytesHuge
1586        | Tag::StringTiny
1587        | Tag::StringShort
1588        | Tag::StringLong
1589        | Tag::StringHuge
1590        | Tag::ListTiny
1591        | Tag::ListShort
1592        | Tag::ListLong
1593        | Tag::ListHuge => read_lengthed_datum(data, tag),
1594        Tag::Uuid => Datum::Uuid(Uuid::from_bytes(read_byte_array(data))),
1595        Tag::Array => {
1596            // See the comment in `Row::push_array` for details on the encoding
1597            // of arrays.
1598            let ndims = read_byte(data);
1599            let dims_size = usize::from(ndims) * size_of::<u64>() * 2;
1600            let (dims, next) = data.split_at(dims_size);
1601            *data = next;
1602            let bytes = read_untagged_bytes(data);
1603            Datum::Array(Array {
1604                dims: ArrayDimensions { data: dims },
1605                elements: DatumList::new(bytes),
1606            })
1607        }
1608        Tag::Dict => {
1609            let bytes = read_untagged_bytes(data);
1610            Datum::Map(DatumMap::new(bytes))
1611        }
1612        Tag::JsonNull => Datum::JsonNull,
1613        Tag::Dummy => Datum::Dummy,
1614        Tag::Numeric => {
1615            let digits = read_byte(data).into();
1616            let exponent = i8::reinterpret_cast(read_byte(data));
1617            let bits = read_byte(data);
1618
1619            let lsu_u16_len = Numeric::digits_to_lsu_elements_len(digits);
1620            let lsu_u8_len = lsu_u16_len * 2;
1621            let (lsu_u8, next) = data.split_at(lsu_u8_len);
1622            *data = next;
1623
1624            // TODO: if we refactor the decimal library to accept the owned
1625            // array as a parameter to `from_raw_parts` below, we could likely
1626            // avoid a copy because it is exactly the value we want
1627            let mut lsu = [0; numeric::NUMERIC_DATUM_WIDTH_USIZE];
1628            for (i, c) in lsu_u8.chunks(2).enumerate() {
1629                lsu[i] = u16::from_le_bytes(c.try_into().unwrap());
1630            }
1631
1632            let d = Numeric::from_raw_parts(digits, exponent.into(), bits, lsu);
1633            Datum::from(d)
1634        }
1635        Tag::MzTimestamp => {
1636            let t = Timestamp::decode(read_byte_array(data));
1637            Datum::MzTimestamp(t)
1638        }
1639        Tag::Range => {
1640            // See notes on `push_range_with` for details about encoding.
1641            let flag_byte = read_byte(data);
1642            let flags = range::InternalFlags::from_bits(flag_byte)
1643                .expect("range flags must be encoded validly");
1644
1645            if flags.contains(range::InternalFlags::EMPTY) {
1646                assert!(
1647                    flags == range::InternalFlags::EMPTY,
1648                    "empty ranges contain only RANGE_EMPTY flag"
1649                );
1650
1651                return Datum::Range(Range { inner: None });
1652            }
1653
1654            let lower_bound = if flags.contains(range::InternalFlags::LB_INFINITE) {
1655                None
1656            } else {
1657                Some(DatumNested::extract(data))
1658            };
1659
1660            let lower = RangeBound {
1661                inclusive: flags.contains(range::InternalFlags::LB_INCLUSIVE),
1662                bound: lower_bound,
1663            };
1664
1665            let upper_bound = if flags.contains(range::InternalFlags::UB_INFINITE) {
1666                None
1667            } else {
1668                Some(DatumNested::extract(data))
1669            };
1670
1671            let upper = RangeBound {
1672                inclusive: flags.contains(range::InternalFlags::UB_INCLUSIVE),
1673                bound: upper_bound,
1674            };
1675
1676            Datum::Range(Range {
1677                inner: Some(RangeInner { lower, upper }),
1678            })
1679        }
1680        Tag::MzAclItem => {
1681            const N: usize = MzAclItem::binary_size();
1682            let mz_acl_item =
1683                MzAclItem::decode_binary(&read_byte_array::<N>(data)).expect("invalid mz_aclitem");
1684            Datum::MzAclItem(mz_acl_item)
1685        }
1686        Tag::AclItem => {
1687            const N: usize = AclItem::binary_size();
1688            let acl_item =
1689                AclItem::decode_binary(&read_byte_array::<N>(data)).expect("invalid aclitem");
1690            Datum::AclItem(acl_item)
1691        }
1692    }
1693}
1694
1695// --------------------------------------------------------------------------------
1696// writing data
1697
1698fn push_untagged_bytes<D>(data: &mut D, bytes: &[u8])
1699where
1700    D: Vector<u8>,
1701{
1702    let len = u64::cast_from(bytes.len());
1703    data.extend_from_slice(&len.to_le_bytes());
1704    data.extend_from_slice(bytes);
1705}
1706
1707fn push_lengthed_bytes<D>(data: &mut D, bytes: &[u8], tag: Tag)
1708where
1709    D: Vector<u8>,
1710{
1711    match tag {
1712        Tag::BytesTiny | Tag::StringTiny | Tag::ListTiny => {
1713            let len = bytes.len().to_le_bytes();
1714            data.push(len[0]);
1715        }
1716        Tag::BytesShort | Tag::StringShort | Tag::ListShort => {
1717            let len = bytes.len().to_le_bytes();
1718            data.extend_from_slice(&len[0..2]);
1719        }
1720        Tag::BytesLong | Tag::StringLong | Tag::ListLong => {
1721            let len = bytes.len().to_le_bytes();
1722            data.extend_from_slice(&len[0..4]);
1723        }
1724        Tag::BytesHuge | Tag::StringHuge | Tag::ListHuge => {
1725            let len = bytes.len().to_le_bytes();
1726            data.extend_from_slice(&len);
1727        }
1728        _ => unreachable!(),
1729    }
1730    data.extend_from_slice(bytes);
1731}
1732
1733pub(super) fn date_to_array(date: Date) -> [u8; size_of::<i32>()] {
1734    i32::to_le_bytes(date.pg_epoch_days())
1735}
1736
1737fn push_date<D>(data: &mut D, date: Date)
1738where
1739    D: Vector<u8>,
1740{
1741    data.extend_from_slice(&date_to_array(date));
1742}
1743
1744pub(super) fn naive_date_to_arrays(
1745    date: NaiveDate,
1746) -> ([u8; size_of::<i32>()], [u8; size_of::<u32>()]) {
1747    (
1748        i32::to_le_bytes(date.year()),
1749        u32::to_le_bytes(date.ordinal()),
1750    )
1751}
1752
1753fn push_naive_date<D>(data: &mut D, date: NaiveDate)
1754where
1755    D: Vector<u8>,
1756{
1757    let (ds1, ds2) = naive_date_to_arrays(date);
1758    data.extend_from_slice(&ds1);
1759    data.extend_from_slice(&ds2);
1760}
1761
1762pub(super) fn time_to_arrays(time: NaiveTime) -> ([u8; size_of::<u32>()], [u8; size_of::<u32>()]) {
1763    (
1764        u32::to_le_bytes(time.num_seconds_from_midnight()),
1765        u32::to_le_bytes(time.nanosecond()),
1766    )
1767}
1768
1769fn push_time<D>(data: &mut D, time: NaiveTime)
1770where
1771    D: Vector<u8>,
1772{
1773    let (ts1, ts2) = time_to_arrays(time);
1774    data.extend_from_slice(&ts1);
1775    data.extend_from_slice(&ts2);
1776}
1777
1778/// Returns an i64 representing a `NaiveDateTime`, if
1779/// said i64 can be round-tripped back to a `NaiveDateTime`.
1780///
1781/// The only exotic NDTs for which this can't happen are those that
1782/// are hundreds of years in the future or past, or those that
1783/// represent a leap second. (Note that Materialize does not support
1784/// leap seconds, but this module does).
1785// This function is inspired by `NaiveDateTime::timestamp_nanos`,
1786// with extra checking.
1787fn checked_timestamp_nanos(dt: NaiveDateTime) -> Option<i64> {
1788    let subsec_nanos = dt.and_utc().timestamp_subsec_nanos();
1789    if subsec_nanos >= 1_000_000_000 {
1790        return None;
1791    }
1792    let as_ns = dt.and_utc().timestamp().checked_mul(1_000_000_000)?;
1793    as_ns.checked_add(i64::from(subsec_nanos))
1794}
1795
1796// This function is extremely hot, so
1797// we just use `as` to avoid the overhead of
1798// `try_into` followed by `unwrap`.
1799// `leading_ones` and `leading_zeros`
1800// can never return values greater than 64, so the conversion is safe.
1801#[inline(always)]
1802#[allow(clippy::as_conversions)]
1803fn min_bytes_signed<T>(i: T) -> u8
1804where
1805    T: Into<i64>,
1806{
1807    let i: i64 = i.into();
1808
1809    // To fit in n bytes, we require that
1810    // everything but the leading sign bits fits in n*8
1811    // bits.
1812    let n_sign_bits = if i.is_negative() {
1813        i.leading_ones() as u8
1814    } else {
1815        i.leading_zeros() as u8
1816    };
1817
1818    (64 - n_sign_bits + 7) / 8
1819}
1820
1821// In principle we could just use `min_bytes_signed`, rather than
1822// having a separate function here, as long as we made that one take
1823// `T: Into<i128>` instead of 64. But LLVM doesn't seem smart enough
1824// to realize that that function is the same as the current version,
1825// and generates worse code.
1826//
1827// Justification for `as` is the same as in `min_bytes_signed`.
1828#[inline(always)]
1829#[allow(clippy::as_conversions)]
1830fn min_bytes_unsigned<T>(i: T) -> u8
1831where
1832    T: Into<u64>,
1833{
1834    let i: u64 = i.into();
1835
1836    let n_sign_bits = i.leading_zeros() as u8;
1837
1838    (64 - n_sign_bits + 7) / 8
1839}
1840
1841const TINY: usize = 1 << 8;
1842const SHORT: usize = 1 << 16;
1843const LONG: usize = 1 << 32;
1844
1845fn push_datum<D>(data: &mut D, datum: Datum)
1846where
1847    D: Vector<u8>,
1848{
1849    match datum {
1850        Datum::Null => data.push(Tag::Null.into()),
1851        Datum::False => data.push(Tag::False.into()),
1852        Datum::True => data.push(Tag::True.into()),
1853        Datum::Int16(i) => {
1854            let mbs = min_bytes_signed(i);
1855            let tag = u8::from(if i.is_negative() {
1856                Tag::NegativeInt16_0
1857            } else {
1858                Tag::NonNegativeInt16_0
1859            }) + mbs;
1860
1861            data.push(tag);
1862            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1863        }
1864        Datum::Int32(i) => {
1865            let mbs = min_bytes_signed(i);
1866            let tag = u8::from(if i.is_negative() {
1867                Tag::NegativeInt32_0
1868            } else {
1869                Tag::NonNegativeInt32_0
1870            }) + mbs;
1871
1872            data.push(tag);
1873            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1874        }
1875        Datum::Int64(i) => {
1876            let mbs = min_bytes_signed(i);
1877            let tag = u8::from(if i.is_negative() {
1878                Tag::NegativeInt64_0
1879            } else {
1880                Tag::NonNegativeInt64_0
1881            }) + mbs;
1882
1883            data.push(tag);
1884            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbs)]);
1885        }
1886        Datum::UInt8(i) => {
1887            let mbu = min_bytes_unsigned(i);
1888            let tag = u8::from(Tag::UInt8_0) + mbu;
1889            data.push(tag);
1890            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1891        }
1892        Datum::UInt16(i) => {
1893            let mbu = min_bytes_unsigned(i);
1894            let tag = u8::from(Tag::UInt16_0) + mbu;
1895            data.push(tag);
1896            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1897        }
1898        Datum::UInt32(i) => {
1899            let mbu = min_bytes_unsigned(i);
1900            let tag = u8::from(Tag::UInt32_0) + mbu;
1901            data.push(tag);
1902            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1903        }
1904        Datum::UInt64(i) => {
1905            let mbu = min_bytes_unsigned(i);
1906            let tag = u8::from(Tag::UInt64_0) + mbu;
1907            data.push(tag);
1908            data.extend_from_slice(&i.to_le_bytes()[0..usize::from(mbu)]);
1909        }
1910        Datum::Float32(f) => {
1911            data.push(Tag::Float32.into());
1912            data.extend_from_slice(&f.to_bits().to_le_bytes());
1913        }
1914        Datum::Float64(f) => {
1915            data.push(Tag::Float64.into());
1916            data.extend_from_slice(&f.to_bits().to_le_bytes());
1917        }
1918        Datum::Date(d) => {
1919            data.push(Tag::Date.into());
1920            push_date(data, d);
1921        }
1922        Datum::Time(t) => {
1923            data.push(Tag::Time.into());
1924            push_time(data, t);
1925        }
1926        Datum::Timestamp(t) => {
1927            let datetime = t.to_naive();
1928            if let Some(nanos) = checked_timestamp_nanos(datetime) {
1929                data.push(Tag::CheapTimestamp.into());
1930                data.extend_from_slice(&nanos.to_le_bytes());
1931            } else {
1932                data.push(Tag::Timestamp.into());
1933                push_naive_date(data, datetime.date());
1934                push_time(data, datetime.time());
1935            }
1936        }
1937        Datum::TimestampTz(t) => {
1938            let datetime = t.to_naive();
1939            if let Some(nanos) = checked_timestamp_nanos(datetime) {
1940                data.push(Tag::CheapTimestampTz.into());
1941                data.extend_from_slice(&nanos.to_le_bytes());
1942            } else {
1943                data.push(Tag::TimestampTz.into());
1944                push_naive_date(data, datetime.date());
1945                push_time(data, datetime.time());
1946            }
1947        }
1948        Datum::Interval(i) => {
1949            data.push(Tag::Interval.into());
1950            data.extend_from_slice(&i.months.to_le_bytes());
1951            data.extend_from_slice(&i.days.to_le_bytes());
1952            data.extend_from_slice(&i.micros.to_le_bytes());
1953        }
1954        Datum::Bytes(bytes) => {
1955            let tag = match bytes.len() {
1956                0..TINY => Tag::BytesTiny,
1957                TINY..SHORT => Tag::BytesShort,
1958                SHORT..LONG => Tag::BytesLong,
1959                _ => Tag::BytesHuge,
1960            };
1961            data.push(tag.into());
1962            push_lengthed_bytes(data, bytes, tag);
1963        }
1964        Datum::String(string) => {
1965            let tag = match string.len() {
1966                0..TINY => Tag::StringTiny,
1967                TINY..SHORT => Tag::StringShort,
1968                SHORT..LONG => Tag::StringLong,
1969                _ => Tag::StringHuge,
1970            };
1971            data.push(tag.into());
1972            push_lengthed_bytes(data, string.as_bytes(), tag);
1973        }
1974        Datum::List(list) => {
1975            let tag = match list.data.len() {
1976                0..TINY => Tag::ListTiny,
1977                TINY..SHORT => Tag::ListShort,
1978                SHORT..LONG => Tag::ListLong,
1979                _ => Tag::ListHuge,
1980            };
1981            data.push(tag.into());
1982            push_lengthed_bytes(data, list.data, tag);
1983        }
1984        Datum::Uuid(u) => {
1985            data.push(Tag::Uuid.into());
1986            data.extend_from_slice(u.as_bytes());
1987        }
1988        Datum::Array(array) => {
1989            // See the comment in `Row::push_array` for details on the encoding
1990            // of arrays.
1991            data.push(Tag::Array.into());
1992            data.push(array.dims.ndims());
1993            data.extend_from_slice(array.dims.data);
1994            push_untagged_bytes(data, array.elements.data);
1995        }
1996        Datum::Map(dict) => {
1997            data.push(Tag::Dict.into());
1998            push_untagged_bytes(data, dict.data);
1999        }
2000        Datum::JsonNull => data.push(Tag::JsonNull.into()),
2001        Datum::MzTimestamp(t) => {
2002            data.push(Tag::MzTimestamp.into());
2003            data.extend_from_slice(&t.encode());
2004        }
2005        Datum::Dummy => data.push(Tag::Dummy.into()),
2006        Datum::Numeric(mut n) => {
2007            // Pseudo-canonical representation of decimal values with
2008            // insignificant zeroes trimmed. This compresses the number further
2009            // than `Numeric::trim` by removing all zeroes, and not only those in
2010            // the fractional component.
2011            numeric::cx_datum().reduce(&mut n.0);
2012            let (digits, exponent, bits, lsu) = n.0.to_raw_parts();
2013            data.push(Tag::Numeric.into());
2014            data.push(u8::try_from(digits).expect("digits to fit within u8; should not exceed 39"));
2015            data.push(
2016                i8::try_from(exponent)
2017                    .expect("exponent to fit within i8; should not exceed +/- 39")
2018                    .to_le_bytes()[0],
2019            );
2020            data.push(bits);
2021
2022            let lsu = &lsu[..Numeric::digits_to_lsu_elements_len(digits)];
2023
2024            // Little endian machines can take the lsu directly from u16 to u8.
2025            if cfg!(target_endian = "little") {
2026                // SAFETY: `lsu` (returned by `coefficient_units()`) is a `&[u16]`, so
2027                // each element can safely be transmuted into two `u8`s.
2028                let (prefix, lsu_bytes, suffix) = unsafe { lsu.align_to::<u8>() };
2029                // The `u8` aligned version of the `lsu` should have twice as many
2030                // elements as we expect for the `u16` version.
2031                soft_assert_no_log!(
2032                    lsu_bytes.len() == Numeric::digits_to_lsu_elements_len(digits) * 2,
2033                    "u8 version of numeric LSU contained the wrong number of elements; expected {}, but got {}",
2034                    Numeric::digits_to_lsu_elements_len(digits) * 2,
2035                    lsu_bytes.len()
2036                );
2037                // There should be no unaligned elements in the prefix or suffix.
2038                soft_assert_no_log!(prefix.is_empty() && suffix.is_empty());
2039                data.extend_from_slice(lsu_bytes);
2040            } else {
2041                for u in lsu {
2042                    data.extend_from_slice(&u.to_le_bytes());
2043                }
2044            }
2045        }
2046        Datum::Range(range) => {
2047            // See notes on `push_range_with` for details about encoding.
2048            data.push(Tag::Range.into());
2049            data.push(range.internal_flag_bits());
2050
2051            if let Some(RangeInner { lower, upper }) = range.inner {
2052                for bound in [lower.bound, upper.bound] {
2053                    if let Some(bound) = bound {
2054                        match bound.datum() {
2055                            Datum::Null => panic!("cannot push Datum::Null into range"),
2056                            d => push_datum::<D>(data, d),
2057                        }
2058                    }
2059                }
2060            }
2061        }
2062        Datum::MzAclItem(mz_acl_item) => {
2063            data.push(Tag::MzAclItem.into());
2064            data.extend_from_slice(&mz_acl_item.encode_binary());
2065        }
2066        Datum::AclItem(acl_item) => {
2067            data.push(Tag::AclItem.into());
2068            data.extend_from_slice(&acl_item.encode_binary());
2069        }
2070    }
2071}
2072
2073/// Return the number of bytes these Datums would use if packed as a Row.
2074pub fn row_size<'a, I>(a: I) -> usize
2075where
2076    I: IntoIterator<Item = Datum<'a>>,
2077{
2078    // Using datums_size instead of a.data().len() here is safer because it will
2079    // return the size of the datums if they were packed into a Row. Although
2080    // a.data().len() happens to give the correct answer (and is faster), data()
2081    // is documented as for debugging only.
2082    let sz = datums_size::<_, _>(a);
2083    let size_of_row = std::mem::size_of::<Row>();
2084    // The Row struct attempts to inline data until it can't fit in the
2085    // preallocated size. Otherwise it spills to heap, and uses the Row to point
2086    // to that.
2087    if sz > Row::SIZE {
2088        sz + size_of_row
2089    } else {
2090        size_of_row
2091    }
2092}
2093
2094/// Number of bytes required by the datum.
2095/// This is used to optimistically pre-allocate buffers for packing rows.
2096pub fn datum_size(datum: &Datum) -> usize {
2097    match datum {
2098        Datum::Null => 1,
2099        Datum::False => 1,
2100        Datum::True => 1,
2101        Datum::Int16(i) => 1 + usize::from(min_bytes_signed(*i)),
2102        Datum::Int32(i) => 1 + usize::from(min_bytes_signed(*i)),
2103        Datum::Int64(i) => 1 + usize::from(min_bytes_signed(*i)),
2104        Datum::UInt8(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2105        Datum::UInt16(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2106        Datum::UInt32(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2107        Datum::UInt64(i) => 1 + usize::from(min_bytes_unsigned(*i)),
2108        Datum::Float32(_) => 1 + size_of::<f32>(),
2109        Datum::Float64(_) => 1 + size_of::<f64>(),
2110        Datum::Date(_) => 1 + size_of::<i32>(),
2111        Datum::Time(_) => 1 + 8,
2112        Datum::Timestamp(t) => {
2113            1 + if checked_timestamp_nanos(t.to_naive()).is_some() {
2114                8
2115            } else {
2116                16
2117            }
2118        }
2119        Datum::TimestampTz(t) => {
2120            1 + if checked_timestamp_nanos(t.naive_utc()).is_some() {
2121                8
2122            } else {
2123                16
2124            }
2125        }
2126        Datum::Interval(_) => 1 + size_of::<i32>() + size_of::<i32>() + size_of::<i64>(),
2127        Datum::Bytes(bytes) => {
2128            // We use a variable length representation of slice length.
2129            let bytes_for_length = match bytes.len() {
2130                0..TINY => 1,
2131                TINY..SHORT => 2,
2132                SHORT..LONG => 4,
2133                _ => 8,
2134            };
2135            1 + bytes_for_length + bytes.len()
2136        }
2137        Datum::String(string) => {
2138            // We use a variable length representation of slice length.
2139            let bytes_for_length = match string.len() {
2140                0..TINY => 1,
2141                TINY..SHORT => 2,
2142                SHORT..LONG => 4,
2143                _ => 8,
2144            };
2145            1 + bytes_for_length + string.len()
2146        }
2147        Datum::Uuid(_) => 1 + size_of::<uuid::Bytes>(),
2148        Datum::Array(array) => {
2149            1 + size_of::<u8>()
2150                + array.dims.data.len()
2151                + size_of::<u64>()
2152                + array.elements.data.len()
2153        }
2154        Datum::List(list) => 1 + size_of::<u64>() + list.data.len(),
2155        Datum::Map(dict) => 1 + size_of::<u64>() + dict.data.len(),
2156        Datum::JsonNull => 1,
2157        Datum::MzTimestamp(_) => 1 + size_of::<Timestamp>(),
2158        Datum::Dummy => 1,
2159        Datum::Numeric(d) => {
2160            let mut d = d.0.clone();
2161            // Values must be reduced to determine appropriate number of
2162            // coefficient units.
2163            numeric::cx_datum().reduce(&mut d);
2164            // 4 = 1 bit each for tag, digits, exponent, bits
2165            4 + (d.coefficient_units().len() * 2)
2166        }
2167        Datum::Range(Range { inner }) => {
2168            // Tag + flags
2169            2 + match inner {
2170                None => 0,
2171                Some(RangeInner { lower, upper }) => [lower.bound, upper.bound]
2172                    .iter()
2173                    .map(|bound| match bound {
2174                        None => 0,
2175                        Some(bound) => bound.val.len(),
2176                    })
2177                    .sum(),
2178            }
2179        }
2180        Datum::MzAclItem(_) => 1 + MzAclItem::binary_size(),
2181        Datum::AclItem(_) => 1 + AclItem::binary_size(),
2182    }
2183}
2184
2185/// Number of bytes required by a sequence of datums.
2186///
2187/// This method can be used to right-size the allocation for a `Row`
2188/// before calling [`RowPacker::extend`].
2189pub fn datums_size<'a, I, D>(iter: I) -> usize
2190where
2191    I: IntoIterator<Item = D>,
2192    D: Borrow<Datum<'a>>,
2193{
2194    iter.into_iter().map(|d| datum_size(d.borrow())).sum()
2195}
2196
2197/// Number of bytes required by a list of datums. This computes the size that would be required if
2198/// the given datums were packed into a list.
2199///
2200/// This is used to optimistically pre-allocate buffers for packing rows.
2201pub fn datum_list_size<'a, I, D>(iter: I) -> usize
2202where
2203    I: IntoIterator<Item = D>,
2204    D: Borrow<Datum<'a>>,
2205{
2206    1 + size_of::<u64>() + datums_size(iter)
2207}
2208
2209impl RowPacker<'_> {
2210    /// Constructs a row packer that will pack additional datums into the
2211    /// provided row.
2212    ///
2213    /// This function is intentionally somewhat inconvenient to call. You
2214    /// usually want to call [`Row::packer`] instead to start packing from
2215    /// scratch.
2216    pub fn for_existing_row(row: &mut Row) -> RowPacker<'_> {
2217        RowPacker { row }
2218    }
2219
2220    /// Extend an existing `Row` with a `Datum`.
2221    #[inline]
2222    pub fn push<'a, D>(&mut self, datum: D)
2223    where
2224        D: Borrow<Datum<'a>>,
2225    {
2226        push_datum(&mut self.row.data, *datum.borrow());
2227    }
2228
2229    /// Extend an existing `Row` with additional `Datum`s.
2230    #[inline]
2231    pub fn extend<'a, I, D>(&mut self, iter: I)
2232    where
2233        I: IntoIterator<Item = D>,
2234        D: Borrow<Datum<'a>>,
2235    {
2236        for datum in iter {
2237            push_datum(&mut self.row.data, *datum.borrow())
2238        }
2239    }
2240
2241    /// Extend an existing `Row` with additional `Datum`s.
2242    ///
2243    /// In the case the iterator produces an error, the pushing of
2244    /// datums in terminated and the error returned. The `Row` will
2245    /// be incomplete, but it will be safe to read datums from it.
2246    #[inline]
2247    pub fn try_extend<'a, I, E, D>(&mut self, iter: I) -> Result<(), E>
2248    where
2249        I: IntoIterator<Item = Result<D, E>>,
2250        D: Borrow<Datum<'a>>,
2251    {
2252        for datum in iter {
2253            push_datum(&mut self.row.data, *datum?.borrow());
2254        }
2255        Ok(())
2256    }
2257
2258    /// Appends the datums of an entire `Row`.
2259    pub fn extend_by_row(&mut self, row: &Row) {
2260        self.row.data.extend_from_slice(row.data.as_slice());
2261    }
2262
2263    /// Appends the datums of an entire `Row`.
2264    pub fn extend_by_row_ref(&mut self, row: &RowRef) {
2265        self.row.data.extend_from_slice(row.data());
2266    }
2267
2268    /// Appends the slice of data representing an entire `Row`. The data is not validated.
2269    ///
2270    /// # Safety
2271    ///
2272    /// The requirements from [`Row::from_bytes_unchecked`] apply here, too:
2273    /// This method relies on `data` being an appropriate row encoding, and can
2274    /// result in unsafety if this is not the case.
2275    #[inline]
2276    pub unsafe fn extend_by_slice_unchecked(&mut self, data: &[u8]) {
2277        self.row.data.extend_from_slice(data)
2278    }
2279
2280    /// Pushes a [`DatumList`] that is built from a closure.
2281    ///
2282    /// The supplied closure will be invoked once with a `Row` that can be used
2283    /// to populate the list. It is valid to call any method on the
2284    /// [`RowPacker`] except for [`RowPacker::clear`], [`RowPacker::truncate`],
2285    /// or [`RowPacker::truncate_datums`].
2286    ///
2287    /// Returns the value returned by the closure, if any.
2288    ///
2289    /// ```
2290    /// # use mz_repr::{Row, Datum};
2291    /// let mut row = Row::default();
2292    /// row.packer().push_list_with(|row| {
2293    ///     row.push(Datum::String("age"));
2294    ///     row.push(Datum::Int64(42));
2295    /// });
2296    /// assert_eq!(
2297    ///     row.unpack_first().unwrap_list().iter().collect::<Vec<_>>(),
2298    ///     vec![Datum::String("age"), Datum::Int64(42)],
2299    /// );
2300    /// ```
2301    #[inline]
2302    pub fn push_list_with<F, R>(&mut self, f: F) -> R
2303    where
2304        F: FnOnce(&mut RowPacker) -> R,
2305    {
2306        // First, assume that the list will fit in 255 bytes, and thus the length will fit in
2307        // 1 byte. If not, we'll fix it up later.
2308        let start = self.row.data.len();
2309        self.row.data.push(Tag::ListTiny.into());
2310        // Write a dummy len, will fix it up later.
2311        self.row.data.push(0);
2312
2313        let out = f(self);
2314
2315        // The `- 1 - 1` is for the tag and the len.
2316        let len = self.row.data.len() - start - 1 - 1;
2317        // We now know the real len.
2318        if len < TINY {
2319            // If the len fits in 1 byte, we just need to fix up the len.
2320            self.row.data[start + 1] = len.to_le_bytes()[0];
2321        } else {
2322            // Note: We move this code path into its own function, so that the common case can be
2323            // inlined.
2324            long_list(&mut self.row.data, start, len);
2325        }
2326
2327        /// 1. Fix up the tag.
2328        /// 2. Move the actual data a bit (for which we also need to make room at the end).
2329        /// 3. Fix up the len.
2330        /// `data`: The row's backing data.
2331        /// `start`: where `push_list_with` started writing in `data`.
2332        /// `len`: the length of the data, excluding the tag and the length.
2333        #[cold]
2334        fn long_list(data: &mut CompactBytes, start: usize, len: usize) {
2335            // `len_len`: the length of the length. (Possible values are: 2, 4, 8. 1 is handled
2336            // elsewhere.) The other parameters are the same as for `long_list`.
2337            let long_list_inner = |data: &mut CompactBytes, len_len| {
2338                // We'll need memory for the new, bigger length, so make the `CompactBytes` bigger.
2339                // The `- 1` is because the old length was 1 byte.
2340                const ZEROS: [u8; 8] = [0; 8];
2341                data.extend_from_slice(&ZEROS[0..len_len - 1]);
2342                // Move the data to the end of the `CompactBytes`, to make space for the new length.
2343                // Originally, it started after the 1-byte tag and the 1-byte length, now it will
2344                // start after the 1-byte tag and the len_len-byte length.
2345                //
2346                // Note that this is the only operation in `long_list` whose cost is proportional
2347                // to `len`. Since `len` is at least 256 here, the other operations' cost are
2348                // negligible. `copy_within` is a memmove, which is probably a fair bit faster per
2349                // Datum than a Datum encoding in the `f` closure.
2350                data.copy_within(start + 1 + 1..start + 1 + 1 + len, start + 1 + len_len);
2351                // Write the new length.
2352                data[start + 1..start + 1 + len_len]
2353                    .copy_from_slice(&len.to_le_bytes()[0..len_len]);
2354            };
2355            match len {
2356                0..TINY => {
2357                    unreachable!()
2358                }
2359                TINY..SHORT => {
2360                    data[start] = Tag::ListShort.into();
2361                    long_list_inner(data, 2);
2362                }
2363                SHORT..LONG => {
2364                    data[start] = Tag::ListLong.into();
2365                    long_list_inner(data, 4);
2366                }
2367                _ => {
2368                    data[start] = Tag::ListHuge.into();
2369                    long_list_inner(data, 8);
2370                }
2371            };
2372        }
2373
2374        out
2375    }
2376
2377    /// Pushes a [`DatumMap`] that is built from a closure.
2378    ///
2379    /// The supplied closure will be invoked once with a `Row` that can be used
2380    /// to populate the dict.
2381    ///
2382    /// The closure **must** alternate pushing string keys and arbitrary values,
2383    /// otherwise reading the dict will cause a panic.
2384    ///
2385    /// The closure **must** push keys in ascending order, otherwise equality
2386    /// checks on the resulting `Row` may be wrong and reading the dict IN DEBUG
2387    /// MODE will cause a panic.
2388    ///
2389    /// The closure **must not** call [`RowPacker::clear`],
2390    /// [`RowPacker::truncate`], or [`RowPacker::truncate_datums`].
2391    ///
2392    /// # Example
2393    ///
2394    /// ```
2395    /// # use mz_repr::{Row, Datum};
2396    /// let mut row = Row::default();
2397    /// row.packer().push_dict_with(|row| {
2398    ///
2399    ///     // key
2400    ///     row.push(Datum::String("age"));
2401    ///     // value
2402    ///     row.push(Datum::Int64(42));
2403    ///
2404    ///     // key
2405    ///     row.push(Datum::String("name"));
2406    ///     // value
2407    ///     row.push(Datum::String("bob"));
2408    /// });
2409    /// assert_eq!(
2410    ///     row.unpack_first().unwrap_map().iter().collect::<Vec<_>>(),
2411    ///     vec![("age", Datum::Int64(42)), ("name", Datum::String("bob"))]
2412    /// );
2413    /// ```
2414    pub fn push_dict_with<F, R>(&mut self, f: F) -> R
2415    where
2416        F: FnOnce(&mut RowPacker) -> R,
2417    {
2418        self.row.data.push(Tag::Dict.into());
2419        let start = self.row.data.len();
2420        // write a dummy len, will fix it up later
2421        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2422
2423        let res = f(self);
2424
2425        let len = u64::cast_from(self.row.data.len() - start - size_of::<u64>());
2426        // fix up the len
2427        self.row.data[start..start + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2428
2429        res
2430    }
2431
2432    /// Like [`RowPacker::push_dict_with`], but accepts a fallible closure.
2433    pub fn try_push_dict_with<F, E>(&mut self, f: F) -> Result<(), E>
2434    where
2435        F: FnOnce(&mut RowPacker) -> Result<(), E>,
2436    {
2437        self.push_dict_with(f)
2438    }
2439
2440    /// Convenience function to construct an array from an iter of `Datum`s.
2441    ///
2442    /// Returns an error if the number of elements in `iter` does not match
2443    /// the cardinality of the array as described by `dims`, or if the
2444    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`]. If an error
2445    /// occurs, the packer's state will be unchanged.
2446    pub fn try_push_array<'a, I, D>(
2447        &mut self,
2448        dims: &[ArrayDimension],
2449        iter: I,
2450    ) -> Result<(), InvalidArrayError>
2451    where
2452        I: IntoIterator<Item = D>,
2453        D: Borrow<Datum<'a>>,
2454    {
2455        // SAFETY: The function returns the exact number of elements pushed into the array.
2456        unsafe {
2457            self.push_array_with_unchecked(dims, |packer| {
2458                let mut nelements = 0;
2459                for datum in iter {
2460                    packer.push(datum);
2461                    nelements += 1;
2462                }
2463                Ok::<_, InvalidArrayError>(nelements)
2464            })
2465        }
2466    }
2467
2468    /// Like [`RowPacker::try_push_array`], but accepts a fallible iterator of
2469    /// elements.
2470    pub fn try_push_array_fallible<'a, I, D, E>(
2471        &mut self,
2472        dims: &[ArrayDimension],
2473        iter: I,
2474    ) -> Result<Result<(), E>, InvalidArrayError>
2475    where
2476        I: IntoIterator<Item = Result<D, E>>,
2477        D: Borrow<Datum<'a>>,
2478    {
2479        enum Error<E> {
2480            Usage(InvalidArrayError),
2481            Inner(E),
2482        }
2483
2484        impl<E> From<InvalidArrayError> for Error<E> {
2485            fn from(e: InvalidArrayError) -> Self {
2486                Self::Usage(e)
2487            }
2488        }
2489
2490        // SAFETY: The function returns the exact number of elements pushed into the array.
2491        let result = unsafe {
2492            self.push_array_with_unchecked(dims, |packer| {
2493                let mut nelements = 0;
2494                for datum in iter {
2495                    packer.push(datum.map_err(Error::Inner)?);
2496                    nelements += 1;
2497                }
2498                Ok(nelements)
2499            })
2500        };
2501        match result {
2502            Ok(()) => Ok(Ok(())),
2503            Err(Error::Usage(e)) => Err(e),
2504            Err(Error::Inner(e)) => Ok(Err(e)),
2505        }
2506    }
2507
2508    /// Convenience function to construct an array from a function. The function must return the
2509    /// number of elements it pushed into the array. It is undefined behavior if the function returns
2510    /// a number different to the number of elements it pushed.
2511    ///
2512    /// Returns an error if the number of elements pushed by `f` does not match
2513    /// the cardinality of the array as described by `dims`, or if the
2514    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`], or if `f` errors. If an error
2515    /// occurs, the packer's state will be unchanged.
2516    pub unsafe fn push_array_with_unchecked<F, E>(
2517        &mut self,
2518        dims: &[ArrayDimension],
2519        f: F,
2520    ) -> Result<(), E>
2521    where
2522        F: FnOnce(&mut RowPacker) -> Result<usize, E>,
2523        E: From<InvalidArrayError>,
2524    {
2525        // Arrays are encoded as follows.
2526        //
2527        // u8    ndims
2528        // u64   dim_0 lower bound
2529        // u64   dim_0 length
2530        // ...
2531        // u64   dim_n lower bound
2532        // u64   dim_n length
2533        // u64   element data size in bytes
2534        // u8    element data, where elements are encoded in row-major order
2535
2536        if dims.len() > usize::from(MAX_ARRAY_DIMENSIONS) {
2537            return Err(InvalidArrayError::TooManyDimensions(dims.len()).into());
2538        }
2539
2540        let start = self.row.data.len();
2541        self.row.data.push(Tag::Array.into());
2542
2543        // Write dimension information.
2544        self.row
2545            .data
2546            .push(dims.len().try_into().expect("ndims verified to fit in u8"));
2547        for dim in dims {
2548            self.row
2549                .data
2550                .extend_from_slice(&i64::cast_from(dim.lower_bound).to_le_bytes());
2551            self.row
2552                .data
2553                .extend_from_slice(&u64::cast_from(dim.length).to_le_bytes());
2554        }
2555
2556        // Write elements.
2557        let off = self.row.data.len();
2558        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2559        let nelements = match f(self) {
2560            Ok(nelements) => nelements,
2561            Err(e) => {
2562                self.row.data.truncate(start);
2563                return Err(e);
2564            }
2565        };
2566        let len = u64::cast_from(self.row.data.len() - off - size_of::<u64>());
2567        self.row.data[off..off + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2568
2569        // Check that the number of elements written matches the dimension
2570        // information.
2571        let cardinality = match dims {
2572            [] => 0,
2573            // Saturate the product: a cardinality that overflows `usize` is
2574            // impossibly large (no array can hold that many elements), so it can
2575            // never equal the actual `nelements` and the check below rejects it as
2576            // `WrongCardinality`. A plain `product()` would panic under overflow
2577            // checks (debug/fuzz) and silently wrap in release — and a wrapped
2578            // value could even spuriously match `nelements`, accepting a corrupt
2579            // array (e.g. dims claiming `[2^32, 2^32]` wrap to 0 elements).
2580            dims => dims
2581                .iter()
2582                .map(|d| d.length)
2583                .fold(1usize, usize::saturating_mul),
2584        };
2585        if nelements != cardinality {
2586            self.row.data.truncate(start);
2587            return Err(InvalidArrayError::WrongCardinality {
2588                actual: nelements,
2589                expected: cardinality,
2590            }
2591            .into());
2592        }
2593
2594        Ok(())
2595    }
2596
2597    /// Pushes an [`Array`] that is built from a closure.
2598    ///
2599    /// __WARNING__: This is fairly "sharp" tool that is easy to get wrong. You
2600    /// should prefer [`RowPacker::try_push_array`] when possible.
2601    ///
2602    /// Returns an error if the number of elements pushed does not match
2603    /// the cardinality of the array as described by `dims`, or if the
2604    /// number of dimensions exceeds [`MAX_ARRAY_DIMENSIONS`]. If an error
2605    /// occurs, the packer's state will be unchanged.
2606    pub fn push_array_with_row_major<F, I>(
2607        &mut self,
2608        dims: I,
2609        f: F,
2610    ) -> Result<(), InvalidArrayError>
2611    where
2612        I: IntoIterator<Item = ArrayDimension>,
2613        F: FnOnce(&mut RowPacker) -> usize,
2614    {
2615        let start = self.row.data.len();
2616        self.row.data.push(Tag::Array.into());
2617
2618        // Write dummy dimension length for now, we'll fix it up.
2619        let dims_start = self.row.data.len();
2620        self.row.data.push(42);
2621
2622        let mut num_dims: u8 = 0;
2623        let mut cardinality: usize = 1;
2624        for dim in dims {
2625            num_dims += 1;
2626            // Saturate: an overflowing cardinality is impossibly large and is
2627            // rejected by the `nelements` check below. See the matching note in
2628            // `push_array_with_unchecked`.
2629            cardinality = cardinality.saturating_mul(dim.length);
2630
2631            self.row
2632                .data
2633                .extend_from_slice(&i64::cast_from(dim.lower_bound).to_le_bytes());
2634            self.row
2635                .data
2636                .extend_from_slice(&u64::cast_from(dim.length).to_le_bytes());
2637        }
2638
2639        if num_dims > MAX_ARRAY_DIMENSIONS {
2640            // Reset the packer state so we don't have invalid data.
2641            self.row.data.truncate(start);
2642            return Err(InvalidArrayError::TooManyDimensions(usize::from(num_dims)));
2643        }
2644        // Fix up our dimension length.
2645        self.row.data[dims_start..dims_start + size_of::<u8>()]
2646            .copy_from_slice(&num_dims.to_le_bytes());
2647
2648        // Write elements.
2649        let off = self.row.data.len();
2650        self.row.data.extend_from_slice(&[0; size_of::<u64>()]);
2651
2652        let nelements = f(self);
2653
2654        let len = u64::cast_from(self.row.data.len() - off - size_of::<u64>());
2655        self.row.data[off..off + size_of::<u64>()].copy_from_slice(&len.to_le_bytes());
2656
2657        // Check that the number of elements written matches the dimension
2658        // information.
2659        let cardinality = match num_dims {
2660            0 => 0,
2661            _ => cardinality,
2662        };
2663        if nelements != cardinality {
2664            self.row.data.truncate(start);
2665            return Err(InvalidArrayError::WrongCardinality {
2666                actual: nelements,
2667                expected: cardinality,
2668            });
2669        }
2670
2671        Ok(())
2672    }
2673
2674    /// Convenience function to push a `DatumList` from an iter of `Datum`s
2675    ///
2676    /// See [`RowPacker::push_dict_with`] if you need to be able to handle errors
2677    pub fn push_list<'a, I, D>(&mut self, iter: I)
2678    where
2679        I: IntoIterator<Item = D>,
2680        D: Borrow<Datum<'a>>,
2681    {
2682        self.push_list_with(|packer| {
2683            for elem in iter {
2684                packer.push(*elem.borrow())
2685            }
2686        });
2687    }
2688
2689    /// Convenience function to push a `DatumMap` from an iter of `(&str, Datum)` pairs
2690    pub fn push_dict<'a, I, D>(&mut self, iter: I)
2691    where
2692        I: IntoIterator<Item = (&'a str, D)>,
2693        D: Borrow<Datum<'a>>,
2694    {
2695        self.push_dict_with(|packer| {
2696            for (k, v) in iter {
2697                packer.push(Datum::String(k));
2698                packer.push(*v.borrow())
2699            }
2700        })
2701    }
2702
2703    /// Pushes a `Datum::Range` derived from the `Range<Datum<'a>`.
2704    ///
2705    /// # Panics
2706    /// - If lower and upper express finite values and they are datums of
2707    ///   different types.
2708    /// - If lower or upper express finite values and are equal to
2709    ///   `Datum::Null`. To handle `Datum::Null` properly, use
2710    ///   [`RangeBound::new`].
2711    ///
2712    /// # Notes
2713    /// - This function canonicalizes the range before pushing it to the row.
2714    /// - Prefer this function over `push_range_with` because of its
2715    ///   canonicaliztion.
2716    /// - Prefer creating [`RangeBound`]s using [`RangeBound::new`], which
2717    ///   handles `Datum::Null` in a SQL-friendly way.
2718    pub fn push_range<'a>(&mut self, mut range: Range<Datum<'a>>) -> Result<(), InvalidRangeError> {
2719        range.canonicalize()?;
2720        match range.inner {
2721            None => {
2722                self.row.data.push(Tag::Range.into());
2723                // Untagged bytes only contains the `RANGE_EMPTY` flag value.
2724                self.row.data.push(range::InternalFlags::EMPTY.bits());
2725                Ok(())
2726            }
2727            Some(inner) => self.push_range_with(
2728                RangeLowerBound {
2729                    inclusive: inner.lower.inclusive,
2730                    bound: inner
2731                        .lower
2732                        .bound
2733                        .map(|value| move |row: &mut RowPacker| Ok(row.push(value))),
2734                },
2735                RangeUpperBound {
2736                    inclusive: inner.upper.inclusive,
2737                    bound: inner
2738                        .upper
2739                        .bound
2740                        .map(|value| move |row: &mut RowPacker| Ok(row.push(value))),
2741                },
2742            ),
2743        }
2744    }
2745
2746    /// Pushes a `DatumRange` built from the specified arguments.
2747    ///
2748    /// # Warning
2749    /// Unlike `push_range`, `push_range_with` _does not_ canonicalize its
2750    /// inputs. Consequentially, this means it's possible to generate ranges
2751    /// that will not reflect the proper ordering and equality.
2752    ///
2753    /// # Panics
2754    /// - If lower or upper expresses a finite value and does not push exactly
2755    ///   one value into the `RowPacker`.
2756    /// - If lower and upper express finite values and they are datums of
2757    ///   different types.
2758    /// - If lower or upper express finite values and push `Datum::Null`.
2759    ///
2760    /// # Notes
2761    /// - Prefer `push_range_with` over this function. This function should be
2762    ///   used only when you are not pushing `Datum`s to the inner row.
2763    /// - Range encoding is `[<flag bytes>,<lower>?,<upper>?]`, where `lower`
2764    ///   and `upper` are optional, contingent on the flag value expressing an
2765    ///   empty range (where neither will be present) or infinite bounds (where
2766    ///   each infinite bound will be absent).
2767    /// - To push an emtpy range, use `push_range` using `Range { inner: None }`.
2768    pub fn push_range_with<L, U, E>(
2769        &mut self,
2770        lower: RangeLowerBound<L>,
2771        upper: RangeUpperBound<U>,
2772    ) -> Result<(), E>
2773    where
2774        L: FnOnce(&mut RowPacker) -> Result<(), E>,
2775        U: FnOnce(&mut RowPacker) -> Result<(), E>,
2776        E: From<InvalidRangeError>,
2777    {
2778        let start = self.row.data.len();
2779        self.row.data.push(Tag::Range.into());
2780
2781        let mut flags = range::InternalFlags::empty();
2782
2783        flags.set(range::InternalFlags::LB_INFINITE, lower.bound.is_none());
2784        flags.set(range::InternalFlags::UB_INFINITE, upper.bound.is_none());
2785        flags.set(range::InternalFlags::LB_INCLUSIVE, lower.inclusive);
2786        flags.set(range::InternalFlags::UB_INCLUSIVE, upper.inclusive);
2787
2788        let mut expected_datums = 0;
2789
2790        self.row.data.push(flags.bits());
2791
2792        let datum_check = self.row.data.len();
2793
2794        if let Some(value) = lower.bound {
2795            let start = self.row.data.len();
2796            value(self)?;
2797            assert!(
2798                start < self.row.data.len(),
2799                "finite values must each push exactly one value; expected 1 but got 0"
2800            );
2801            expected_datums += 1;
2802        }
2803
2804        if let Some(value) = upper.bound {
2805            let start = self.row.data.len();
2806            value(self)?;
2807            assert!(
2808                start < self.row.data.len(),
2809                "finite values must each push exactly one value; expected 1 but got 0"
2810            );
2811            expected_datums += 1;
2812        }
2813
2814        // Validate the invariants that 0, 1, or 2 elements were pushed, none are Null,
2815        // and if two are pushed then the second is not less than the first. Panic in
2816        // some cases and error in others.
2817        let mut actual_datums = 0;
2818        let mut seen = None;
2819        let mut dataz = &self.row.data[datum_check..];
2820        while !dataz.is_empty() {
2821            let d = unsafe { read_datum(&mut dataz) };
2822            // These checks only fail when decoding untrusted/corrupted bytes;
2823            // valid callers always push consistent, non-null bounds. Return an
2824            // error rather than asserting so a crafted proto doesn't panic.
2825            if d == Datum::Null {
2826                self.row.data.truncate(start);
2827                return Err(InvalidRangeError::InvalidRangeData.into());
2828            }
2829
2830            match seen {
2831                None => seen = Some(d),
2832                Some(seen) => {
2833                    let seen_kind = DatumKind::from(seen);
2834                    let d_kind = DatumKind::from(d);
2835                    if seen_kind != d_kind {
2836                        self.row.data.truncate(start);
2837                        return Err(InvalidRangeError::InvalidRangeData.into());
2838                    }
2839
2840                    if seen > d {
2841                        self.row.data.truncate(start);
2842                        return Err(InvalidRangeError::MisorderedRangeBounds.into());
2843                    }
2844                }
2845            }
2846            actual_datums += 1;
2847        }
2848
2849        if actual_datums != expected_datums {
2850            self.row.data.truncate(start);
2851            return Err(InvalidRangeError::InvalidRangeData.into());
2852        }
2853
2854        Ok(())
2855    }
2856
2857    /// Clears the contents of the packer without de-allocating its backing memory.
2858    pub fn clear(&mut self) {
2859        self.row.data.clear();
2860    }
2861
2862    /// Truncates the underlying storage to the specified byte position.
2863    ///
2864    /// # Safety
2865    ///
2866    /// `pos` MUST specify a byte offset that lies on a datum boundary.
2867    /// If `pos` specifies a byte offset that is *within* a datum, the row
2868    /// packer will produce an invalid row, the unpacking of which may
2869    /// trigger undefined behavior!
2870    ///
2871    /// To find the byte offset of a datum boundary, inspect the packer's
2872    /// byte length by calling `packer.data().len()` after pushing the desired
2873    /// number of datums onto the packer.
2874    pub unsafe fn truncate(&mut self, pos: usize) {
2875        self.row.data.truncate(pos)
2876    }
2877
2878    /// Truncates the underlying row to contain at most the first `n` datums.
2879    pub fn truncate_datums(&mut self, n: usize) {
2880        let prev_len = self.row.data.len();
2881        let mut iter = self.row.iter();
2882        for _ in iter.by_ref().take(n) {}
2883        let next_len = iter.data.len();
2884        // SAFETY: iterator offsets always lie on a datum boundary.
2885        unsafe { self.truncate(prev_len - next_len) }
2886    }
2887
2888    /// Returns the total amount of bytes used by the underlying row.
2889    pub fn byte_len(&self) -> usize {
2890        self.row.byte_len()
2891    }
2892}
2893
2894impl<'a> IntoIterator for &'a Row {
2895    type Item = Datum<'a>;
2896    type IntoIter = DatumListIter<'a>;
2897    fn into_iter(self) -> DatumListIter<'a> {
2898        self.iter()
2899    }
2900}
2901
2902impl fmt::Debug for Row {
2903    /// Debug representation using the internal datums
2904    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2905        f.write_str("Row{")?;
2906        f.debug_list().entries(self.iter()).finish()?;
2907        f.write_str("}")
2908    }
2909}
2910
2911impl fmt::Display for Row {
2912    /// Display representation using the internal datums
2913    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2914        f.write_str("(")?;
2915        for (i, datum) in self.iter().enumerate() {
2916            if i != 0 {
2917                f.write_str(", ")?;
2918            }
2919            write!(f, "{}", datum)?;
2920        }
2921        f.write_str(")")
2922    }
2923}
2924
2925impl<'a, T> DatumList<'a, T> {
2926    pub fn iter(&self) -> DatumListIter<'a> {
2927        DatumListIter { data: self.data }
2928    }
2929
2930    /// Iterate elements as typed `T` values rather than raw `Datum`s.
2931    ///
2932    /// Each datum is decoded and converted via [`FromDatum`]. Since generic
2933    /// type parameters in `#[sqlfunc]` are erased to `Datum<'a>` before code
2934    /// generation, this is monomorphized to an identity conversion at runtime.
2935    pub fn typed_iter(&self) -> DatumListTypedIter<'a, T>
2936    where
2937        T: FromDatum<'a>,
2938    {
2939        DatumListTypedIter {
2940            inner: self.iter(),
2941            _phantom: PhantomData,
2942        }
2943    }
2944
2945    /// For debugging only
2946    pub fn data(&self) -> &'a [u8] {
2947        self.data
2948    }
2949}
2950
2951impl<T> DatumList<'static, T> {
2952    pub fn empty() -> Self {
2953        DatumList::new(&[])
2954    }
2955}
2956
2957impl<'a> IntoIterator for DatumList<'a> {
2958    type Item = Datum<'a>;
2959    type IntoIter = DatumListIter<'a>;
2960    fn into_iter(self) -> DatumListIter<'a> {
2961        self.iter()
2962    }
2963}
2964
2965impl<'a> Iterator for DatumListIter<'a> {
2966    type Item = Datum<'a>;
2967    fn next(&mut self) -> Option<Self::Item> {
2968        if self.data.is_empty() {
2969            None
2970        } else {
2971            Some(unsafe { read_datum(&mut self.data) })
2972        }
2973    }
2974}
2975
2976impl<'a, T: FromDatum<'a>> Iterator for DatumListTypedIter<'a, T> {
2977    type Item = T;
2978    fn next(&mut self) -> Option<Self::Item> {
2979        self.inner.next().map(T::from_datum)
2980    }
2981}
2982
2983impl<'a, T> DatumMap<'a, T> {
2984    pub fn iter(&self) -> DatumDictIter<'a> {
2985        DatumDictIter {
2986            data: self.data,
2987            prev_key: None,
2988        }
2989    }
2990
2991    /// Iterate entries as `(&str, T)` pairs rather than `(&str, Datum)`.
2992    ///
2993    /// Each value datum is converted via [`FromDatum`]. Since generic type
2994    /// parameters in `#[sqlfunc]` are erased to `Datum<'a>` before code
2995    /// generation, this is monomorphized to an identity conversion at runtime.
2996    pub fn typed_iter(&self) -> DatumDictTypedIter<'a, T>
2997    where
2998        T: FromDatum<'a>,
2999    {
3000        DatumDictTypedIter {
3001            inner: self.iter(),
3002            _phantom: PhantomData,
3003        }
3004    }
3005
3006    /// For debugging only
3007    pub fn data(&self) -> &'a [u8] {
3008        self.data
3009    }
3010}
3011
3012impl<T> DatumMap<'static, T> {
3013    pub fn empty() -> Self {
3014        DatumMap::new(&[])
3015    }
3016}
3017
3018impl<'a, T> Debug for DatumMap<'a, T> {
3019    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
3020        f.debug_map().entries(self.iter()).finish()
3021    }
3022}
3023
3024impl<'a> IntoIterator for &'a DatumMap<'a> {
3025    type Item = (&'a str, Datum<'a>);
3026    type IntoIter = DatumDictIter<'a>;
3027    fn into_iter(self) -> DatumDictIter<'a> {
3028        self.iter()
3029    }
3030}
3031
3032impl<'a> Iterator for DatumDictIter<'a> {
3033    type Item = (&'a str, Datum<'a>);
3034    fn next(&mut self) -> Option<Self::Item> {
3035        if self.data.is_empty() {
3036            None
3037        } else {
3038            let key_tag =
3039                Tag::try_from_primitive(read_byte(&mut self.data)).expect("unknown row tag");
3040            assert!(
3041                key_tag == Tag::StringTiny
3042                    || key_tag == Tag::StringShort
3043                    || key_tag == Tag::StringLong
3044                    || key_tag == Tag::StringHuge,
3045                "Dict keys must be strings, got {:?}",
3046                key_tag
3047            );
3048            let key = unsafe { read_lengthed_datum(&mut self.data, key_tag).unwrap_str() };
3049            let val = unsafe { read_datum(&mut self.data) };
3050
3051            // if in debug mode, sanity check keys
3052            if cfg!(debug_assertions) {
3053                if let Some(prev_key) = self.prev_key {
3054                    debug_assert!(
3055                        prev_key < key,
3056                        "Dict keys must be unique and given in ascending order: {} came before {}",
3057                        prev_key,
3058                        key
3059                    );
3060                }
3061                self.prev_key = Some(key);
3062            }
3063
3064            Some((key, val))
3065        }
3066    }
3067}
3068
3069impl<'a, T: FromDatum<'a>> Iterator for DatumDictTypedIter<'a, T> {
3070    type Item = (&'a str, T);
3071    fn next(&mut self) -> Option<Self::Item> {
3072        self.inner.next().map(|(k, v)| (k, T::from_datum(v)))
3073    }
3074}
3075
3076impl RowArena {
3077    pub fn new() -> Self {
3078        RowArena {
3079            inner: RefCell::new(vec![]),
3080            scratch: RefCell::new(None),
3081        }
3082    }
3083
3084    /// Creates a `RowArena` with an initial region sized to hold `capacity` bytes, to avoid
3085    /// reallocations as the first datums are created in the arena.
3086    pub fn with_capacity(capacity: usize) -> Self {
3087        let mut inner = Vec::new();
3088        if capacity > 0 {
3089            inner.push(Vec::with_capacity(capacity));
3090        }
3091        RowArena {
3092            inner: RefCell::new(inner),
3093            scratch: RefCell::new(None),
3094        }
3095    }
3096
3097    /// Ensures the active region can hold at least `additional` more bytes without allocating a
3098    /// new region. Call this when you expect to push roughly `additional` bytes next.
3099    pub fn reserve(&self, additional: usize) {
3100        if additional == 0 {
3101            return;
3102        }
3103        let mut inner = self.inner.borrow_mut();
3104        match inner.last_mut() {
3105            // The active region is empty, so nothing references it yet and it is safe to grow it
3106            // in place (a reallocation cannot dangle a live reference).
3107            Some(active) if active.is_empty() => {
3108                if active.capacity() < additional {
3109                    active.reserve_exact(additional);
3110                }
3111            }
3112            // The active region holds live data; we cannot grow it without moving those bytes, so
3113            // stage a fresh region. Size it like `push_bytes` does (at least double the current
3114            // region) so a sequence of small `reserve`s still yields at most log-many regions
3115            // rather than many small ones.
3116            Some(active) => {
3117                let new_cap = std::cmp::max(additional, active.capacity().saturating_mul(2));
3118                inner.push(Vec::with_capacity(new_cap));
3119            }
3120            None => inner.push(Vec::with_capacity(additional)),
3121        }
3122    }
3123
3124    /// Copies `bytes` into the arena and returns a reference valid for its lifetime.
3125    ///
3126    /// Accepts anything that derefs to `[u8]` (e.g. `Vec<u8>`, `&[u8]`); the bytes are copied, so
3127    /// the caller's allocation is not retained.
3128    #[allow(clippy::transmute_ptr_to_ptr)]
3129    pub fn push_bytes<'a, B: Deref<Target = [u8]>>(&'a self, bytes: B) -> &'a [u8] {
3130        let bytes: &[u8] = &bytes;
3131        let need = bytes.len();
3132        if need == 0 {
3133            return &[];
3134        }
3135        let mut inner = self.inner.borrow_mut();
3136
3137        // Find or create a region with spare capacity for `need` bytes, never growing a region
3138        // that already holds data (see the type-level comment for why this preserves references).
3139        let has_room = inner
3140            .last()
3141            .map_or(false, |region| region.capacity() - region.len() >= need);
3142        if !has_room {
3143            let last_cap = inner.last().map_or(0, |region| region.capacity());
3144            let new_cap = std::cmp::max(need, last_cap.saturating_mul(2));
3145            inner.push(Vec::with_capacity(new_cap));
3146        }
3147
3148        let region = inner.last_mut().expect("region present");
3149        let start = region.len();
3150        region.extend_from_slice(bytes);
3151        let copied = &region[start..];
3152        unsafe {
3153            // This is safe because:
3154            //   * `copied` references bytes inside `region`'s heap buffer, which we just sized to
3155            //     fit without reallocating; that buffer is never resized again while it holds data
3156            //     (we allocate a new region instead), so the reference stays valid.
3157            //   * The buffer lives as long as the arena: regions are only dropped by `clear`/`drop`,
3158            //     both of which take `&mut`/ownership, so no `&'a self`-tied reference can outlive
3159            //     them.
3160            //   * Pushing further regions may reallocate `self.inner`, but that moves only the
3161            //     `Vec<u8>` headers, not the heap buffers they own.
3162            transmute::<&[u8], &'a [u8]>(copied)
3163        }
3164    }
3165
3166    /// Copies `string` into the arena and returns a reference valid for its lifetime.
3167    pub fn push_string<'a>(&'a self, string: String) -> &'a str {
3168        let copied = self.push_bytes(string.as_bytes());
3169        unsafe {
3170            // This is safe because we just copied the bytes of a valid `String`.
3171            std::str::from_utf8_unchecked(copied)
3172        }
3173    }
3174
3175    /// Returns a growable, writeable byte buffer for assembling a value incrementally.
3176    ///
3177    /// Write into it with [`RowArenaBuf::push`], [`RowArenaBuf::extend_from_slice`], or
3178    /// [`std::io::Write`], then call [`RowArenaBuf::finish`] to copy the result into the arena and
3179    /// obtain a reference valid for the arena's lifetime. The backing buffer is a single scratch
3180    /// allocation reused across writers, so this lets a producer that builds bytes piecewise (e.g.
3181    /// decoding a row) avoid managing its own scratch.
3182    ///
3183    /// Nested writers are sound but not free: a writer obtained while another is still live can't
3184    /// reuse the (in-use) scratch, so it allocates its own buffer. Steady-state, non-nested use
3185    /// stays allocation-free.
3186    pub fn writer(&self) -> RowArenaBuf<'_> {
3187        // Take the recycled buffer if one is available, else allocate a fresh one. The cell is
3188        // borrowed only for this `take`, never for the writer's lifetime, so a nested `writer` call
3189        // doesn't double-borrow: it simply finds the slot empty and allocates its own buffer.
3190        let mut buf = self.scratch.borrow_mut().take().unwrap_or_default();
3191        buf.clear();
3192        RowArenaBuf { arena: self, buf }
3193    }
3194
3195    /// Take ownership of `row` for the lifetime of the arena, returning a
3196    /// reference to the first datum in the row.
3197    ///
3198    /// If we had an owned datum type, this method would be much clearer, and
3199    /// would be called `push_owned_datum`.
3200    pub fn push_unary_row<'a>(&'a self, row: Row) -> Datum<'a> {
3201        let copied = self.push_bytes(row.data());
3202        unsafe {
3203            // This is safe because `copied` is a valid encoding of a single datum (we just packed
3204            // it into `row`), backed by the arena for the lifetime `'a`. Copying the bytes also
3205            // sidesteps the `Row`'s inline (`SmallVec`) storage entirely.
3206            let datum = read_datum(&mut &copied[..]);
3207            transmute::<Datum<'_>, Datum<'a>>(datum)
3208        }
3209    }
3210
3211    /// Equivalent to `push_unary_row` but returns a `DatumNested` rather than a
3212    /// `Datum`.
3213    fn push_unary_row_datum_nested<'a>(&'a self, row: Row) -> DatumNested<'a> {
3214        let copied = self.push_bytes(row.data());
3215        unsafe {
3216            // Safe for the same reasons as `push_unary_row`.
3217            let nested = DatumNested::extract(&mut &copied[..]);
3218            transmute::<DatumNested<'_>, DatumNested<'a>>(nested)
3219        }
3220    }
3221
3222    /// Convenience function to make a new `Row` containing a single datum, and
3223    /// take ownership of it for the lifetime of the arena
3224    ///
3225    /// ```
3226    /// # use mz_repr::{RowArena, Datum};
3227    /// let arena = RowArena::new();
3228    /// let datum = arena.make_datum(|packer| {
3229    ///   packer.push_list(&[Datum::String("hello"), Datum::String("world")]);
3230    /// });
3231    /// assert_eq!(datum.unwrap_list().iter().collect::<Vec<_>>(), vec![Datum::String("hello"), Datum::String("world")]);
3232    /// ```
3233    pub fn make_datum<'a, F>(&'a self, f: F) -> Datum<'a>
3234    where
3235        F: FnOnce(&mut RowPacker),
3236    {
3237        let mut row = Row::default();
3238        f(&mut row.packer());
3239        self.push_unary_row(row)
3240    }
3241
3242    /// Convenience function to build a list datum from an iterator of typed
3243    /// elements and return it as a `DatumList<'a, T>`.
3244    ///
3245    /// By accepting an iterator of `T: Borrow<Datum>` instead of a raw
3246    /// `RowPacker` closure, this guarantees that only elements of type `T`
3247    /// are pushed.
3248    pub fn make_datum_list<'a, T: std::borrow::Borrow<Datum<'a>>>(
3249        &'a self,
3250        iter: impl IntoIterator<Item = T>,
3251    ) -> DatumList<'a, T> {
3252        let datum = self.make_datum(|packer| {
3253            packer.push_list_with(|packer| {
3254                for elem in iter {
3255                    packer.push(*elem.borrow());
3256                }
3257            });
3258        });
3259        DatumList::new(datum.unwrap_list().data())
3260    }
3261
3262    /// Convenience function identical to `make_datum` but instead returns a
3263    /// `DatumNested`.
3264    pub fn make_datum_nested<'a, F>(&'a self, f: F) -> DatumNested<'a>
3265    where
3266        F: FnOnce(&mut RowPacker),
3267    {
3268        let mut row = Row::default();
3269        f(&mut row.packer());
3270        self.push_unary_row_datum_nested(row)
3271    }
3272
3273    /// Like [`RowArena::make_datum`], but the provided closure can return an error.
3274    pub fn try_make_datum<'a, F, E>(&'a self, f: F) -> Result<Datum<'a>, E>
3275    where
3276        F: FnOnce(&mut RowPacker) -> Result<(), E>,
3277    {
3278        let mut row = Row::default();
3279        f(&mut row.packer())?;
3280        Ok(self.push_unary_row(row))
3281    }
3282
3283    /// Clear the contents of the arena.
3284    ///
3285    /// Retains the single largest region (emptied) so the arena can be reused without
3286    /// reallocating; a workload that clears between uses of similar size becomes allocation-free.
3287    pub fn clear(&mut self) {
3288        let inner = self.inner.get_mut();
3289        // Keep only the largest-capacity region, reset to empty, and drop the rest. Because region
3290        // capacities only ever grow (each new region at least doubles the previous), the largest is
3291        // normally the last; we scan for it defensively, which is cheap given log-many regions.
3292        if let Some(largest) = (0..inner.len()).max_by_key(|&i| inner[i].capacity()) {
3293            inner.swap(0, largest);
3294            inner.truncate(1);
3295            inner[0].clear();
3296        }
3297    }
3298}
3299
3300impl Default for RowArena {
3301    fn default() -> RowArena {
3302        RowArena::new()
3303    }
3304}
3305
3306/// A growable, writeable byte buffer that builds a value into a [`RowArena`].
3307///
3308/// Obtained from [`RowArena::writer`]. Behaves like a writeable byte slice (push/extend bytes,
3309/// read back as `&[u8]`); [`RowArenaBuf::finish`] copies the assembled bytes into the arena and
3310/// returns a reference valid for the arena's lifetime. The buffer is owned for the writer's
3311/// lifetime and, on drop, returned to the arena to be reused by the next writer.
3312#[derive(Debug)]
3313pub struct RowArenaBuf<'a> {
3314    arena: &'a RowArena,
3315    buf: Vec<u8>,
3316}
3317
3318impl<'a> RowArenaBuf<'a> {
3319    /// Appends a single byte.
3320    pub fn push(&mut self, byte: u8) {
3321        self.buf.push(byte);
3322    }
3323
3324    /// Appends a slice of bytes.
3325    pub fn extend_from_slice(&mut self, bytes: &[u8]) {
3326        self.buf.extend_from_slice(bytes);
3327    }
3328
3329    /// The bytes written so far.
3330    pub fn as_slice(&self) -> &[u8] {
3331        &self.buf
3332    }
3333
3334    /// The number of bytes written so far.
3335    pub fn len(&self) -> usize {
3336        self.buf.len()
3337    }
3338
3339    /// Whether no bytes have been written.
3340    pub fn is_empty(&self) -> bool {
3341        self.buf.is_empty()
3342    }
3343
3344    /// Copies the written bytes into the arena, returning a reference valid for its lifetime.
3345    pub fn finish(self) -> &'a [u8] {
3346        // `self` is dropped at the end of this call, returning `buf` to the arena for reuse; the
3347        // returned reference points into a committed region, not `buf`, so it stays valid.
3348        self.arena.push_bytes(self.buf.as_slice())
3349    }
3350
3351    /// Like [`RowArenaBuf::finish`], but returns the bytes as a `&str`.
3352    ///
3353    /// Intended for buffers written via [`std::fmt::Write`] (e.g. `write!`), whose contents are
3354    /// valid UTF-8. Panics if the bytes are not valid UTF-8.
3355    pub fn finish_str(self) -> &'a str {
3356        let bytes = self.arena.push_bytes(self.buf.as_slice());
3357        std::str::from_utf8(bytes).expect("RowArenaBuf::finish_str on non-UTF-8 contents")
3358    }
3359}
3360
3361impl<'a> Drop for RowArenaBuf<'a> {
3362    fn drop(&mut self) {
3363        // Return the buffer to the arena so the next writer can reuse its allocation. We keep only
3364        // one buffer: if the slot is already occupied — an outer writer is still live, or a nested
3365        // writer beat us to it — we drop ours rather than growing an unbounded pool. The borrow is
3366        // transient and never overlaps a live writer's, so this can't double-borrow.
3367        let mut slot = self.arena.scratch.borrow_mut();
3368        if slot.is_none() {
3369            *slot = Some(std::mem::take(&mut self.buf));
3370        }
3371    }
3372}
3373
3374impl<'a> std::ops::Deref for RowArenaBuf<'a> {
3375    type Target = [u8];
3376    fn deref(&self) -> &[u8] {
3377        &self.buf
3378    }
3379}
3380
3381impl<'a> std::io::Write for RowArenaBuf<'a> {
3382    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
3383        self.buf.extend_from_slice(bytes);
3384        Ok(bytes.len())
3385    }
3386
3387    fn flush(&mut self) -> std::io::Result<()> {
3388        Ok(())
3389    }
3390}
3391
3392impl<'a> std::fmt::Write for RowArenaBuf<'a> {
3393    fn write_str(&mut self, s: &str) -> std::fmt::Result {
3394        self.buf.extend_from_slice(s.as_bytes());
3395        Ok(())
3396    }
3397}
3398
3399/// A thread-local row, which can be borrowed and returned.
3400/// # Example
3401///
3402/// Use this type instead of creating a new row:
3403/// ```
3404/// use mz_repr::SharedRow;
3405///
3406/// let mut row_builder = SharedRow::get();
3407/// ```
3408///
3409/// This allows us to reuse an existing row allocation instead of creating a new one or retaining
3410/// an allocation locally. Additionally, we can observe the size of the local row in a central
3411/// place and potentially reallocate to reduce memory needs.
3412///
3413/// # Panic
3414///
3415/// [`SharedRow::get`] panics when trying to obtain multiple references to the shared row.
3416#[derive(Debug)]
3417pub struct SharedRow(Row);
3418
3419impl SharedRow {
3420    thread_local! {
3421        /// A thread-local slot containing a shared Row that can be temporarily used by a function.
3422        /// There can be at most one active user of this Row, which is tracked by the state of the
3423        /// `Option<_>` wrapper. When it is `Some(..)`, the row is available for using. When it
3424        /// is `None`, it is not, and the constructor will panic if a thread attempts to use it.
3425        static SHARED_ROW: Cell<Option<Row>> = const { Cell::new(Some(Row::empty())) }
3426    }
3427
3428    /// Get the shared row.
3429    ///
3430    /// The row's contents are cleared before returning it.
3431    ///
3432    /// # Panic
3433    ///
3434    /// Panics when the row is already borrowed elsewhere.
3435    pub fn get() -> Self {
3436        let mut row = Self::SHARED_ROW
3437            .take()
3438            .expect("attempted to borrow already borrowed SharedRow");
3439        // Clear row
3440        row.packer();
3441        Self(row)
3442    }
3443
3444    /// Gets the shared row and uses it to pack `iter`.
3445    pub fn pack<'a, I, D>(iter: I) -> Row
3446    where
3447        I: IntoIterator<Item = D>,
3448        D: Borrow<Datum<'a>>,
3449    {
3450        let mut row_builder = Self::get();
3451        let mut row_packer = row_builder.packer();
3452        row_packer.extend(iter);
3453        row_builder.clone()
3454    }
3455}
3456
3457impl std::ops::Deref for SharedRow {
3458    type Target = Row;
3459
3460    fn deref(&self) -> &Self::Target {
3461        &self.0
3462    }
3463}
3464
3465impl std::ops::DerefMut for SharedRow {
3466    fn deref_mut(&mut self) -> &mut Self::Target {
3467        &mut self.0
3468    }
3469}
3470
3471impl Drop for SharedRow {
3472    fn drop(&mut self) {
3473        // Take the Row allocation from this instance and put it back in the thread local slot for
3474        // the next user. The Row in `self` is replaced with an empty Row which does not allocate.
3475        Self::SHARED_ROW.set(Some(std::mem::take(&mut self.0)))
3476    }
3477}
3478
3479#[cfg(test)]
3480mod tests {
3481    use std::cmp::Ordering;
3482    use std::collections::hash_map::DefaultHasher;
3483    use std::hash::{Hash, Hasher};
3484
3485    use chrono::{DateTime, NaiveDate};
3486    use itertools::Itertools;
3487    use mz_ore::{assert_err, assert_none};
3488    use ordered_float::OrderedFloat;
3489
3490    use crate::SqlScalarType;
3491
3492    use super::*;
3493
3494    // Regression: comparing deeply nested list values must not overflow the
3495    // stack (STACK-7). `Datum` ordering recurses once per nesting level.
3496    #[mz_ore::test]
3497    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
3498    fn cmp_deep_nested_list_does_not_overflow() {
3499        fn deep() -> Row {
3500            // `push_list` byte-copies the inner value, so building does not recurse.
3501            let mut row = Row::pack_slice(&[Datum::Int64(1)]);
3502            for _ in 0..50_000 {
3503                let mut next = Row::default();
3504                next.packer().push_list([row.unpack_first()]);
3505                row = next;
3506            }
3507            row
3508        }
3509        let a = deep();
3510        let b = deep();
3511        assert_eq!(a.unpack_first().cmp(&b.unpack_first()), Ordering::Equal);
3512    }
3513
3514    fn hash<T: Hash>(t: &T) -> u64 {
3515        let mut hasher = DefaultHasher::new();
3516        t.hash(&mut hasher);
3517        hasher.finish()
3518    }
3519
3520    #[mz_ore::test]
3521    fn test_assumptions() {
3522        assert_eq!(size_of::<Tag>(), 1);
3523        #[cfg(target_endian = "big")]
3524        {
3525            // if you want to run this on a big-endian cpu, we'll need big-endian versions of the serialization code
3526            assert!(false);
3527        }
3528    }
3529
3530    #[mz_ore::test]
3531    fn miri_test_arena() {
3532        let arena = RowArena::new();
3533
3534        assert_eq!(arena.push_string("".to_owned()), "");
3535        assert_eq!(arena.push_string("العَرَبِيَّة".to_owned()), "العَرَبِيَّة");
3536
3537        let empty: &[u8] = &[];
3538        assert_eq!(arena.push_bytes(vec![]), empty);
3539        assert_eq!(arena.push_bytes(vec![0, 2, 1, 255]), &[0, 2, 1, 255]);
3540
3541        let mut row = Row::default();
3542        let mut packer = row.packer();
3543        packer.push_dict_with(|row| {
3544            row.push(Datum::String("a"));
3545            row.push_list_with(|row| {
3546                row.push(Datum::String("one"));
3547                row.push(Datum::String("two"));
3548                row.push(Datum::String("three"));
3549            });
3550            row.push(Datum::String("b"));
3551            row.push(Datum::String("c"));
3552        });
3553        assert_eq!(arena.push_unary_row(row.clone()), row.unpack_first());
3554    }
3555
3556    #[mz_ore::test]
3557    fn miri_test_arena_growth_keeps_references() {
3558        // References returned by `push_bytes` must stay valid as later pushes allocate new
3559        // regions; this exercises the "never resize a region that holds data" invariant.
3560        let arena = RowArena::new();
3561        let chunks: Vec<Vec<u8>> = (0..128u16)
3562            .map(|i| vec![u8::try_from(i % 256).unwrap(); usize::from(i % 13) + 1])
3563            .collect();
3564        let refs: Vec<&[u8]> = chunks
3565            .iter()
3566            .map(|c| arena.push_bytes(c.as_slice()))
3567            .collect();
3568        for (i, r) in refs.iter().enumerate() {
3569            assert_eq!(*r, chunks[i].as_slice());
3570        }
3571    }
3572
3573    #[mz_ore::test]
3574    fn miri_test_arena_unary_row_at_offset() {
3575        // A row pushed after other bytes lands at a non-zero offset within a region; reading it
3576        // back must not depend on the row starting at offset zero or on any alignment.
3577        let arena = RowArena::new();
3578        arena.reserve(4096);
3579        let _pad = arena.push_bytes(vec![0xAB; 5]);
3580        let row = Row::pack_slice(&[Datum::String("hello"), Datum::Int64(42), Datum::True]);
3581        assert_eq!(arena.push_unary_row(row.clone()), row.unpack_first());
3582    }
3583
3584    #[mz_ore::test]
3585    fn miri_test_arena_clear_reuse() {
3586        // After `clear` the arena retains a region and remains usable across cycles.
3587        let mut arena = RowArena::new();
3588        for i in 0..100u8 {
3589            let _ = arena.push_bytes(vec![i; 16]);
3590        }
3591        arena.clear();
3592        assert_eq!(arena.push_bytes(vec![7u8; 8]), &[7u8; 8]);
3593        assert_eq!(arena.push_string("after clear".to_owned()), "after clear");
3594        arena.clear();
3595        let empty: &[u8] = &[];
3596        assert_eq!(arena.push_bytes(Vec::<u8>::new()), empty);
3597    }
3598
3599    #[mz_ore::test]
3600    fn miri_test_arena_writer() {
3601        use std::io::Write;
3602
3603        let arena = RowArena::new();
3604
3605        // Build a value incrementally and commit it.
3606        let mut w = arena.writer();
3607        let mut expected = Vec::new();
3608        for i in 0..1000u16 {
3609            let byte = u8::try_from(i % 256).unwrap();
3610            w.push(byte);
3611            expected.push(byte);
3612            w.extend_from_slice(&[byte, byte]);
3613            expected.extend_from_slice(&[byte, byte]);
3614        }
3615        assert_eq!(w.as_slice(), expected.as_slice());
3616        assert_eq!(w.len(), expected.len());
3617        let first = w.finish();
3618        assert_eq!(first, expected.as_slice());
3619
3620        // A second writer reuses the scratch; its result is independent of the first, which stays
3621        // valid because `finish` copied it into the arena.
3622        let mut w2 = arena.writer();
3623        write!(w2, "hello").unwrap();
3624        let second = w2.finish();
3625        assert_eq!(second, b"hello");
3626        assert_eq!(first, expected.as_slice());
3627
3628        // An empty writer commits to an empty slice.
3629        let empty: &[u8] = &[];
3630        assert_eq!(arena.writer().finish(), empty);
3631
3632        // Abandoning a writer without finishing is fine; the next writer starts empty.
3633        {
3634            let mut w3 = arena.writer();
3635            w3.extend_from_slice(b"discarded");
3636        }
3637        assert_eq!(arena.writer().as_slice(), empty);
3638    }
3639
3640    #[mz_ore::test]
3641    fn miri_test_arena_writer_nested() {
3642        // Reentrancy: a writer obtained while another is still live must not panic (no `RefCell`
3643        // double-borrow) and must not disturb the outer writer. The nested writer just gets its own
3644        // buffer; the outer one keeps building independently.
3645        let arena = RowArena::new();
3646
3647        let mut outer = arena.writer();
3648        outer.extend_from_slice(b"outer-before-");
3649
3650        // Take a second writer while `outer` is still live -- the case that double-borrowed before.
3651        let inner_bytes = {
3652            let mut inner = arena.writer();
3653            inner.extend_from_slice(b"inner");
3654            // The outer writer is unaffected by the nested one.
3655            assert_eq!(outer.as_slice(), b"outer-before-");
3656            inner.finish()
3657        };
3658        assert_eq!(inner_bytes, b"inner");
3659
3660        // `outer` is intact and still writable after the nested writer committed.
3661        outer.extend_from_slice(b"after");
3662        let outer_bytes = outer.finish();
3663        assert_eq!(outer_bytes, b"outer-before-after");
3664        // Both committed slices stay valid and independent.
3665        assert_eq!(inner_bytes, b"inner");
3666
3667        // Once all writers have dropped, the recycled buffer is reusable (and cleared on acquire).
3668        let mut again = arena.writer();
3669        again.extend_from_slice(b"reused");
3670        assert_eq!(again.finish(), b"reused");
3671    }
3672
3673    #[mz_ore::test]
3674    fn miri_test_arena_writer_fmt() {
3675        use std::fmt::Write;
3676
3677        // Format text into the writer (e.g. building a cast-to-string result) and commit as `&str`.
3678        let arena = RowArena::new();
3679        let mut w = arena.writer();
3680        for i in 0..5 {
3681            write!(w, "{i},").unwrap();
3682        }
3683        assert_eq!(w.finish_str(), "0,1,2,3,4,");
3684    }
3685
3686    #[mz_ore::test]
3687    fn miri_test_round_trip() {
3688        fn round_trip(datums: Vec<Datum>) {
3689            let row = Row::pack(datums.clone());
3690
3691            // When run under miri this catches undefined bytes written to data
3692            // eg by calling push_copy! on a type which contains undefined padding values
3693            println!("{:?}", row.data());
3694
3695            let datums2 = row.iter().collect::<Vec<_>>();
3696            let datums3 = row.unpack();
3697            assert_eq!(datums, datums2);
3698            assert_eq!(datums, datums3);
3699        }
3700
3701        round_trip(vec![]);
3702        round_trip(
3703            SqlScalarType::enumerate()
3704                .iter()
3705                .flat_map(|r#type| r#type.interesting_datums())
3706                .collect(),
3707        );
3708        round_trip(vec![
3709            Datum::Null,
3710            Datum::Null,
3711            Datum::False,
3712            Datum::True,
3713            Datum::Int16(-21),
3714            Datum::Int32(-42),
3715            Datum::Int64(-2_147_483_648 - 42),
3716            Datum::UInt8(0),
3717            Datum::UInt8(1),
3718            Datum::UInt16(0),
3719            Datum::UInt16(1),
3720            Datum::UInt16(1 << 8),
3721            Datum::UInt32(0),
3722            Datum::UInt32(1),
3723            Datum::UInt32(1 << 8),
3724            Datum::UInt32(1 << 16),
3725            Datum::UInt32(1 << 24),
3726            Datum::UInt64(0),
3727            Datum::UInt64(1),
3728            Datum::UInt64(1 << 8),
3729            Datum::UInt64(1 << 16),
3730            Datum::UInt64(1 << 24),
3731            Datum::UInt64(1 << 32),
3732            Datum::UInt64(1 << 40),
3733            Datum::UInt64(1 << 48),
3734            Datum::UInt64(1 << 56),
3735            Datum::Float32(OrderedFloat::from(-42.12)),
3736            Datum::Float64(OrderedFloat::from(-2_147_483_648.0 - 42.12)),
3737            Datum::Date(Date::from_pg_epoch(365 * 45 + 21).unwrap()),
3738            Datum::Timestamp(
3739                CheckedTimestamp::from_timestamplike(
3740                    NaiveDate::from_isoywd_opt(2019, 30, chrono::Weekday::Wed)
3741                        .unwrap()
3742                        .and_hms_opt(14, 32, 11)
3743                        .unwrap(),
3744                )
3745                .unwrap(),
3746            ),
3747            Datum::TimestampTz(
3748                CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(61, 0).unwrap())
3749                    .unwrap(),
3750            ),
3751            Datum::Interval(Interval {
3752                months: 312,
3753                ..Default::default()
3754            }),
3755            Datum::Interval(Interval::new(0, 0, 1_012_312)),
3756            Datum::Bytes(&[]),
3757            Datum::Bytes(&[0, 2, 1, 255]),
3758            Datum::String(""),
3759            Datum::String("العَرَبِيَّة"),
3760        ]);
3761    }
3762
3763    #[mz_ore::test]
3764    fn test_array() {
3765        // Construct an array using `Row::push_array` and verify that it unpacks
3766        // correctly.
3767        const DIM: ArrayDimension = ArrayDimension {
3768            lower_bound: 2,
3769            length: 2,
3770        };
3771        let mut row = Row::default();
3772        let mut packer = row.packer();
3773        packer
3774            .try_push_array(&[DIM], vec![Datum::Int32(1), Datum::Int32(2)])
3775            .unwrap();
3776        let arr1 = row.unpack_first().unwrap_array();
3777        assert_eq!(arr1.dims().into_iter().collect::<Vec<_>>(), vec![DIM]);
3778        assert_eq!(
3779            arr1.elements().into_iter().collect::<Vec<_>>(),
3780            vec![Datum::Int32(1), Datum::Int32(2)]
3781        );
3782
3783        // Pack a previously-constructed `Datum::Array` and verify that it
3784        // unpacks correctly.
3785        let row = Row::pack_slice(&[Datum::Array(arr1)]);
3786        let arr2 = row.unpack_first().unwrap_array();
3787        assert_eq!(arr1, arr2);
3788    }
3789
3790    #[mz_ore::test]
3791    fn test_multidimensional_array() {
3792        let datums = vec![
3793            Datum::Int32(1),
3794            Datum::Int32(2),
3795            Datum::Int32(3),
3796            Datum::Int32(4),
3797            Datum::Int32(5),
3798            Datum::Int32(6),
3799            Datum::Int32(7),
3800            Datum::Int32(8),
3801        ];
3802
3803        let mut row = Row::default();
3804        let mut packer = row.packer();
3805        packer
3806            .try_push_array(
3807                &[
3808                    ArrayDimension {
3809                        lower_bound: 1,
3810                        length: 1,
3811                    },
3812                    ArrayDimension {
3813                        lower_bound: 1,
3814                        length: 4,
3815                    },
3816                    ArrayDimension {
3817                        lower_bound: 1,
3818                        length: 2,
3819                    },
3820                ],
3821                &datums,
3822            )
3823            .unwrap();
3824        let array = row.unpack_first().unwrap_array();
3825        assert_eq!(array.elements().into_iter().collect::<Vec<_>>(), datums);
3826    }
3827
3828    #[mz_ore::test]
3829    fn test_array_max_dimensions() {
3830        let mut row = Row::default();
3831        let max_dims = usize::from(MAX_ARRAY_DIMENSIONS);
3832
3833        // An array with one too many dimensions should be rejected.
3834        let res = row.packer().try_push_array(
3835            &vec![
3836                ArrayDimension {
3837                    lower_bound: 1,
3838                    length: 1
3839                };
3840                max_dims + 1
3841            ],
3842            vec![Datum::Int32(4)],
3843        );
3844        assert_eq!(res, Err(InvalidArrayError::TooManyDimensions(max_dims + 1)));
3845        assert!(row.data.is_empty());
3846
3847        // An array with exactly the maximum allowable dimensions should be
3848        // accepted.
3849        row.packer()
3850            .try_push_array(
3851                &vec![
3852                    ArrayDimension {
3853                        lower_bound: 1,
3854                        length: 1
3855                    };
3856                    max_dims
3857                ],
3858                vec![Datum::Int32(4)],
3859            )
3860            .unwrap();
3861    }
3862
3863    #[mz_ore::test]
3864    fn test_array_wrong_cardinality() {
3865        let mut row = Row::default();
3866        let res = row.packer().try_push_array(
3867            &[
3868                ArrayDimension {
3869                    lower_bound: 1,
3870                    length: 2,
3871                },
3872                ArrayDimension {
3873                    lower_bound: 1,
3874                    length: 3,
3875                },
3876            ],
3877            vec![Datum::Int32(1), Datum::Int32(2)],
3878        );
3879        assert_eq!(
3880            res,
3881            Err(InvalidArrayError::WrongCardinality {
3882                actual: 2,
3883                expected: 6,
3884            })
3885        );
3886        assert!(row.data.is_empty());
3887    }
3888
3889    #[mz_ore::test]
3890    fn test_array_cardinality_overflow() {
3891        // Dimension lengths whose product overflows `usize` must be rejected as
3892        // a `WrongCardinality` error, not panic (under overflow checks) or wrap
3893        // (in release, which could spuriously accept a corrupt array). The
3894        // product saturates to `usize::MAX`, which no real element count matches.
3895        let mut row = Row::default();
3896        let res = row.packer().try_push_array(
3897            &[
3898                ArrayDimension {
3899                    lower_bound: 1,
3900                    length: usize::MAX,
3901                },
3902                ArrayDimension {
3903                    lower_bound: 1,
3904                    length: 2,
3905                },
3906            ],
3907            vec![Datum::Int32(1), Datum::Int32(2)],
3908        );
3909        assert_eq!(
3910            res,
3911            Err(InvalidArrayError::WrongCardinality {
3912                actual: 2,
3913                expected: usize::MAX,
3914            })
3915        );
3916        assert!(row.data.is_empty());
3917    }
3918
3919    #[mz_ore::test]
3920    fn test_nesting() {
3921        let mut row = Row::default();
3922        row.packer().push_dict_with(|row| {
3923            row.push(Datum::String("favourites"));
3924            row.push_list_with(|row| {
3925                row.push(Datum::String("ice cream"));
3926                row.push(Datum::String("oreos"));
3927                row.push(Datum::String("cheesecake"));
3928            });
3929            row.push(Datum::String("name"));
3930            row.push(Datum::String("bob"));
3931        });
3932
3933        let mut iter = row.unpack_first().unwrap_map().iter();
3934
3935        let (k, v) = iter.next().unwrap();
3936        assert_eq!(k, "favourites");
3937        assert_eq!(
3938            v.unwrap_list().iter().collect::<Vec<_>>(),
3939            vec![
3940                Datum::String("ice cream"),
3941                Datum::String("oreos"),
3942                Datum::String("cheesecake"),
3943            ]
3944        );
3945
3946        let (k, v) = iter.next().unwrap();
3947        assert_eq!(k, "name");
3948        assert_eq!(v, Datum::String("bob"));
3949    }
3950
3951    #[mz_ore::test]
3952    fn test_dict_errors() -> Result<(), Box<dyn std::error::Error>> {
3953        let pack = |ok| {
3954            let mut row = Row::default();
3955            row.packer().push_dict_with(|row| {
3956                if ok {
3957                    row.push(Datum::String("key"));
3958                    row.push(Datum::Int32(42));
3959                    Ok(7)
3960                } else {
3961                    Err("fail")
3962                }
3963            })?;
3964            Ok(row)
3965        };
3966
3967        assert_eq!(pack(false), Err("fail"));
3968
3969        let row = pack(true)?;
3970        let mut dict = row.unpack_first().unwrap_map().iter();
3971        assert_eq!(dict.next(), Some(("key", Datum::Int32(42))));
3972        assert_eq!(dict.next(), None);
3973
3974        Ok(())
3975    }
3976
3977    #[mz_ore::test]
3978    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `decNumberFromInt32` on OS `linux`
3979    fn test_datum_sizes() {
3980        let arena = RowArena::new();
3981
3982        // Test the claims about various datum sizes.
3983        let values_of_interest = vec![
3984            Datum::Null,
3985            Datum::False,
3986            Datum::Int16(0),
3987            Datum::Int32(0),
3988            Datum::Int64(0),
3989            Datum::UInt8(0),
3990            Datum::UInt8(1),
3991            Datum::UInt16(0),
3992            Datum::UInt16(1),
3993            Datum::UInt16(1 << 8),
3994            Datum::UInt32(0),
3995            Datum::UInt32(1),
3996            Datum::UInt32(1 << 8),
3997            Datum::UInt32(1 << 16),
3998            Datum::UInt32(1 << 24),
3999            Datum::UInt64(0),
4000            Datum::UInt64(1),
4001            Datum::UInt64(1 << 8),
4002            Datum::UInt64(1 << 16),
4003            Datum::UInt64(1 << 24),
4004            Datum::UInt64(1 << 32),
4005            Datum::UInt64(1 << 40),
4006            Datum::UInt64(1 << 48),
4007            Datum::UInt64(1 << 56),
4008            Datum::Float32(OrderedFloat(0.0)),
4009            Datum::Float64(OrderedFloat(0.0)),
4010            Datum::from(numeric::Numeric::from(0)),
4011            Datum::from(numeric::Numeric::from(1000)),
4012            Datum::from(numeric::Numeric::from(9999)),
4013            Datum::Date(
4014                NaiveDate::from_ymd_opt(1, 1, 1)
4015                    .unwrap()
4016                    .try_into()
4017                    .unwrap(),
4018            ),
4019            Datum::Timestamp(
4020                CheckedTimestamp::from_timestamplike(
4021                    DateTime::from_timestamp(0, 0).unwrap().naive_utc(),
4022                )
4023                .unwrap(),
4024            ),
4025            Datum::TimestampTz(
4026                CheckedTimestamp::from_timestamplike(DateTime::from_timestamp(0, 0).unwrap())
4027                    .unwrap(),
4028            ),
4029            Datum::Interval(Interval::default()),
4030            Datum::Bytes(&[]),
4031            Datum::String(""),
4032            Datum::JsonNull,
4033            Datum::Range(Range { inner: None }),
4034            arena.make_datum(|packer| {
4035                packer
4036                    .push_range(Range::new(Some((
4037                        RangeLowerBound::new(Datum::Int32(-1), true),
4038                        RangeUpperBound::new(Datum::Int32(1), true),
4039                    ))))
4040                    .unwrap();
4041            }),
4042        ];
4043        for value in values_of_interest {
4044            if datum_size(&value) != Row::pack_slice(&[value]).data.len() {
4045                panic!("Disparity in claimed size for {:?}", value);
4046            }
4047        }
4048    }
4049
4050    #[mz_ore::test]
4051    fn test_range_errors() {
4052        fn test_range_errors_inner<'a>(
4053            datums: Vec<Vec<Datum<'a>>>,
4054        ) -> Result<(), InvalidRangeError> {
4055            let mut row = Row::default();
4056            let row_len = row.byte_len();
4057            let mut packer = row.packer();
4058            let r = packer.push_range_with(
4059                RangeLowerBound {
4060                    inclusive: true,
4061                    bound: Some(|row: &mut RowPacker| {
4062                        for d in &datums[0] {
4063                            row.push(d);
4064                        }
4065                        Ok(())
4066                    }),
4067                },
4068                RangeUpperBound {
4069                    inclusive: true,
4070                    bound: Some(|row: &mut RowPacker| {
4071                        for d in &datums[1] {
4072                            row.push(d);
4073                        }
4074                        Ok(())
4075                    }),
4076                },
4077            );
4078
4079            assert_eq!(row_len, row.byte_len());
4080
4081            r
4082        }
4083
4084        // A finite bound whose closure pushes zero values violates the
4085        // `push_range_with` caller contract and still panics. This is
4086        // unreachable when decoding a `ProtoRow`: each decoded bound pushes
4087        // exactly one datum (or fails), so only an in-process caller can hit it.
4088        for panicking_case in [
4089            vec![vec![Datum::Int32(1)], vec![]],
4090            vec![vec![Datum::Int32(1), Datum::Int32(2)], vec![]],
4091        ] {
4092            #[allow(clippy::disallowed_methods)] // not using enhanced panic handler in tests
4093            let result = std::panic::catch_unwind(|| test_range_errors_inner(panicking_case));
4094            assert_err!(result);
4095        }
4096
4097        // Inconsistent bound counts, mismatched datum kinds, and Null bounds are
4098        // all reachable from a crafted/corrupted `ProtoRow`, so they return an
4099        // error instead of panicking.
4100        for error_case in [
4101            vec![
4102                vec![Datum::Int32(1), Datum::Int32(2)],
4103                vec![Datum::Int32(3)],
4104            ],
4105            vec![
4106                vec![Datum::Int32(1)],
4107                vec![Datum::Int32(2), Datum::Int32(3)],
4108            ],
4109            vec![vec![Datum::Int32(1)], vec![Datum::UInt16(2)]],
4110            vec![vec![Datum::Null], vec![Datum::Int32(2)]],
4111            vec![vec![Datum::Int32(1)], vec![Datum::Null]],
4112        ] {
4113            assert_eq!(
4114                test_range_errors_inner(error_case),
4115                Err(InvalidRangeError::InvalidRangeData)
4116            );
4117        }
4118
4119        let e = test_range_errors_inner(vec![vec![Datum::Int32(2)], vec![Datum::Int32(1)]]);
4120        assert_eq!(e, Err(InvalidRangeError::MisorderedRangeBounds));
4121    }
4122
4123    /// Lists have a variable-length encoding for their lengths. We test each case here.
4124    #[mz_ore::test]
4125    #[cfg_attr(miri, ignore)] // slow
4126    fn test_list_encoding() {
4127        fn test_list_encoding_inner(len: usize) {
4128            let list_elem = |i: usize| {
4129                if i % 2 == 0 {
4130                    Datum::False
4131                } else {
4132                    Datum::True
4133                }
4134            };
4135            let mut row = Row::default();
4136            {
4137                // Push some stuff.
4138                let mut packer = row.packer();
4139                packer.push(Datum::String("start"));
4140                packer.push_list_with(|packer| {
4141                    for i in 0..len {
4142                        packer.push(list_elem(i));
4143                    }
4144                });
4145                packer.push(Datum::String("end"));
4146            }
4147            // Check that we read back exactly what we pushed.
4148            let mut row_it = row.iter();
4149            assert_eq!(row_it.next().unwrap(), Datum::String("start"));
4150            match row_it.next().unwrap() {
4151                Datum::List(list) => {
4152                    let mut list_it = list.iter();
4153                    for i in 0..len {
4154                        assert_eq!(list_it.next().unwrap(), list_elem(i));
4155                    }
4156                    assert_none!(list_it.next());
4157                }
4158                _ => panic!("expected Datum::List"),
4159            }
4160            assert_eq!(row_it.next().unwrap(), Datum::String("end"));
4161            assert_none!(row_it.next());
4162        }
4163
4164        test_list_encoding_inner(0);
4165        test_list_encoding_inner(1);
4166        test_list_encoding_inner(10);
4167        test_list_encoding_inner(TINY - 1); // tiny
4168        test_list_encoding_inner(TINY + 1); // short
4169        test_list_encoding_inner(SHORT + 1); // long
4170
4171        // The biggest one takes 40 s on my laptop, probably not worth it.
4172        //test_list_encoding_inner(LONG + 1); // huge
4173    }
4174
4175    /// Demonstrates that DatumList's Eq (bytewise) and Ord (datum-by-datum) are now consistent.
4176    /// A list containing -0.0 and one containing +0.0 have different byte representations
4177    /// (IEEE 754 distinguishes them), originally Eq says they are not equal. But after
4178    /// using the new Datum::cmp, Eq says they are equal, which matches what Ord
4179    /// compares via iter().cmp(other.iter()), and them as equal.
4180    #[mz_ore::test]
4181    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4182    fn test_datum_list_eq_ord_consistency() {
4183        // Build list containing +0.0
4184        let mut row_pos = Row::default();
4185        row_pos.packer().push_list_with(|p| {
4186            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4187        });
4188        let list_pos = row_pos.unpack_first().unwrap_list();
4189
4190        // Build list containing -0.0 (distinct bit pattern from +0.0)
4191        let mut row_neg = Row::default();
4192        row_neg.packer().push_list_with(|p| {
4193            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4194        });
4195        let list_neg = row_neg.unpack_first().unwrap_list();
4196
4197        // Eq is bytewise: different encodings => not equal
4198        // This was a bug in the past, so we test it.
4199        assert_eq!(
4200            list_pos, list_neg,
4201            "Eq should see different encodings as equal"
4202        );
4203
4204        // Ord is datum-by-datum: -0.0 and +0.0 compare equal as Datums
4205        assert_eq!(
4206            list_pos.cmp(&list_neg),
4207            Ordering::Equal,
4208            "Ord (datum-by-datum) should see -0.0 and +0.0 as equal"
4209        );
4210    }
4211
4212    /// Demonstrates that DatumMap's derived Eq (bytewise) can make maps with equal keys and
4213    /// values compare equal when values have different encodings (e.g. -0.0 vs +0.0).
4214    #[mz_ore::test]
4215    fn test_datum_map_eq_bytewise_consistency() {
4216        // Build map {"k": +0.0}
4217        let mut row_pos = Row::default();
4218        row_pos.packer().push_dict_with(|p| {
4219            p.push(Datum::String("k"));
4220            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4221        });
4222        let map_pos = row_pos.unpack_first().unwrap_map();
4223
4224        // Build map {"k": -0.0}
4225        let mut row_neg = Row::default();
4226        row_neg.packer().push_dict_with(|p| {
4227            p.push(Datum::String("k"));
4228            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4229        });
4230        let map_neg = row_neg.unpack_first().unwrap_map();
4231
4232        // Same keys and semantically equal values, but Eq (bytewise) says not equal
4233        assert_eq!(
4234            map_pos, map_neg,
4235            "DatumMap Eq is semantic; -0.0 and +0.0 have different encodings but are equal"
4236        );
4237        // Verify they have the same logical content
4238        let entries_pos: Vec<_> = map_pos.iter().collect();
4239        let entries_neg: Vec<_> = map_neg.iter().collect();
4240        assert_eq!(entries_pos.len(), entries_neg.len());
4241        for ((k1, v1), (k2, v2)) in entries_pos.iter().zip_eq(entries_neg.iter()) {
4242            assert_eq!(k1, k2);
4243            assert_eq!(
4244                v1, v2,
4245                "Datum-level comparison treats -0.0 and +0.0 as equal"
4246            );
4247        }
4248    }
4249
4250    /// Hash must agree with Eq: equal lists must have the same hash.
4251    #[mz_ore::test]
4252    fn test_datum_list_hash_consistency() {
4253        // Equal lists (including -0.0 vs +0.0) must hash the same
4254        let mut row_pos = Row::default();
4255        row_pos.packer().push_list_with(|p| {
4256            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4257        });
4258        let list_pos = row_pos.unpack_first().unwrap_list();
4259
4260        let mut row_neg = Row::default();
4261        row_neg.packer().push_list_with(|p| {
4262            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4263        });
4264        let list_neg = row_neg.unpack_first().unwrap_list();
4265
4266        assert_eq!(list_pos, list_neg);
4267        assert_eq!(
4268            hash(&list_pos),
4269            hash(&list_neg),
4270            "equal lists must have same hash"
4271        );
4272
4273        // Unequal lists should have different hashes (with asymptotic probability 1)
4274        let mut row_a = Row::default();
4275        row_a.packer().push_list_with(|p| {
4276            p.push(Datum::Int32(1));
4277            p.push(Datum::Int32(2));
4278        });
4279        let list_a = row_a.unpack_first().unwrap_list();
4280
4281        let mut row_b = Row::default();
4282        row_b.packer().push_list_with(|p| {
4283            p.push(Datum::Int32(1));
4284            p.push(Datum::Int32(3));
4285        });
4286        let list_b = row_b.unpack_first().unwrap_list();
4287
4288        assert_ne!(list_a, list_b);
4289        assert_ne!(
4290            hash(&list_a),
4291            hash(&list_b),
4292            "unequal lists must have different hashes"
4293        );
4294    }
4295
4296    /// Ord/PartialOrd for DatumList: less, equal, greater.
4297    #[mz_ore::test]
4298    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4299    fn test_datum_list_ordering() {
4300        let mut row_12 = Row::default();
4301        row_12.packer().push_list_with(|p| {
4302            p.push(Datum::Int32(1));
4303            p.push(Datum::Int32(2));
4304        });
4305        let list_12 = row_12.unpack_first().unwrap_list();
4306
4307        let mut row_13 = Row::default();
4308        row_13.packer().push_list_with(|p| {
4309            p.push(Datum::Int32(1));
4310            p.push(Datum::Int32(3));
4311        });
4312        let list_13 = row_13.unpack_first().unwrap_list();
4313
4314        let mut row_123 = Row::default();
4315        row_123.packer().push_list_with(|p| {
4316            p.push(Datum::Int32(1));
4317            p.push(Datum::Int32(2));
4318            p.push(Datum::Int32(3));
4319        });
4320        let list_123 = row_123.unpack_first().unwrap_list();
4321
4322        // [1, 2] < [1, 3] due to the second element being different
4323        assert_eq!(list_12.cmp(&list_13), Ordering::Less);
4324        assert_eq!(list_13.cmp(&list_12), Ordering::Greater);
4325        assert_eq!(list_12.cmp(&list_12), Ordering::Equal);
4326        // shorter prefix compares less
4327        assert_eq!(list_12.cmp(&list_123), Ordering::Less);
4328    }
4329
4330    /// Hash must agree with Eq: equal maps must have the same hash.
4331    #[mz_ore::test]
4332    fn test_datum_map_hash_consistency() {
4333        let mut row_pos = Row::default();
4334        row_pos.packer().push_dict_with(|p| {
4335            p.push(Datum::String("x"));
4336            p.push(Datum::Float64(OrderedFloat::from(0.0)));
4337        });
4338        let map_pos = row_pos.unpack_first().unwrap_map();
4339
4340        let mut row_neg = Row::default();
4341        row_neg.packer().push_dict_with(|p| {
4342            p.push(Datum::String("x"));
4343            p.push(Datum::Float64(OrderedFloat::from(-0.0)));
4344        });
4345        let map_neg = row_neg.unpack_first().unwrap_map();
4346
4347        assert_eq!(map_pos, map_neg);
4348        assert_eq!(
4349            hash(&map_pos),
4350            hash(&map_neg),
4351            "equal maps must have same hash"
4352        );
4353
4354        let mut row_a = Row::default();
4355        row_a.packer().push_dict_with(|p| {
4356            p.push(Datum::String("a"));
4357            p.push(Datum::Int32(1));
4358        });
4359        let map_a = row_a.unpack_first().unwrap_map();
4360
4361        let mut row_b = Row::default();
4362        row_b.packer().push_dict_with(|p| {
4363            p.push(Datum::String("a"));
4364            p.push(Datum::Int32(2));
4365        });
4366        let map_b = row_b.unpack_first().unwrap_map();
4367
4368        assert_ne!(map_a, map_b);
4369        assert_ne!(
4370            hash(&map_a),
4371            hash(&map_b),
4372            "unequal maps must have different hashes"
4373        );
4374    }
4375
4376    /// Ord/PartialOrd for DatumMap: less, equal, greater (by key then value).
4377    #[mz_ore::test]
4378    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4379    fn test_datum_map_ordering() {
4380        let mut row_a1 = Row::default();
4381        row_a1.packer().push_dict_with(|p| {
4382            p.push(Datum::String("a"));
4383            p.push(Datum::Int32(1));
4384        });
4385        let map_a1 = row_a1.unpack_first().unwrap_map();
4386
4387        let mut row_a2 = Row::default();
4388        row_a2.packer().push_dict_with(|p| {
4389            p.push(Datum::String("a"));
4390            p.push(Datum::Int32(2));
4391        });
4392        let map_a2 = row_a2.unpack_first().unwrap_map();
4393
4394        let mut row_b1 = Row::default();
4395        row_b1.packer().push_dict_with(|p| {
4396            p.push(Datum::String("b"));
4397            p.push(Datum::Int32(1));
4398        });
4399        let map_b1 = row_b1.unpack_first().unwrap_map();
4400
4401        assert_eq!(map_a1.cmp(&map_a2), Ordering::Less);
4402        assert_eq!(map_a2.cmp(&map_a1), Ordering::Greater);
4403        assert_eq!(map_a1.cmp(&map_a1), Ordering::Equal);
4404        assert_eq!(map_a1.cmp(&map_b1), Ordering::Less); // "a" < "b"
4405    }
4406
4407    /// Datum puts Null last in the enum so that nulls sort last (PostgreSQL default).
4408    /// This ordering is used when comparing DatumList/DatumMap (e.g. jsonb_agg tiebreaker).
4409    #[mz_ore::test]
4410    #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
4411    fn test_datum_list_and_map_null_sorts_last() {
4412        // DatumList: [1] < [null] so non-null sorts before null
4413        let mut row_list_1 = Row::default();
4414        row_list_1
4415            .packer()
4416            .push_list_with(|p| p.push(Datum::Int32(1)));
4417        let list_1 = row_list_1.unpack_first().unwrap_list();
4418
4419        let mut row_list_null = Row::default();
4420        row_list_null
4421            .packer()
4422            .push_list_with(|p| p.push(Datum::Null));
4423        let list_null = row_list_null.unpack_first().unwrap_list();
4424
4425        assert_eq!(list_1.cmp(&list_null), Ordering::Less);
4426        assert_eq!(list_null.cmp(&list_1), Ordering::Greater);
4427
4428        // DatumMap: {"k": 1} < {"k": null} so non-null sorts before null (same as jsonb_agg)
4429        let mut row_map_1 = Row::default();
4430        row_map_1.packer().push_dict_with(|p| {
4431            p.push(Datum::String("k"));
4432            p.push(Datum::Int32(1));
4433        });
4434        let map_1 = row_map_1.unpack_first().unwrap_map();
4435
4436        let mut row_map_null = Row::default();
4437        row_map_null.packer().push_dict_with(|p| {
4438            p.push(Datum::String("k"));
4439            p.push(Datum::Null);
4440        });
4441        let map_null = row_map_null.unpack_first().unwrap_map();
4442
4443        assert_eq!(map_1.cmp(&map_null), Ordering::Less);
4444        assert_eq!(map_null.cmp(&map_1), Ordering::Greater);
4445    }
4446}