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