diff --git a/Cargo.lock b/Cargo.lock index ceda6800e..4493ba502 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2149,6 +2149,14 @@ dependencies = [ "hmac", ] +[[package]] +name = "pbox_slice" +version = "0.0.0" +dependencies = [ + "pgrx", + "pgrx-tests", +] + [[package]] name = "pem-rfc7468" version = "1.0.0" diff --git a/pgrx-examples/pbox_slice/.gitignore b/pgrx-examples/pbox_slice/.gitignore new file mode 100644 index 000000000..d4de4f351 --- /dev/null +++ b/pgrx-examples/pbox_slice/.gitignore @@ -0,0 +1,7 @@ +.DS_Store +.idea/ +/target +*.iml +**/*.rs.bk +Cargo.lock +sql/pbox_slice-*.sql diff --git a/pgrx-examples/pbox_slice/Cargo.toml b/pgrx-examples/pbox_slice/Cargo.toml new file mode 100644 index 000000000..109edf354 --- /dev/null +++ b/pgrx-examples/pbox_slice/Cargo.toml @@ -0,0 +1,35 @@ +#LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +#LICENSE +#LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +#LICENSE +#LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +#LICENSE +#LICENSE All rights reserved. +#LICENSE +#LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. + +[package] +name = "pbox_slice" +version = "0.0.0" +edition.workspace = true +rust-version.workspace = true +publish = false + +[lib] +crate-type = ["cdylib"] + +[features] +default = ["pg13"] +pg13 = ["pgrx/pg13", "pgrx-tests/pg13"] +pg14 = ["pgrx/pg14", "pgrx-tests/pg14"] +pg15 = ["pgrx/pg15", "pgrx-tests/pg15"] +pg16 = ["pgrx/pg16", "pgrx-tests/pg16"] +pg17 = ["pgrx/pg17", "pgrx-tests/pg17"] +pg18 = ["pgrx/pg18", "pgrx-tests/pg18"] +pg_test = [] + +[dependencies] +pgrx = { path = "../../pgrx" } + +[dev-dependencies] +pgrx-tests = { path = "../../pgrx-tests" } diff --git a/pgrx-examples/pbox_slice/README.md b/pgrx-examples/pbox_slice/README.md new file mode 100644 index 000000000..0dda2a440 --- /dev/null +++ b/pgrx-examples/pbox_slice/README.md @@ -0,0 +1,72 @@ +# pbox_slice — `MemCx::alloc_varlena` shown in a real workload + +A teaching extension that demonstrates `MemCx::alloc_varlena` and `PBox` in code that mirrors a real use case rather than a synthetic benchmark. + +The shape is a **minimal vector-embedding store**: pack a SQL `real[]` into a 4-byte-per-dimension binary `bytea` (matching pgvector's binary convention), and provide the distance / arithmetic primitives a smallRAG or k-NN application would actually run. + +| Function | What it does | +|----------|-------------| +| `embedding_pack(real[]) -> bytea` | Pack a float array as binary embedding (4 bytes/dim). NULL → 0.0. | +| `embedding_dims(bytea) -> int` | Number of dimensions (or `-1` if layout is invalid). | +| `embedding_l2_squared(bytea, bytea) -> float8` | Squared Euclidean distance — skip the sqrt because nearest-neighbour ordering only needs the monotonic transform. | +| `embedding_dot(bytea, bytea) -> float8` | Dot product. For unit-length vectors this equals cosine similarity. | +| `embedding_cosine(bytea, bytea) -> float8` | Cosine similarity. NaN on dim mismatch or zero norm. | +| `embedding_mean(bytea, bytea) -> bytea` | Element-wise mean. Useful for centroids and ensembling. | +| `embedding_normalize(bytea) -> bytea` | L2-normalize to unit length. Lets downstream cosine become a plain dot product. | +| `embedding_seq_mean(bytea, int) -> bytea` | Mean-pool a packed sequence of N embeddings. **Demonstrates child-MemCx scratch + `'mcx` lifetime protection** via `PBox::from_iter_in`. | +| `build_bytea(int) -> bytea` | Smoke-test fixture for SQL-level checks; unrelated to embeddings. | + +## Example workflow + +```sql +CREATE TABLE documents ( + id bigserial PRIMARY KEY, + text text, + emb bytea -- packed embedding from external embedder +); + +-- Insert: external embedder returns real[]; we pack to compact binary. +INSERT INTO documents (text, emb) +VALUES ( + 'rust is fun', + embedding_pack(ARRAY[0.13, 0.42, 0.78, /* ... */]::real[]) +); + +-- Verify the encoding round-trips. +SELECT embedding_dims(emb) FROM documents LIMIT 1; + +-- Top-K nearest neighbours by L2 distance. +SELECT id, text, + embedding_l2_squared(emb, $query_emb) AS dist +FROM documents +ORDER BY dist +LIMIT 10; + +-- Pre-normalize so cosine = dot, run faster searches downstream. +UPDATE documents SET emb = embedding_normalize(emb); +``` + +## Why this matters for the API surface + +Every function above is a real workload an extension author would +write. They exercise the new APIs in the shapes that matter: + +- **`embedding_pack`, `_mean`, `_normalize`**: build a new bytea from + scratch with exact size known up-front. No intermediate `Vec`, + no `memcpy` from Rust heap into a varlena. The output goes straight + to Postgres. +- **`embedding_l2_squared`, `_dot`, `_cosine`, `_dims`**: read existing + byteas via `&[u8]` (zero-copy in) and produce a scalar. No + allocation needed; tests confirm dimension-mismatch / zero-norm + edge cases return NaN rather than crashing. + +For embeddings sized 384–1536 floats (≈1.5–6 KB each) over millions of +rows, eliminating one `Vec` allocation + memcpy per query pays off +visibly. That is the point of `MemCx::alloc_varlena`. + +## Build & test + +```bash +cd pgrx-examples/pbox_slice +cargo pgrx test pg18 +``` diff --git a/pgrx-examples/pbox_slice/pbox_slice.control b/pgrx-examples/pbox_slice/pbox_slice.control new file mode 100644 index 000000000..cb11d1aab --- /dev/null +++ b/pgrx-examples/pbox_slice/pbox_slice.control @@ -0,0 +1,5 @@ +comment = 'pbox_slice: demonstrates MemCx slice + varlena APIs' +default_version = '@CARGO_VERSION@' +module_pathname = '$libdir/pbox_slice' +relocatable = false +superuser = false diff --git a/pgrx-examples/pbox_slice/src/lib.rs b/pgrx-examples/pbox_slice/src/lib.rs new file mode 100644 index 000000000..5998edced --- /dev/null +++ b/pgrx-examples/pbox_slice/src/lib.rs @@ -0,0 +1,478 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//LICENSE All rights reserved. +//LICENSE +//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. + +//! `pbox_slice` — a *teaching* extension showing how `MemCx::alloc_varlena` +//! and `PBox` enable zero-copy bytea outputs in real workloads. +//! +//! The example shape is a *minimal vector-embedding store*: pack/unpack +//! `float[]` to a binary `bytea` representation (4 bytes per dimension, +//! native endian — same convention as pgvector's `vector` type) and +//! provide the distance / arithmetic primitives a small RAG or k-NN +//! search application would actually run. +//! +//! For production vector workloads use [pgvector]; this crate exists to +//! demonstrate the `MemCx` / `PBox` API surface in code that +//! mirrors a real use case rather than a synthetic benchmark. +//! +//! [pgvector]: https://github.com/pgvector/pgvector + +use pgrx::datum::varlena_buf::{RawVarlena, VarlenaBuf}; +use pgrx::memcx::MemCx; +use pgrx::palloc::PBox; +use pgrx::prelude::*; + +pgrx::pg_module_magic!(name, version); + +const F32_BYTES: usize = core::mem::size_of::(); + +/// Smoke-test fixture: build a bytea of the requested length filled with a `(i % 256)` byte pattern. Useful for SQL-level smoke tests that don't need a real embedding shape. +#[pg_extern] +fn build_bytea<'mcx>(len: i32, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let n = len.max(0) as usize; + let mut buf: VarlenaBuf<'_> = cx.alloc_varlena(n).expect("OOM"); + for (i, b) in buf.payload_mut().iter_mut().enumerate() { + *b = (i & 0xff) as u8; + } + buf.into_pbox() +} + +/// Pack a SQL `real[]` into a binary embedding (4 bytes per dimension, native endian). NULL elements are encoded as `0.0`. +/// +/// The output layout matches the convention used by pgvector's `vector` type and is what an external embedder service would typically write to disk for compact storage. +#[pg_extern] +fn embedding_pack<'mcx>(coords: Array<'mcx, f32>, cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let n = coords.len(); + let mut buf = cx.alloc_varlena(n.saturating_mul(F32_BYTES)).expect("OOM"); + let dst = buf.payload_mut(); + for (i, v) in coords.iter().enumerate() { + let bytes = v.unwrap_or(0.0).to_ne_bytes(); + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&bytes); + } + buf.into_pbox() +} + +/// Number of dimensions stored in a packed embedding. Returns `-1` if the byte length is not a multiple of `sizeof(f32)`, which signals a corrupted or wrong-typed payload. +#[pg_extern] +fn embedding_dims(emb: &[u8]) -> i32 { + if emb.len() % F32_BYTES != 0 { -1 } else { (emb.len() / F32_BYTES) as i32 } +} + +/// Squared L2 (Euclidean) distance. Skips the sqrt because ordering by distance only needs the monotonic transform — the standard optimization in nearest-neighbour search. +/// +/// Returns `NaN` if the two embeddings have different dimensions. +#[pg_extern] +fn embedding_l2_squared(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + + av.iter() + .zip(bv.iter()) + .map(|(&x, &y)| { + let d = (x - y) as f64; + d * d + }) + .sum() +} + +/// Dot product. For unit-length embeddings this equals cosine similarity. +/// Returns `NaN` on dimension mismatch. +#[pg_extern] +fn embedding_dot(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + av.iter().zip(bv.iter()).map(|(&x, &y)| x as f64 * y as f64).sum() +} + +/// Cosine similarity (dot product over product of magnitudes). Returns `NaN` on dimension mismatch or when either embedding has zero norm. +#[pg_extern] +fn embedding_cosine(a: &[u8], b: &[u8]) -> f64 { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return f64::NAN, + }; + let mut dot = 0.0; + let mut na = 0.0; + let mut nb = 0.0; + for (&x, &y) in av.iter().zip(bv.iter()) { + let xf = x as f64; + let yf = y as f64; + dot += xf * yf; + na += xf * xf; + nb += yf * yf; + } + if na == 0.0 || nb == 0.0 { f64::NAN } else { dot / (na.sqrt() * nb.sqrt()) } +} + +/// Element-wise mean of two equal-dimension embeddings. Useful for simple ensembling, smoothing, or computing centroids without a full aggregate. Returns a zero-length bytea on dimension mismatch. +#[pg_extern] +fn embedding_mean<'mcx>(a: &[u8], b: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let (av, bv) = match decode_pair(a, b) { + Some(p) => p, + None => return cx.alloc_varlena(0).expect("OOM").into_pbox(), + }; + let n = av.len(); + let mut buf = cx.alloc_varlena(n.saturating_mul(F32_BYTES)).expect("OOM"); + let dst = buf.payload_mut(); + for i in 0..n { + let m = (av[i] + bv[i]) * 0.5; + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&m.to_ne_bytes()); + } + buf.into_pbox() +} + +/// L2-normalize an embedding to unit length. Common preprocessing step that lets downstream cosine similarity be expressed as a plain dot product (faster). Returns a zero-length bytea if the input has zero norm or invalid layout. +#[pg_extern] +fn embedding_normalize<'mcx>(emb: &[u8], cx: &MemCx<'mcx>) -> PBox<'mcx, RawVarlena> { + let v = match decode(emb) { + Some(v) => v, + None => return cx.alloc_varlena(0).expect("OOM").into_pbox(), + }; + let norm: f32 = v.iter().map(|&x| (x as f64) * (x as f64)).sum::().sqrt() as f32; + if norm == 0.0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + let mut buf = cx.alloc_varlena(v.len().saturating_mul(F32_BYTES)).expect("OOM"); + let dst = buf.payload_mut(); + for (i, &x) in v.iter().enumerate() { + let n = x / norm; + dst[i * F32_BYTES..(i + 1) * F32_BYTES].copy_from_slice(&n.to_ne_bytes()); + } + buf.into_pbox() +} + +/// Decode a packed embedding into a `Vec`, or `None` on invalid layout. +/// Copies element-by-element via `from_ne_bytes` because the bytea payload is only 1-byte aligned, so a zero-copy `&[f32]` reinterpret would be unsound. +fn decode(emb: &[u8]) -> Option> { + if emb.len() % F32_BYTES != 0 { + return None; + } + let n = emb.len() / F32_BYTES; + let mut out = Vec::with_capacity(n); + for i in 0..n { + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&emb[i * F32_BYTES..(i + 1) * F32_BYTES]); + out.push(f32::from_ne_bytes(buf)); + } + Some(out) +} + +fn decode_pair(a: &[u8], b: &[u8]) -> Option<(Vec, Vec)> { + let av = decode(a)?; + let bv = decode(b)?; + if av.len() != bv.len() || av.is_empty() { + return None; + } + Some((av, bv)) +} + +/// # Memory-context discipline for zero-copy varlena builders +/// +/// 1. Final accumulator → `cx.alloc_varlena(...)` in the caller's MemCx; survives function return. +/// 2. Per-row scratch → `from_iter_in(.., child_cx)` against a child MemCx that we delete before returning. All scratch freed in O(1). +/// 3. The borrow checker enforces the split: the scratch `PBox<[f32]>` has lifetime `'child_mcx` (shorter than `'mcx`) and **cannot escape** into the return value. See the compile_fail block below. +/// +/// ```compile_fail +/// // Pseudocode — what Rust rejects when you try to leak scratch: +/// fn cannot_leak<'mcx>(cx: &pgrx::memcx::MemCx<'mcx>) +/// -> pgrx::palloc::PBox<'mcx, [f32]> +/// { +/// unsafe { +/// let child = pgrx::PgMemoryContexts::new("scratch"); +/// let _prev = pgrx::pg_sys::MemoryContextSwitchTo(child.value()); +/// let smuggled = pgrx::memcx::current_context(|child_cx| { +/// pgrx::palloc::PBox::from_iter_in([1.0f32].into_iter(), child_cx) +/// .unwrap() +/// }); +/// pgrx::pg_sys::MemoryContextSwitchTo(_prev); +/// smuggled // lifetime may not live long enough +/// } +/// } +/// ``` +#[pg_extern] +fn embedding_seq_mean<'mcx>( + sequence: &[u8], + dims: i32, + cx: &MemCx<'mcx>, +) -> PBox<'mcx, RawVarlena> { + let dims = dims.max(0) as usize; + let stride = dims.saturating_mul(F32_BYTES); + if stride == 0 || sequence.len() % stride != 0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + let n = sequence.len() / stride; + if n == 0 { + return cx.alloc_varlena(0).expect("OOM").into_pbox(); + } + + // Final accumulator — lives in caller's MemCx, returned to PG, alloc_varlena already zeroes the payload. + let mut acc_buf = cx.alloc_varlena(stride).expect("OOM"); + + // Scratch path — child memory context, deleted before return. + // + // Panic-safety: `child` (PgMemoryContexts::Owned) deletes the underlying context in its `Drop`, and the `CxRestoreGuard` below restores the previous `CurrentMemoryContext` in its `Drop`. The guard is declared AFTER `child`, so on unwind it drops first (LIFO) — first the parent CMC is restored, then the child context is deleted. Either way (normal return or panic) `CurrentMemoryContext` is left exactly as we found it. + struct CxRestoreGuard { + previous: pg_sys::MemoryContext, + } + impl Drop for CxRestoreGuard { + fn drop(&mut self) { + // SAFETY: `previous` was the value of CurrentMemoryContext at the time we switched away; restoring it cannot fail. + unsafe { + pg_sys::MemoryContextSwitchTo(self.previous); + } + } + } + + let child = pgrx::PgMemoryContexts::new("embedding_seq_mean_scratch"); + let _restore = unsafe { + let previous = pg_sys::MemoryContextSwitchTo(child.value()); + CxRestoreGuard { previous } + }; + + pgrx::memcx::current_context(|child_cx| { + let inv_n = 1.0f32 / (n as f32); + for i in 0..n { + let row = &sequence[i * stride..(i + 1) * stride]; + + // from_iter_in: build a PBox<[f32]> in child_cx straight from an ExactSizeIterator. No MaybeUninit wrangling, no raw pointer writes. + // This PBox cannot outlive child_cx — borrow checker blocks any attempt to return it from the closure. + let scratch: PBox<[f32]> = PBox::from_iter_in( + (0..dims).map(|d| { + let off = d * F32_BYTES; + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&row[off..off + F32_BYTES]); + f32::from_ne_bytes(buf) + }), + child_cx, + ) + .expect("OOM"); + + // Borrow scratch (Deref<[f32]>) and accumulate into the output buffer (which lives in the OUTER cx). + let acc_bytes = acc_buf.payload_mut(); + for (d, &v) in scratch.iter().enumerate() { + let off = d * F32_BYTES; + let mut buf = [0u8; F32_BYTES]; + buf.copy_from_slice(&acc_bytes[off..off + F32_BYTES]); + let acc_v = f32::from_ne_bytes(buf) + v * inv_n; + acc_bytes[off..off + F32_BYTES].copy_from_slice(&acc_v.to_ne_bytes()); + } + // `scratch` falls out of scope here; its bytes stay in the child context until the bulk delete below. + } + }); + + acc_buf.into_pbox() +} + +#[cfg(any(test, feature = "pg_test"))] +#[pg_schema] +mod tests { + use pgrx::prelude::*; + + #[pg_test] + fn test_build_bytea_length() { + let r = Spi::get_one::("SELECT length(build_bytea(1024))"); + assert_eq!(r, Ok(Some(1024))); + } + + #[pg_test] + fn test_pack_dims() { + let r = Spi::get_one::( + "SELECT embedding_dims(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert_eq!(r, Ok(Some(3))); + } + + #[pg_test] + fn test_pack_byte_length() { + // 3 dims × 4 bytes = 12. + let r = Spi::get_one::("SELECT length(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))"); + assert_eq!(r, Ok(Some(12))); + } + + #[pg_test] + fn test_dims_invalid_layout() { + // 5 bytes is not a multiple of 4 — invalid embedding. + let r = Spi::get_one::(r"SELECT embedding_dims('\x0102030405'::bytea)"); + assert_eq!(r, Ok(Some(-1))); + } + + #[pg_test] + fn test_l2_self_distance_is_zero() { + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert_eq!(r, Ok(Some(0.0))); + } + + #[pg_test] + fn test_l2_known_distance() { + // ||(0,0) - (3,4)||^2 = 9 + 16 = 25 + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[0.0, 0.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[]))", + ); + assert_eq!(r, Ok(Some(25.0))); + } + + #[pg_test] + fn test_l2_dim_mismatch_is_nan() { + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + assert!(r.unwrap().unwrap().is_nan()); + } + + #[pg_test] + fn test_dot_product() { + // (1,2,3) · (4,5,6) = 4 + 10 + 18 = 32 + let r = Spi::get_one::( + "SELECT embedding_dot( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[4.0, 5.0, 6.0]::real[]))", + ); + assert_eq!(r, Ok(Some(32.0))); + } + + #[pg_test] + fn test_cosine_orthogonal_is_zero() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[1.0, 0.0]::real[]), + embedding_pack(ARRAY[0.0, 1.0]::real[]))", + ); + assert_eq!(r, Ok(Some(0.0))); + } + + #[pg_test] + fn test_cosine_parallel_is_one() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), + embedding_pack(ARRAY[2.0, 4.0, 6.0]::real[]))", + ); + // floating point: equal direction → 1.0 within fp precision. + let v = r.unwrap().unwrap(); + assert!((v - 1.0).abs() < 1e-6, "got {}", v); + } + + #[pg_test] + fn test_cosine_zero_vector_is_nan() { + let r = Spi::get_one::( + "SELECT embedding_cosine( + embedding_pack(ARRAY[0.0, 0.0]::real[]), + embedding_pack(ARRAY[1.0, 2.0]::real[]))", + ); + assert!(r.unwrap().unwrap().is_nan()); + } + + #[pg_test] + fn test_mean_simple() { + // mean of (1,2) and (3,4) is (2,3) → packed as 8 bytes + let r = Spi::get_one::( + "SELECT length(embedding_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[])))", + ); + assert_eq!(r, Ok(Some(8))); + } + + #[pg_test] + fn test_mean_value_check() { + // L2² of mean against expected (2,3) should be 0. + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]), + embedding_pack(ARRAY[3.0, 4.0]::real[])), + embedding_pack(ARRAY[2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); + } + + #[pg_test] + fn test_normalize_unit_length() { + // After normalize, dot(v,v) ≈ 1. + let r = Spi::get_one::( + "SELECT embedding_dot( + embedding_normalize(embedding_pack(ARRAY[3.0, 4.0]::real[])), + embedding_normalize(embedding_pack(ARRAY[3.0, 4.0]::real[])))", + ); + let v = r.unwrap().unwrap(); + assert!((v - 1.0).abs() < 1e-5, "got {}", v); + } + + #[pg_test] + fn test_normalize_zero_vector_returns_empty() { + let r = Spi::get_one::( + "SELECT length(embedding_normalize(embedding_pack(ARRAY[0.0, 0.0]::real[])))", + ); + assert_eq!(r, Ok(Some(0))); + } + + // ─── Memory-context demo tests ─────────────────────────────────────────── + + #[pg_test] + fn test_seq_mean_single_row_is_identity() { + // n=1: mean of one embedding equals itself. + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_seq_mean(embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]), 3), + embedding_pack(ARRAY[1.0, 2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); + } + + #[pg_test] + fn test_seq_mean_two_rows() { + // Concatenate (1,2) and (3,4); mean should be (2,3). + let r = Spi::get_one::( + "SELECT embedding_l2_squared( + embedding_seq_mean( + embedding_pack(ARRAY[1.0, 2.0]::real[]) + || embedding_pack(ARRAY[3.0, 4.0]::real[]), + 2), + embedding_pack(ARRAY[2.0, 3.0]::real[]))", + ); + let v = r.unwrap().unwrap(); + assert!(v < 1e-6, "expected ~0, got {}", v); + } + + #[pg_test] + fn test_seq_mean_invalid_dims() { + let r = Spi::get_one::( + "SELECT length(embedding_seq_mean(embedding_pack(ARRAY[1.0, 2.0]::real[]), 0))", + ); + assert_eq!(r, Ok(Some(0))); + } + + #[pg_test] + fn test_seq_mean_misaligned_input() { + let r = Spi::get_one::(r"SELECT length(embedding_seq_mean('\x010203'::bytea, 2))"); + assert_eq!(r, Ok(Some(0))); + } +} + +#[cfg(test)] +pub mod pg_test { + pub fn setup(_options: Vec<&str>) {} + pub fn postgresql_conf_options() -> Vec<&'static str> { + vec![] + } +} diff --git a/pgrx-unit-tests/src/tests/mod.rs b/pgrx-unit-tests/src/tests/mod.rs index a2da900f5..17c35bbe8 100644 --- a/pgrx-unit-tests/src/tests/mod.rs +++ b/pgrx-unit-tests/src/tests/mod.rs @@ -44,6 +44,7 @@ mod memcxt_tests; mod name_tests; mod numeric_tests; mod oid_tests; +mod pbox_slice_tests; mod pg_cast_tests; mod pg_extern_tests; mod pg_guard_tests; @@ -70,6 +71,7 @@ mod trigger_tests; mod type_ident_tests; mod uuid_tests; mod variadic_tests; +mod varlena_buf_tests; mod xact_callback_tests; mod xid64_tests; mod zero_datum_edge_cases; diff --git a/pgrx-unit-tests/src/tests/pbox_slice_tests.rs b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs new file mode 100644 index 000000000..4af6ee896 --- /dev/null +++ b/pgrx-unit-tests/src/tests/pbox_slice_tests.rs @@ -0,0 +1,358 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//LICENSE All rights reserved. +//LICENSE +//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use core::alloc::Layout; + use core::mem::MaybeUninit; + use pgrx::PgMemoryContexts; + use pgrx::memcx; + use pgrx::palloc::PBox; + use pgrx::prelude::*; + + /// Run `f` inside a freshly-created child memory context, then delete the + /// child and assert the parent's `MemoryContextMemAllocated` returned to + /// its baseline. If anything `f` allocated escaped the child context, + /// the parent's byte count grew and we fail. + /// + /// We use `PgMemoryContexts::new` (which wraps `AllocSetContextCreateExtended`) + /// and rely on its `Drop` impl to call `MemoryContextDelete`, so the child + /// context is freed even on unwind. The `recurse=false` argument to + /// `MemoryContextMemAllocated` measures only the parent's own blocks, + /// so child allocations don't pollute the baseline. + /// + /// Panic-safety: a `CxRestoreGuard` (Drop-based) restores the previous + /// `CurrentMemoryContext` even if `f` panics. The guard is declared + /// *after* `child` so on unwind it drops first (LIFO): CMC is restored, + /// then the child context is deleted. + /// + /// Caveat: `MemoryContextMemAllocated` reports the AllocSet's *block* + /// count, not byte-precise usage. Allocations that fit inside the + /// initial 8 KiB block (`ALLOCSET_DEFAULT_INITSIZE`) won't move the + /// counter, so a hypothetical leak smaller than that is invisible. + /// Tests that need leak detection should allocate well past one block + /// to force a fresh AllocSet block to be claimed. + fn with_scoped_memcx(f: F) + where + F: for<'mcx> FnOnce(&pgrx::memcx::MemCx<'mcx>), + { + struct CxRestoreGuard { + previous: pg_sys::MemoryContext, + } + impl Drop for CxRestoreGuard { + fn drop(&mut self) { + // SAFETY: `previous` was the value of CurrentMemoryContext at + // the time we switched away; restoring it cannot fail. + unsafe { + pg_sys::MemoryContextSwitchTo(self.previous); + } + } + } + + unsafe { + let parent = pg_sys::CurrentMemoryContext; + let before = pg_sys::MemoryContextMemAllocated(parent, false); + + let child = PgMemoryContexts::new("pbox_slice_test_scope"); + // Switch CMC and arm the restore guard. The guard is held in a + // named binding (NOT `let _ = ...`, which would drop immediately) + // so it lives until the end of the unsafe block, then drops in + // LIFO order before `child` — restoring CMC before the child + // context is deleted. `#[allow(unused_variables)]` silences the + // dead-binding lint since the guard's effect is in its Drop impl. + #[allow(unused_variables)] + let restore_guard = { + let previous = pg_sys::MemoryContextSwitchTo(child.value()); + CxRestoreGuard { previous } + }; + memcx::current_context(|cx| f(cx)); + + let after = pg_sys::MemoryContextMemAllocated(parent, false); + assert_eq!( + before, + after, + "leak: parent context grew by {} bytes after child deletion", + after.saturating_sub(before), + ); + } + } + + #[pg_test] + fn alloc_layout_returns_aligned_nonnull() { + memcx::current_context(|cx| { + let layout = Layout::from_size_align(128, 8).unwrap(); + let p = cx.alloc_layout(layout).expect("alloc"); + assert_eq!(p.as_ptr() as usize % 8, 0); + }); + } + + #[pg_test] + fn alloc_layout_zeroed_is_zero() { + memcx::current_context(|cx| { + let layout = Layout::from_size_align(64, 8).unwrap(); + let p = cx.alloc_layout_zeroed(layout).expect("alloc"); + let s = unsafe { core::slice::from_raw_parts(p.as_ptr(), 64) }; + assert!(s.iter().all(|&b| b == 0)); + }); + } + + #[pg_test] + fn new_uninit_slice_in_len_and_align() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = PBox::new_uninit_slice_in(16, cx).expect("alloc"); + assert_eq!(s.len(), 16); + assert_eq!(s.as_ptr() as *const u32 as usize % core::mem::align_of::(), 0); + }); + } + + #[pg_test] + fn new_zeroed_slice_in_is_zero() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = PBox::new_zeroed_slice_in(32, cx).expect("alloc"); + // SAFETY: zero is a valid bit pattern for u8. + let init: PBox<[u8]> = unsafe { s.assume_init() }; + assert_eq!(init.len(), 32); + assert!(init.iter().all(|&b| b == 0)); + }); + } + + #[pg_test] + fn from_slice_in_roundtrip() { + memcx::current_context(|cx| { + let src = [1u32, 2, 3, 4, 5]; + let dst: PBox<[u32]> = PBox::from_slice_in(&src, cx).expect("alloc"); + assert_eq!(&*dst, &src); + }); + } + + #[pg_test] + fn from_iter_in_happy_path() { + memcx::current_context(|cx| { + let v: Vec = (0..10).collect(); + let dst: PBox<[i32]> = PBox::from_iter_in(v.iter().copied(), cx).expect("alloc"); + assert_eq!(&*dst, v.as_slice()); + }); + } + + #[pg_test] + fn from_iter_in_under_yield_returns_err() { + // ExactSizeIterator promises 10, yields 3 — must return + // FromIterError::IteratorUnderYield (not OOM) and leave the buffer to the context (no panic, no UB). `T: Copy` so there is no Drop + // concern on the partial prefix. + struct ShortIter { + i: usize, + promised: usize, + real: usize, + } + impl Iterator for ShortIter { + type Item = u32; + fn next(&mut self) -> Option { + if self.i >= self.real { + return None; + } + let v = self.i as u32; + self.i += 1; + Some(v) + } + } + impl ExactSizeIterator for ShortIter { + fn len(&self) -> usize { + self.promised - self.i + } + } + memcx::current_context(|cx| { + let r = PBox::<[u32]>::from_iter_in(ShortIter { i: 0, promised: 10, real: 3 }, cx); + match r { + Err(pgrx::palloc::FromIterError::IteratorUnderYield { promised, yielded }) => { + assert_eq!(promised, 10); + assert_eq!(yielded, 3); + } + Err(pgrx::palloc::FromIterError::OutOfMemory) => { + panic!("expected IteratorUnderYield, got OutOfMemory") + } + Ok(_) => panic!("expected IteratorUnderYield, got Ok"), + } + }); + } + + #[pg_test] + fn from_iter_in_over_yield_truncates() { + struct OverIter { + i: u32, + promised: usize, + } + impl Iterator for OverIter { + type Item = u32; + fn next(&mut self) -> Option { + let v = self.i; + self.i += 1; + Some(v) // never returns None + } + } + impl ExactSizeIterator for OverIter { + fn len(&self) -> usize { + self.promised + } + } + memcx::current_context(|cx| { + let r = PBox::<[u32]>::from_iter_in(OverIter { i: 0, promised: 3 }, cx).expect("alloc"); + assert_eq!(&*r, &[0u32, 1, 2]); + }); + } + + #[pg_test] + fn zero_length_slice_is_safe() { + memcx::current_context(|cx| { + let s: PBox<[i64]> = PBox::from_slice_in(&[], cx).expect("alloc"); + assert_eq!(s.len(), 0); + assert!(s.is_empty()); + // Deref to an empty slice must not UB. + #[allow(clippy::explicit_auto_deref)] + let r: &[i64] = &*s; + assert_eq!(r.len(), 0); + }); + } + + #[pg_test] + fn new_uninit_slice_in_zero_length() { + // palloc(0) is implementation-defined across PG versions; the wrapper + // must still hand back a usable empty PBox without UB. + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = + PBox::new_uninit_slice_in(0, cx).expect("zero-length alloc"); + assert_eq!(s.len(), 0); + // SAFETY: zero elements means nothing to initialize; assume_init is trivially sound. + let init: PBox<[u64]> = unsafe { s.assume_init() }; + assert!(init.is_empty()); + let r: &[u64] = &*init; + assert_eq!(r.len(), 0); + }); + } + + #[pg_test] + fn new_zeroed_slice_in_zero_length() { + memcx::current_context(|cx| { + let s: PBox<[MaybeUninit]> = + PBox::new_zeroed_slice_in(0, cx).expect("zero-length alloc"); + assert_eq!(s.len(), 0); + // SAFETY: zero elements, nothing to read. + let init: PBox<[u32]> = unsafe { s.assume_init() }; + assert!(init.is_empty()); + }); + } + + #[pg_test] + fn oom_returns_err() { + memcx::current_context(|cx| { + // Request a length so large that `Layout::array` overflows + // (length * size > isize::MAX). PBox maps that to OutOfMemory + // before ever reaching palloc, so no ereport can occur. + let r = PBox::<[core::mem::MaybeUninit]>::new_uninit_slice_in(usize::MAX, cx); + assert!(r.is_err(), "expected OOM, got Ok"); + }); + } + + #[pg_test] + fn no_leak_after_context_delete() { + with_scoped_memcx(|cx| { + // Allocate a slice large enough to force the AllocSet to grab + // a fresh block (well past ALLOCSET_DEFAULT_INITSIZE = 8 KB), + // so a hypothetical leak would visibly grow the parent's + // MemoryContextMemAllocated. Intentionally do NOT pfree — + // the harness verifies that deleting the child reclaims it. + let _s: PBox<[u64]> = PBox::from_slice_in(&[1u64; 131_072], cx).expect("alloc"); + }); + } + + #[pg_test] + fn alloc_layout_rejects_over_max_alloc_size() { + memcx::current_context(|cx| { + // MaxAllocSize = 1 GiB - 1. Request 1 GiB + 1 byte through the + // non-huge path: must be rejected up front as OutOfMemory + // (not via PG ereport longjmp). + let layout = Layout::from_size_align((1usize << 30) + 1, 8).unwrap(); + assert!(cx.alloc_layout(layout).is_err()); + assert!(cx.alloc_layout_zeroed(layout).is_err()); + }); + } + + #[pg_test] + fn alloc_huge_layout_accepts_over_max_alloc_size() { + memcx::current_context(|cx| { + // 1.5 GiB — above MaxAllocSize, well below MaxAllocHugeSize. + // Succeeds on a machine with sufficient RAM; on a tight test + // environment we accept OutOfMemory as a valid response. + // The critical property: no longjmp, no panic — Rust gets + // back a Result either way. + let layout = Layout::from_size_align(3usize << 29, 8).unwrap(); + let _ = cx.alloc_huge_layout(layout); + }); + } + + #[pg_test] + fn alloc_huge_layout_zeroed_is_zero() { + memcx::current_context(|cx| { + // Small probe — verifies the ZERO flag still flows through + // when HUGE is set. No need to actually allocate huge here. + let layout = Layout::from_size_align(64, 8).unwrap(); + let p = cx.alloc_huge_layout_zeroed(layout).expect("alloc"); + let s = unsafe { core::slice::from_raw_parts(p.as_ptr(), 64) }; + assert!(s.iter().all(|&b| b == 0)); + }); + } + + #[pg_test] + fn new_uninit_slice_in_rejects_over_max_alloc_size() { + const MAX_ALLOC_SIZE: usize = 0x3fff_ffff; + let len = (MAX_ALLOC_SIZE / core::mem::size_of::()) + 1; + memcx::current_context(|cx| { + let r = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx); + assert!(r.is_err(), "expected OOM at MaxAllocSize boundary, got Ok"); + let r0 = PBox::<[MaybeUninit]>::new_zeroed_slice_in(len, cx); + assert!(r0.is_err(), "expected OOM at MaxAllocSize boundary (zeroed), got Ok"); + }); + } + + #[pg_test] + fn alloc_varlena_zero_payload() { + memcx::current_context(|cx| { + let mut buf = cx.alloc_varlena(0).expect("alloc"); + assert_eq!(buf.total_size(), pg_sys::VARHDRSZ); + assert_eq!(buf.payload().len(), 0); + assert_eq!(buf.payload_mut().len(), 0); + let ptr = buf.into_raw(); + // Round-trip: the empty payload must still be readable via inspect. + unsafe { + let view = cx.inspect_varlena(ptr); + assert_eq!(view.total_size(), pg_sys::VARHDRSZ); + assert_eq!(view.payload().len(), 0); + } + }); + } + + #[pg_test] + fn from_huge_slice_in_roundtrip() { + memcx::current_context(|cx| { + let src = [10u32, 20, 30, 40, 50]; + let dst: PBox<[u32]> = PBox::from_huge_slice_in(&src, cx).expect("alloc"); + assert_eq!(&*dst, &src); + }); + } + + #[pg_test] + fn from_huge_iter_in_roundtrip() { + memcx::current_context(|cx| { + let v: Vec = (100..110).collect(); + let dst: PBox<[i64]> = PBox::from_huge_iter_in(v.iter().copied(), cx).expect("alloc"); + assert_eq!(&*dst, v.as_slice()); + }); + } +} diff --git a/pgrx-unit-tests/src/tests/varlena_buf_tests.rs b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs new file mode 100644 index 000000000..c69d32f26 --- /dev/null +++ b/pgrx-unit-tests/src/tests/varlena_buf_tests.rs @@ -0,0 +1,140 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//LICENSE All rights reserved. +//LICENSE +//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::datum::varlena_buf::VarlenaBuf; + use pgrx::memcx; + use pgrx::pg_sys; + use pgrx::prelude::*; + + #[pg_test] + fn alloc_varlena_header_is_valid() { + memcx::current_context(|cx| { + let buf: VarlenaBuf = cx.alloc_varlena(64).expect("alloc"); + assert_eq!(buf.payload().len(), 64); + // payload should be zeroed + assert!(buf.payload().iter().all(|&b| b == 0)); + // total_size = header + payload + assert_eq!(buf.total_size(), 64 + pg_sys::VARHDRSZ); + // hand off to PG and check via varsize_any + let raw = buf.into_raw(); + unsafe { + assert_eq!(pgrx::varlena::varsize_any(raw), 64 + pg_sys::VARHDRSZ,); + } + }); + } + + #[pg_test] + fn write_then_read_payload() { + memcx::current_context(|cx| { + let mut buf: VarlenaBuf = cx.alloc_varlena(16).expect("alloc"); + for (i, b) in buf.payload_mut().iter_mut().enumerate() { + *b = i as u8; + } + for (i, &b) in buf.payload().iter().enumerate() { + assert_eq!(b as usize, i); + } + }); + } + + #[pg_test] + fn inspect_varlena_roundtrip() { + memcx::current_context(|cx| { + let mut buf: VarlenaBuf = cx.alloc_varlena(11).expect("alloc"); + buf.payload_mut().copy_from_slice(b"hello world"); + let ptr = buf.into_raw(); + // SAFETY: `ptr` came from alloc_varlena above, still in `cx`, header valid. + let view = unsafe { cx.inspect_varlena(ptr) }; + assert_eq!(view.payload(), b"hello world"); + }); + } + + /// Exercises `RawVarlena::header_size`'s short-header branch, which the 4-byte-emitting `alloc_varlena` never reaches on its own. We/ hand-craft a 1-byte (short) header buffer and verify `inspect_varlena` reads the payload boundary correctly. + #[pg_test] + fn inspect_varlena_handles_short_header() { + memcx::current_context(|cx| { + // Short header encodes total size (header + payload) in 1 byte: + // total fits in 7 bits => max 127 bytes total => max payload 126. + const PAYLOAD_LEN: usize = 5; + const TOTAL: usize = pg_sys::VARHDRSZ_SHORT + PAYLOAD_LEN; + + // Allocate raw bytes in cx, then set a 1-byte header by hand. + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(TOTAL, 1).unwrap()) + .expect("alloc"); + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + unsafe { + // SAFETY: `raw` is `TOTAL` bytes, writable, exclusively ours until we expose it via `inspect_varlena`. + pgrx::varlena::set_varsize_short(varlena_ptr, TOTAL as i32); + // Write a recognizable payload pattern past the 1-byte header. + let payload = raw.as_ptr().add(pg_sys::VARHDRSZ_SHORT); + for i in 0..PAYLOAD_LEN { + payload.add(i).write(0xA0 | (i as u8)); + } + } + + // SAFETY: header is now a valid in-line 1-byte header; pointee lives in `cx` so `'mcx`-bound borrow is sound. + let view = unsafe { cx.inspect_varlena(varlena_ptr) }; + assert_eq!(view.total_size(), TOTAL); + assert_eq!(view.payload().len(), PAYLOAD_LEN); + assert_eq!(view.payload(), &[0xA0, 0xA1, 0xA2, 0xA3, 0xA4]); + }); + } + + /// `inspect_varlena` must reject TOAST/external (1B_E) varlenas — they reference out-of-line storage that `RawVarlena::payload` cannot safely expose. We hand-craft the external header byte and confirm the assertion fires. + #[pg_test] + fn inspect_varlena_rejects_external() { + let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + memcx::current_context(|cx| { + // VARTAG_EXTERNAL byte: first byte == 0x01 marks "1B_E". + // We only need enough room for the tag byte; `inspect_varlena` panics before reading further bytes. + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(16, 4).unwrap()) + .expect("alloc"); + unsafe { + // 1B_E header byte: 0x01 on little-endian, 0x80 on big-endian. + #[cfg(target_endian = "little")] + let tag: u8 = 0x01; + #[cfg(target_endian = "big")] + let tag: u8 = 0x80; + raw.as_ptr().write(tag); + } + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + // SAFETY (panic expected): we deliberately violate the "in-line varlena" precondition to verify the runtime guard catches it. + let _ = unsafe { cx.inspect_varlena(varlena_ptr) }; + }); + })); + assert!(caught.is_err(), "expected panic on external varlena"); + } + + /// `inspect_varlena` must also reject 4-byte-headered *compressed* varlenas (VARATT_IS_4B_C). Same rationale: the payload bytes are notdirectly inspectable without decompression. + #[pg_test] + fn inspect_varlena_rejects_compressed() { + let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + memcx::current_context(|cx| { + let raw = cx + .alloc_layout_zeroed(core::alloc::Layout::from_size_align(16, 4).unwrap()) + .expect("alloc"); + unsafe { + #[cfg(target_endian = "little")] + let header_word: u32 = (16u32 << 2) | 0b10; + #[cfg(target_endian = "big")] + let header_word: u32 = 0x4000_0000 | 16u32; // top 2 bits = 01 + (raw.as_ptr() as *mut u32).write(header_word); + } + let varlena_ptr = raw.as_ptr() as *mut pg_sys::varlena; + let _ = unsafe { cx.inspect_varlena(varlena_ptr) }; + }); + })); + assert!(caught.is_err(), "expected panic on compressed varlena"); + } +} diff --git a/pgrx/src/datum/mod.rs b/pgrx/src/datum/mod.rs index 7dd3f38fa..bccdf6d4e 100644 --- a/pgrx/src/datum/mod.rs +++ b/pgrx/src/datum/mod.rs @@ -32,6 +32,7 @@ mod tuples; mod unbox; mod uuid; mod varlena; +pub mod varlena_buf; pub use self::uuid::*; pub use crate::datetime::support as datetime_support; @@ -50,6 +51,7 @@ pub use numeric::{AnyNumeric, Numeric}; pub use range::*; pub use unbox::*; pub use varlena::*; +pub use varlena_buf::{RawVarlena, VarlenaBuf}; use crate::memcx::MemCx; use crate::pg_sys; diff --git a/pgrx/src/datum/varlena_buf.rs b/pgrx/src/datum/varlena_buf.rs new file mode 100644 index 000000000..60fdf8fb7 --- /dev/null +++ b/pgrx/src/datum/varlena_buf.rs @@ -0,0 +1,202 @@ +//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC. +//LICENSE +//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc. +//LICENSE +//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. +//LICENSE +//LICENSE All rights reserved. +//LICENSE +//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file. + +//! Safe, lifetime-bound varlena allocation. +//! +//! [`MemCx::alloc_varlena`] returns a [`VarlenaBuf<'mcx>`] whose 4-byte header +//! is set to a valid total size before the wrapper exists. Writing into the +//! payload via [`VarlenaBuf::payload_mut`] cannot resize, so the header +//! invariant is preserved. +//! +//! # Examples +//! +//! Building a bytea and handing it to Postgres: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::pg_sys; +//! +//! memcx::current_context(|cx| { +//! let mut buf = cx.alloc_varlena(11).expect("OOM"); +//! buf.payload_mut().copy_from_slice(b"hello world"); +//! let ptr: *mut pg_sys::varlena = buf.into_raw(); +//! // `ptr` can be returned from a #[pg_extern] as a varlena Datum. +//! // The allocation lives in `cx` until the memory context resets. +//! let _ = ptr; +//! }); +//! ``` +//! +//! See: `docs/superpowers/specs/2026-06-15-memcx-slice-varlena-parity-design.md` + +use crate::callconv::{BoxRet, FcInfo}; +use crate::datum::Datum; +use crate::memcx::MemCx; +use crate::palloc::PBox; +use crate::pg_sys; +use core::ptr::NonNull; +use pgrx_sql_entity_graph::metadata::{ + ArgumentError, ReturnsError, ReturnsRef, SqlMappingRef, SqlTranslatable, +}; + +/// A `varlena` allocation with a valid in-line header (1-byte short or +/// 4-byte uncompressed), maintained as a type invariant. +/// +/// Safely obtainable only through [`MemCx::alloc_varlena`] (always emits a +/// 4-byte header) or [`MemCx::inspect_varlena`] (caller-asserted; rejects +/// TOAST/compressed varlenas at runtime). +/// +/// # Layout +/// +/// `#[repr(transparent)]` over `[u8]`. Header sits at byte offset 0; payload +/// starts at offset 1 (short header) or 4 (full header). The total byte +/// length matches what `varsize_any` returns for the header. +#[repr(transparent)] +pub struct RawVarlena { + bytes: [u8], +} + +impl RawVarlena { + /// Total size in bytes (header + payload), as recorded in the header. + pub fn total_size(&self) -> usize { + // SAFETY: type invariant — header is a valid in-line varlena header(4-byte uncompressed or 1-byte short; TOAST/compressed rejected at construction). + unsafe { crate::varlena::varsize_any(self.as_ptr()) } + } + + /// Byte offset of the payload from the start of the allocation, derived from the header type (1 for short header, 4 for full header). + #[inline] + fn header_size(&self) -> usize { + // SAFETY: invariant — header is in-line; non-1B_E. + if unsafe { crate::varlena::varatt_is_1b(self.as_ptr()) } { + pg_sys::VARHDRSZ_SHORT + } else { + pg_sys::VARHDRSZ + } + } + + /// Read-only access to the payload bytes (excludes the header). + pub fn payload(&self) -> &[u8] { + let hdr = self.header_size(); + let total = self.total_size(); + + let payload_len = + total.checked_sub(hdr).expect("RawVarlena invariant: total_size >= header_size"); + // SAFETY: payload starts at `hdr` bytes past the allocation start; the type invariant guarantees the allocation is at least `total` bytes. + unsafe { + core::slice::from_raw_parts((self as *const Self as *const u8).add(hdr), payload_len) + } + } + + /// Mutable access to the payload bytes. Cannot be used to grow or shrink: + /// the header (and therefore `total_size`) is unchanged. + pub fn payload_mut(&mut self) -> &mut [u8] { + let hdr = self.header_size(); + let total = self.total_size(); + let payload_len = + total.checked_sub(hdr).expect("RawVarlena invariant: total_size >= header_size"); + // SAFETY: as `payload`; exclusive access via `&mut self`. + unsafe { + core::slice::from_raw_parts_mut((self as *mut Self as *mut u8).add(hdr), payload_len) + } + } + + /// Raw `*const pg_sys::varlena` pointer. The pointee remains valid for + /// the lifetime of `&self`. + pub fn as_ptr(&self) -> *const pg_sys::varlena { + self as *const Self as *const pg_sys::varlena + } + + /// Raw `*mut pg_sys::varlena` pointer. + pub fn as_mut_ptr(&mut self) -> *mut pg_sys::varlena { + self as *mut Self as *mut pg_sys::varlena + } +} + +/// A varlena buffer: a `PBox`-owned `RawVarlena`, scoped to a memory context. +pub struct VarlenaBuf<'mcx> { + inner: PBox<'mcx, RawVarlena>, +} + +impl<'mcx> VarlenaBuf<'mcx> { + /// # Safety + /// `data_ptr` must point to a valid varlena allocation in `cx` whose header has + /// already been set, and the total byte length must equal `total`. + pub(crate) unsafe fn from_raw(data_ptr: NonNull, total: usize, cx: &MemCx<'mcx>) -> Self { + // Build a *mut [u8] fat pointer of the right length, then cast to + // *mut RawVarlena (sound: RawVarlena is #[repr(transparent)] over [u8]). + let slice_ptr: *mut [u8] = core::ptr::slice_from_raw_parts_mut(data_ptr.as_ptr(), total); + let raw_ptr: *mut RawVarlena = slice_ptr as *mut RawVarlena; + // SAFETY: data_ptr was non-null, so the fat pointer's data pointer is non-null. + let nn = NonNull::new_unchecked(raw_ptr); + let inner = PBox::from_raw_in(nn, cx); + VarlenaBuf { inner } + } + + /// Read-only payload access. + pub fn payload(&self) -> &[u8] { + // SAFETY: `inner.ptr` is a valid fat pointer to a RawVarlena whose header + // is valid by construction. + unsafe { (&*self.inner.as_ptr()).payload() } + } + + /// Mutable payload access. + pub fn payload_mut(&mut self) -> &mut [u8] { + // SAFETY: `inner` exclusively owns the allocation; we hold `&mut self`. + unsafe { (&mut *self.inner.as_mut_ptr()).payload_mut() } + } + + /// Total size (header + payload) as recorded in the header. + pub fn total_size(&self) -> usize { + // SAFETY: same as `payload`. + unsafe { (&*self.inner.as_ptr()).total_size() } + } + + /// Convert into the underlying `PBox`. Ownership of the + /// memory-context-bound allocation is preserved. + pub fn into_pbox(self) -> PBox<'mcx, RawVarlena> { + self.inner + } + + /// Leak the `PBox` and return a raw `*mut pg_sys::varlena` for handoff + /// to Postgres (e.g., as a function return). The allocation remains in + /// `cx` and is reclaimed at context reset. + pub fn into_raw(self) -> *mut pg_sys::varlena { + let mut me = core::mem::ManuallyDrop::new(self); + me.inner.as_mut_ptr() as *mut pg_sys::varlena + } +} + +// --- Returning `PBox` directly from a `#[pg_extern]` ----------- +// +// `RawVarlena` deliberately does NOT implement `BorrowDatum` (it has no fixed in-Datum representation: a varlena is always pass-by-reference and its size is header-driven). The blanket `BoxRet`/`SqlTranslatable` impls for `PBox` in `palloc/pbox.rs` therefore do not cover it. +/// +/// Returning a `PBox<'mcx, RawVarlena>` from a `#[pg_extern]` hands the +/// underlying varlena pointer to Postgres unchanged. No copy: the +/// allocation stays in the caller's `MemCx` and is reclaimed on context +/// reset. +unsafe impl<'mcx> BoxRet for PBox<'mcx, RawVarlena> { + unsafe fn box_into<'fcx>(self, fcinfo: &mut FcInfo<'fcx>) -> Datum<'fcx> { + // Leak the PBox; ownership of the allocation transfers to Postgres. + let mut me = core::mem::ManuallyDrop::new(self); + let ptr = me.as_mut_ptr() as *mut pg_sys::varlena; + // SAFETY: `ptr` points at a valid varlena allocation in the caller's + // memory context; Postgres owns it from here on. + unsafe { fcinfo.return_raw_datum(pg_sys::Datum::from(ptr)) } + } +} + +/// `PBox<'mcx, RawVarlena>` maps to SQL `bytea`. +unsafe impl<'mcx> SqlTranslatable for PBox<'mcx, RawVarlena> { + const TYPE_IDENT: &'static str = "bytea"; + const TYPE_ORIGIN: pgrx_sql_entity_graph::metadata::TypeOrigin = + pgrx_sql_entity_graph::metadata::TypeOrigin::External; + const ARGUMENT_SQL: Result = Ok(SqlMappingRef::literal("bytea")); + const RETURN_SQL: Result = + Ok(ReturnsRef::One(SqlMappingRef::literal("bytea"))); +} diff --git a/pgrx/src/memcx.rs b/pgrx/src/memcx.rs index d2fd5177f..ee6795857 100644 --- a/pgrx/src/memcx.rs +++ b/pgrx/src/memcx.rs @@ -13,8 +13,18 @@ use pgrx_sql_entity_graph::metadata::{ use crate::callconv::{Arg, ArgAbi}; use crate::nullable::Nullable; use crate::pg_sys; +use core::alloc::Layout; use core::{marker::PhantomData, ptr::NonNull}; +/// Postgres's `MaxAllocSize` (1 GiB - 1). `bindgen` does not expose it +/// (see `pgrx/src/array.rs:219`), so we mirror the definition here. Any +/// `palloc` call exceeding this size requires `MCXT_ALLOC_HUGE`; without +/// that flag, Postgres `ereport`s — which longjmps out of Rust frames. +const MAX_ALLOC_SIZE: usize = 0x3fff_ffff; + +/// Postgres's `MaxAllocHugeSize` (`SIZE_MAX / 2`). `bindgen` records this as`0` because the C macro depends on `SIZE_MAX`; we mirror the real value so the huge-allocation entrypoints can reject sizes that would otherwise trigger Postgres's `MemoryContextCheckSize` ereport and longjmp out of Rust frames. +const MAX_ALLOC_HUGE_SIZE: usize = usize::MAX / 2; + /// A borrowed memory context. #[repr(transparent)] pub struct MemCx<'mcx> { @@ -50,6 +60,102 @@ impl<'mcx> MemCx<'mcx> { NonNull::new(ptr.cast()).ok_or(OutOfMemory::new()) } + /// Allocate a buffer matching `layout`. + /// + /// Sizes above Postgres's `MaxAllocSize` (1 GiB - 1) are rejected up + /// front with `OutOfMemory`. Without this guard, the underlying palloc + /// call would `ereport` and longjmp out of the Rust frame. For genuine + /// huge allocations, use [`MemCx::alloc_huge_layout`] instead. + /// + /// On PG >= 16 this uses `MemoryContextAllocAligned` for arbitrary power-of-two alignment. On earlier versions this falls back to size-only `MemoryContextAllocExtended`, which only guarantees `MAXIMUM_ALIGNOF` (8 on 64-bit, 4 on 32-bit); requests for stronger alignment are rejected with `OutOfMemory` rather than silently returning an under-aligned pointer. + pub fn alloc_layout(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_SIZE)?; + self.alloc_layout_with_flags(layout, pg_sys::MCXT_ALLOC_NO_OOM as i32) + } + + /// Like [`MemCx::alloc_layout`] but zeroes the allocation. + pub fn alloc_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_SIZE)?; + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_ZERO) as i32, + ) + } + + /// Like [`MemCx::alloc_layout`] but permits sizes above `MaxAllocSize` + /// (1 GiB - 1), up to `MaxAllocHugeSize` (`SIZE_MAX / 2`). Sizes beyond + /// that bound are rejected up front with `OutOfMemory`; without that + /// guard Postgres's `MemoryContextCheckSize` would `ereport` (which + /// `MCXT_ALLOC_NO_OOM` does **not** suppress) and longjmp out of the + /// Rust frame. + /// + /// Use only when the allocation is genuinely known to be huge. The 1 GiB + /// cap on [`MemCx::alloc_layout`] is a deliberate sanity check that + /// catches sizing bugs early; opting into `_huge_` should be a + /// considered choice, not a routine workaround. + /// + /// Note: varlena values cannot exceed 1 GiB by design (the 4-byte + /// header encodes 30-bit size). Huge allocations are only useful for + /// Rust-side scratch buffers, not for values returned to Postgres. + pub fn alloc_huge_layout(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_HUGE_SIZE)?; + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE) as i32, + ) + } + + /// Like [`MemCx::alloc_huge_layout`] but zeroes the allocation. + pub fn alloc_huge_layout_zeroed(&self, layout: Layout) -> Result, OutOfMemory> { + Self::check_layout(layout, MAX_ALLOC_HUGE_SIZE)?; + self.alloc_layout_with_flags( + layout, + (pg_sys::MCXT_ALLOC_NO_OOM | pg_sys::MCXT_ALLOC_HUGE | pg_sys::MCXT_ALLOC_ZERO) as i32, + ) + } + + /// Validate a `Layout` against the per-context size cap and the alignment guarantees actually delivered by the underlying palloc call. + /// + /// On PG >= 16 `MemoryContextAllocAligned` honors arbitrary power-of-two alignments. On older versions the implementation falls back to size-only `MemoryContextAllocExtended`, which only guarantees`MAXIMUM_ALIGNOF`. Rejecting over-aligned requests there prevents safe callers from receiving an under-aligned `NonNull` that silently violates the requested `Layout`. + #[inline] + fn check_layout(layout: Layout, max_size: usize) -> Result<(), OutOfMemory> { + if layout.size() > max_size { + return Err(OutOfMemory::new()); + } + #[cfg(not(any(feature = "pg16", feature = "pg17", feature = "pg18")))] + { + if layout.align() > pg_sys::MAXIMUM_ALIGNOF as usize { + return Err(OutOfMemory::new()); + } + } + Ok(()) + } + + #[inline] + fn alloc_layout_with_flags( + &self, + layout: Layout, + flags: i32, + ) -> Result, OutOfMemory> { + let ptr = unsafe { + #[cfg(any(feature = "pg16", feature = "pg17", feature = "pg18"))] + { + pg_sys::MemoryContextAllocAligned( + self.ptr.as_ptr(), + layout.size(), + layout.align(), + flags, + ) + } + #[cfg(not(any(feature = "pg16", feature = "pg17", feature = "pg18")))] + { + let _ = layout.align(); + pg_sys::MemoryContextAllocExtended(self.ptr.as_ptr(), layout.size(), flags) + } + }; + NonNull::new(ptr.cast()).ok_or_else(OutOfMemory::new) + } + /// Stores the current memory context, switches to *this* memory context, /// and executes the closure `f`. /// Once `f` completes, the previous current memory context is restored. @@ -66,6 +172,76 @@ impl<'mcx> MemCx<'mcx> { pg_sys::MemoryContextSwitchTo(remembered); res } + + /// Allocate a varlena with a valid 4-byte header and `payload_len` zeroed payload bytes. The returned [`VarlenaBuf`] maintains the header as a type invariant. + /// + /// [`VarlenaBuf`]: crate::datum::varlena_buf::VarlenaBuf + pub fn alloc_varlena( + &self, + payload_len: usize, + ) -> Result, OutOfMemory> { + let total = (pg_sys::VARHDRSZ).checked_add(payload_len).ok_or_else(OutOfMemory::new)?; + let layout = core::alloc::Layout::from_size_align(total, core::mem::align_of::()) + .map_err(|_| OutOfMemory::new())?; + let raw = self.alloc_layout_zeroed(layout)?; + // Set the 4-byte header. The total_size encoded in the header is `total` (header bytes + payload bytes), per Postgres convention for 4-byte headers. + // SAFETY: `raw` points at `total` zeroed bytes, large enough for a 4-byte header. + unsafe { + crate::varlena::set_varsize_4b(raw.as_ptr() as *mut pg_sys::varlena, total as i32); + } + // SAFETY: `raw` is non-null, sized at `total`, header is now valid. + Ok(unsafe { crate::datum::varlena_buf::VarlenaBuf::from_raw(raw, total, self) }) + } + + /// Inspect an existing varlena pointer as a `&RawVarlena`. + /// + /// The returned reference is bound to `&self`, not to the full `'mcx` + /// lifetime. This narrows the borrow so that other handles in the same + /// `MemCx` (e.g. a `PBox` returned by + /// `VarlenaBuf::into_pbox`) cannot be re-borrowed as `&mut` while the + /// `&RawVarlena` is alive. + /// + /// Lifetime caveat: the `&'a self` tie does **not** prove the pointee + /// outlives `'a`. The pointer is raw and originates outside the type + /// system. Use-after-free protection relies entirely on the unsafe + /// preconditions below — specifically that `*ptr` lives for at least + /// `'mcx`. The lifetime is purely an aliasing knob, not a UAF guard. + /// + /// # Panics + /// Panics if `ptr` references an external/TOAST (`VARATT_IS_1B_E`) or + /// 4-byte compressed (`VARATT_IS_4B_C`) varlena. Both classes carry + /// out-of-line / compressed payloads that `RawVarlena::payload` + /// cannot soundly expose; detoast first (e.g. `pg_detoast_datum`). + /// Inside a `#[pg_extern]` the panic is converted to `ereport ERROR` + /// by pgrx's panic handler. + /// + /// # Safety + /// - `ptr` must point to a valid varlena whose header size is correctly set + /// - `ptr` must be 4-byte aligned when the header is a 4-byte header (uncompressed or compressed); `varsize_4b` reads a `u32` directly from `ptr`. 1-byte (short) headers carry no alignment requirement. + /// - `ptr` must point to an in-line (non-TOAST, non-compressed) varlena; + /// `RawVarlena::payload` can only expose bytes that live in the same + /// allocation. TOAST/compressed varlenas must be detoasted first + /// (e.g. via `pg_detoast_datum`). + /// - the pointee must live for at least `'mcx` + /// - no `&mut` reference to the same allocation may exist while the + /// returned `&RawVarlena` is alive + pub unsafe fn inspect_varlena<'a>( + &'a self, + ptr: *const pg_sys::varlena, + ) -> &'a crate::datum::varlena_buf::RawVarlena { + assert!( + !crate::varlena::varatt_is_1b_e(ptr), + "inspect_varlena: external/TOASTed varlena must be detoasted first", + ); + assert!( + !crate::varlena::varatt_is_4b_c(ptr), + "inspect_varlena: compressed varlena must be detoasted first", + ); + let total = crate::varlena::varsize_any(ptr); + let slice_ptr: *const [u8] = core::ptr::slice_from_raw_parts(ptr as *const u8, total); + let raw_ptr: *const crate::datum::varlena_buf::RawVarlena = slice_ptr as *const _; + &*raw_ptr + } } #[derive(Debug)] diff --git a/pgrx/src/palloc.rs b/pgrx/src/palloc.rs index 12950f766..07a4f7c1b 100644 --- a/pgrx/src/palloc.rs +++ b/pgrx/src/palloc.rs @@ -1,3 +1,3 @@ mod pbox; -pub use pbox::PBox; +pub use pbox::{FromIterError, PBox}; diff --git a/pgrx/src/palloc/pbox.rs b/pgrx/src/palloc/pbox.rs index 69e181d67..37649cd88 100644 --- a/pgrx/src/palloc/pbox.rs +++ b/pgrx/src/palloc/pbox.rs @@ -1,9 +1,48 @@ +//! Lifetime-bound boxed allocations in a Postgres memory context. +//! +//! [`PBox<'mcx, T>`] is the safe, lifetime-bound analogue of [`alloc::boxed::Box`] +//! for memory allocated inside a [`MemCx<'mcx>`]. Unlike `Box`, the value is +//! reclaimed when the memory context is reset — not when `PBox` is dropped. +//! +//! # Examples +//! +//! Allocating a zero-initialized slice and reading it back: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::palloc::PBox; +//! use core::mem::MaybeUninit; +//! +//! memcx::current_context(|cx| { +//! let s = PBox::<[MaybeUninit]>::new_zeroed_slice_in(8, cx) +//! .expect("OOM"); +//! // SAFETY: zero is a valid bit pattern for u32. +//! let s: PBox<[u32]> = unsafe { s.assume_init() }; +//! assert_eq!(s.len(), 8); +//! assert!(s.iter().all(|&n| n == 0)); +//! }); +//! ``` +//! +//! Copying from an existing slice: +//! +//! ```no_run +//! use pgrx::memcx; +//! use pgrx::palloc::PBox; +//! +//! memcx::current_context(|cx| { +//! let src = [1u32, 2, 3, 4]; +//! let dst: PBox<[u32]> = PBox::from_slice_in(&src, cx).expect("OOM"); +//! assert_eq!(&*dst, &src); +//! }); +//! ``` + use crate::callconv::{BoxRet, FcInfo}; use crate::datum::{BorrowDatum, Datum}; use crate::layout::PassBy; use crate::memcx::{MemCx, OutOfMemory}; use crate::pg_sys; use core::marker::PhantomData; +use core::mem::MaybeUninit; use core::ops::{Deref, DerefMut}; use core::ptr::NonNull; @@ -15,6 +54,19 @@ use pgrx_sql_entity_graph::metadata::{ [stdbox]: alloc::boxed::Box + +Lifetime escape is statically rejected. The following snippet must fail to type-check because the `MemCx`-bound lifetime cannot outlive the context it was borrowed from: + +```compile_fail +use pgrx::memcx::MemCx; +use pgrx::palloc::PBox; + +fn smuggle<'long, 'short>(cx: &MemCx<'short>) -> PBox<'long, [u32]> { + PBox::<[core::mem::MaybeUninit]>::new_uninit_slice_in(4, cx) + .map(|p| unsafe { p.assume_init() }) + .unwrap() +} +``` */ #[repr(transparent)] pub struct PBox<'mcx, T: ?Sized> { @@ -24,7 +76,7 @@ pub struct PBox<'mcx, T: ?Sized> { impl<'mcx, T: ?Sized> PBox<'mcx, T> { // # Safety - // The same constraints as [`Box::from_raw`], AND + // The same constraints as [`Box::from_raw`] // - you assert the pointer was allocated in the `MemCx` // - you assert the pointer may be freed by `pfree` pub unsafe fn from_raw_in(ptr: NonNull, _cx: &MemCx<'mcx>) -> PBox<'mcx, T> { @@ -41,13 +93,234 @@ impl<'mcx, T: Sized> PBox<'mcx, T> { pub fn try_new_in(val: T, memcx: &MemCx<'mcx>) -> Result { const { assert!(align_of::() <= size_of::()) }; let ptr = memcx.alloc_bytes(size_of::())?.cast(); - // SAFETY: We were guaranteed an appropriately sized allocation to write to, - // and we have asserted our alignment maximum was upheld + // SAFETY: We were guaranteed an appropriately sized allocation to write to, and we have asserted our alignment maximum was upheld unsafe { ptr.write(val) }; Ok(PBox { ptr, _cx: PhantomData }) } } +impl<'mcx, T: ?Sized> PBox<'mcx, T> { + /// Raw pointer to the boxed value. Does not dereference. + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() as *const T + } + /// Raw mutable pointer to the boxed value. Does not dereference. + pub fn as_mut_ptr(&mut self) -> *mut T { + self.ptr.as_ptr() + } +} + +impl<'mcx, T> PBox<'mcx, [MaybeUninit]> { + /// Allocate a slice of `len` uninitialized `T` inside `cx`. + /// + /// Mirrors [`alloc::boxed::Box::new_uninit_slice`] but the resulting slice is bound to the lifetime of `cx`. + pub fn new_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + // Match the alignment bound on `new_in`: palloc only guarantees `MAXIMUM_ALIGNOF` (= sizeof(Datum) = 8 on 64-bit), and on PG<16 `alloc_layout` falls back to size-only `MemoryContextAllocExtended` which silently drops `layout.align()`. + const { assert!(align_of::() <= size_of::()) }; + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_layout(layout)?; + // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no initialization, so building the fat pointer over `len` elements is sound. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// Allocate a slice of `len` zero-initialized `T` (kept as `MaybeUninit`) inside `cx`. Caller can transmute via [`PBox::assume_init`] when zero is a valid bit pattern for `T`. + pub fn new_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_layout_zeroed(layout)?; + // SAFETY: `raw` is non-null, sized by `layout`; MaybeUninit needs no init. + let slice_ptr: *mut [core::mem::MaybeUninit] = core::ptr::slice_from_raw_parts_mut( + raw.as_ptr().cast::>(), + len, + ); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// # Safety + /// Caller asserts every element of the slice has been initialized. + pub unsafe fn assume_init(self) -> PBox<'mcx, [T]> { + let p = self.ptr.as_ptr() as *mut [T]; + // SAFETY: caller upholds the init invariant; the data pointer was already non-null (it came from PBox::new_uninit_slice_in / new_zeroed_slice_in). + let ptr = unsafe { NonNull::new_unchecked(p) }; + // `PBox` has no Drop impl today (palloc memory is reclaimed by MemoryContext reset, not Drop), so `forget` is a no-op. Kept as a defensive forward-compatibility guard: if `Drop` is ever added, this prevents a double-free between `self` and the returned PBox. + core::mem::forget(self); + PBox { ptr, _cx: PhantomData } + } + + /// Like [`new_uninit_slice_in`] but permits allocations above Postgres's 1 GiB `MaxAllocSize`. Only use for genuinely huge Rust-side scratch buffers — values returned to Postgres still cap at 1 GiB regardless of how they were allocated. + /// + /// [`new_uninit_slice_in`]: PBox::new_uninit_slice_in + pub fn new_huge_uninit_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_huge_layout(layout)?; + // SAFETY: same as new_uninit_slice_in. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } + + /// Zeroed counterpart of [`new_huge_uninit_slice_in`]. + /// + /// [`new_huge_uninit_slice_in`]: PBox::new_huge_uninit_slice_in + pub fn new_huge_zeroed_slice_in(len: usize, cx: &MemCx<'mcx>) -> Result { + const { assert!(align_of::() <= size_of::()) }; + let layout = core::alloc::Layout::array::(len).map_err(|_| OutOfMemory::new())?; + let raw = cx.alloc_huge_layout_zeroed(layout)?; + // SAFETY: same as new_zeroed_slice_in. + let slice_ptr: *mut [MaybeUninit] = + core::ptr::slice_from_raw_parts_mut(raw.as_ptr().cast::>(), len); + let ptr = unsafe { NonNull::new_unchecked(slice_ptr) }; + Ok(PBox { ptr, _cx: PhantomData }) + } +} + +impl<'mcx, T> PBox<'mcx, [T]> { + /// Number of elements in the slice. + pub fn len(&self) -> usize { + self.ptr.len() + } + + /// `true` if the slice has zero elements. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl<'mcx, T: Copy> PBox<'mcx, [T]> { + /// Allocate a slice of `src.len()` elements in `cx` and copy `src` into it. + pub fn from_slice_in(src: &[T], cx: &MemCx<'mcx>) -> Result { + let mut uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(src.len(), cx)?; + // SAFETY: `uninit.len() == src.len()` by construction. T: Copy, so a bytewise copy fully initializes every element with no Drop concerns. + unsafe { + let dst = (uninit.as_mut_ptr() as *mut MaybeUninit).cast::(); + core::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len()); + Ok(uninit.assume_init()) + } + } + + /// Like [`from_slice_in`] but permits allocations above Postgres's 1 GiB `MaxAllocSize`. Use only for genuinely huge Rust-side scratch buffers; values returned to Postgres still cap at 1 GiB. + /// + /// [`from_slice_in`]: PBox::from_slice_in + pub fn from_huge_slice_in(src: &[T], cx: &MemCx<'mcx>) -> Result { + let mut uninit = PBox::<[MaybeUninit]>::new_huge_uninit_slice_in(src.len(), cx)?; + // SAFETY: as `from_slice_in`. + unsafe { + let dst = (uninit.as_mut_ptr() as *mut MaybeUninit).cast::(); + core::ptr::copy_nonoverlapping(src.as_ptr(), dst, src.len()); + Ok(uninit.assume_init()) + } + } +} + +impl<'mcx, T: Copy> PBox<'mcx, [T]> { + /// Allocate a slice of `iter.len()` elements in `cx` and write each element produced by `iter`. Requires `ExactSizeIterator` so the allocation size is known upfront with no resizing. + /// + /// `T: Copy` is required to match [`from_slice_in`]: `PBox` is not dropped when it goes out of scope (the backing memory is reclaimed in bulk when the `MemCx` resets), so `Drop` impls on `T` would never run. Restricting to `Copy` makes that asymmetry impossible to express by mistake. + /// + /// # Errors + /// + /// - [`FromIterError::OutOfMemory`] if the initial allocation fails. + /// - [`FromIterError::IteratorUnderYield`] if `iter` yields **fewer** + /// items than its `ExactSizeIterator::len()` advertised. This is + /// surfaced as a distinct variant (not folded into `OutOfMemory`) + /// because it's an iterator-contract bug, not an allocator failure; + /// callers logging or branching on OOM should not see this case + /// masquerade as one. The partially-written buffer stays in `cx` + /// until the context resets. + /// + /// If `iter` yields **more** items than promised, the iterator is dropped immediately after the `len`-th element is written. The function then returns `Ok` with the first `len` elements. + /// + /// [`from_slice_in`]: PBox::from_slice_in + pub fn from_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + { + let iter = iter.into_iter(); + let len = iter.len(); + let uninit = PBox::<[MaybeUninit]>::new_uninit_slice_in(len, cx) + .map_err(|_| FromIterError::OutOfMemory)?; + Self::fill_from_iter(uninit, iter, len) + } + + /// Huge counterpart of [`from_iter_in`]. See [`from_huge_slice_in`] for + /// when to use the huge variants. + /// + /// [`from_iter_in`]: PBox::from_iter_in + /// [`from_huge_slice_in`]: PBox::from_huge_slice_in + pub fn from_huge_iter_in(iter: I, cx: &MemCx<'mcx>) -> Result + where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, + { + let iter = iter.into_iter(); + let len = iter.len(); + let uninit = PBox::<[MaybeUninit]>::new_huge_uninit_slice_in(len, cx) + .map_err(|_| FromIterError::OutOfMemory)?; + Self::fill_from_iter(uninit, iter, len) + } + + fn fill_from_iter( + mut uninit: PBox<'mcx, [MaybeUninit]>, + iter: I, + len: usize, + ) -> Result + where + I: Iterator, + { + // Thin pointer to the first element of the uninit buffer. + let base: *mut T = uninit.as_mut_ptr().cast::(); + + let mut count = 0usize; + for item in iter { + if count >= len { + // Iterator over-yielded vs ExactSizeIterator promise; stop here. + break; + } + // SAFETY: `count < len`, so `base.add(count)` is in-bounds for the freshly allocated slice. The slot is uninit MaybeUninit; writing a T is sound. On panic the buffer stays uninitialized — safe because `T: Copy` has no Drop and `PBox<[MaybeUninit]>` has no Drop either; the memory is reclaimed when `cx` resets. + unsafe { core::ptr::write(base.add(count), item) }; + count += 1; + } + + if count < len { + // Iterator under-yielded; the buffer stays in `cx` until reset. + return Err(FromIterError::IteratorUnderYield { promised: len, yielded: count }); + } + + // SAFETY: every slot in the slice was written above (count == len). + Ok(unsafe { uninit.assume_init() }) + } +} + +/// Error returned by [`PBox::from_iter_in`]. +#[derive(Debug)] +pub enum FromIterError { + /// The initial slice allocation in the [`MemCx`] failed. + OutOfMemory, + /// The iterator yielded fewer items than its [`ExactSizeIterator::len`] advertised. This indicates a buggy iterator implementation; the buffer was allocated but only partially initialized. + IteratorUnderYield { promised: usize, yielded: usize }, +} + +impl core::fmt::Display for FromIterError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + FromIterError::OutOfMemory => f.write_str("out of memory"), + FromIterError::IteratorUnderYield { promised, yielded } => write!( + f, + "ExactSizeIterator under-yielded: promised {promised} items, got {yielded}", + ), + } + } +} + +impl std::error::Error for FromIterError {} + unsafe impl<'mcx, T> BoxRet for PBox<'mcx, T> where T: ?Sized + BorrowDatum, @@ -57,11 +330,15 @@ where PassBy::Value => { // start with a zeroed Datum, just to minimize funny business let mut datum = pg_sys::Datum::null(); - // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum, - // and PBox must have an initialized pointee, so a copy is sound-by-construction. + // SAFETY: Due to BorrowDatum, this type has a definite size less than a Datum, and PBox must have an initialized pointee, so a copy is sound-by-construction. unsafe { let size = size_of_val(&*self.ptr.as_ptr()); - debug_assert!(size <= size_of::()); + // Hard assert (not debug_assert): a `?Sized` `BorrowDatum` whose dynamic size exceeds `Datum` would otherwise overrun the on-stack `Datum` slot via `copy_from_nonoverlapping`. Real PassBy::Value types are all `Sized` and ≤ 8 bytes today, but a future fat-pointee implementing `BorrowDatum` with `PASS = PassBy::Value` must crash, not corrupt the stack. + assert!( + size <= size_of::(), + "PBox::box_into: PassBy::Value pointee size {size} exceeds Datum size {}", + size_of::() + ); // using `BorrowDatum::point_from` handles endianness let datum_ptr = T::point_from(NonNull::from_mut(&mut datum).cast::()); datum_ptr.cast::().copy_from_nonoverlapping(self.ptr.cast(), size); @@ -75,8 +352,7 @@ where } } -/// SAFETY: SQL has no "pointers" so by-val and by-ref calling conventions are identical, -/// and all `PBox` truly does is enable pass-by-ref returns of unsized values. +/// SAFETY: SQL has no "pointers" so by-val and by-ref calling conventions are identical,and all `PBox` truly does is enable pass-by-ref returns of unsized values. unsafe impl<'mcx, T> SqlTranslatable for PBox<'mcx, T> where T: SqlTranslatable + ?Sized, @@ -108,3 +384,19 @@ where unsafe { self.ptr.as_mut() } } } + +impl<'mcx, T> Deref for PBox<'mcx, [T]> { + type Target = [T]; + + fn deref(&self) -> &[T] { + // SAFETY: `self.ptr` is a valid fat pointer to an initialized [T] slice (the Sized constructors only return PBox<[T]> after assume_init or from_slice_in/from_iter_in, all of which uphold initialization). + unsafe { self.ptr.as_ref() } + } +} + +impl<'mcx, T> DerefMut for PBox<'mcx, [T]> { + fn deref_mut(&mut self) -> &mut [T] { + // SAFETY: same as Deref, plus we hold `&mut self` so aliasing is exclusive. + unsafe { self.ptr.as_mut() } + } +} diff --git a/pgrx/src/varlena.rs b/pgrx/src/varlena.rs index 9cbe732da..c4c27cb80 100644 --- a/pgrx/src/varlena.rs +++ b/pgrx/src/varlena.rs @@ -14,8 +14,7 @@ use core::{ops::DerefMut, slice, str}; /// # Safety /// -/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer -/// that is aligned to 4 bytes, and that the `len` is a half of [`i32::MAX`] +/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer that is aligned to 4 bytes, and that the `len` is a half of [`i32::MAX`] #[inline(always)] pub unsafe fn set_varsize_4b(ptr: *mut pg_sys::varlena, len: i32) { // #ifdef WORDS_BIGENDIAN @@ -28,8 +27,7 @@ pub unsafe fn set_varsize_4b(ptr: *mut pg_sys::varlena, len: i32) { // SAFETY: A varlena can be safely cast to a varattrib_4b let header = &mut (*ptr.cast::()).va_4byte.deref_mut().va_header; - // Using core::ptr::write(), which never calls drop(), to prevent - // automatically dropping a field of a ManuallyDrop + // Using core::ptr::write(), which never calls drop(), to prevent automatically dropping a field of a ManuallyDrop core::ptr::write(header, encode_vlen_4b(len)) } @@ -51,8 +49,7 @@ pub(crate) fn encode_vlen_1b(len: i32) -> u8 { /// # Safety /// -/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer -/// that is aligned to 4 bytes. +/// The caller asserts the specified `ptr` really is a non-null, palloc'd [`pg_sys::varlena`] pointer that is aligned to 4 bytes. #[inline(always)] #[deprecated(since = "0.12.0", note = "you probably meant set_varsize_4b")] pub unsafe fn set_varsize(ptr: *mut pg_sys::varlena, len: i32) { @@ -417,8 +414,7 @@ pub unsafe fn vardata_any(ptr: *const pg_sys::varlena) -> *const std::os::raw::c /// /// This function is unsafe because it blindly assumes the provided varlena pointer is non-null. /// -/// Note also that this function is zero-copy and the underlying Rust &str is backed by Postgres-allocated -/// memory. As such, the return value will become invalid the moment Postgres frees the varlena +/// Note also that this function is zero-copy and the underlying Rust &str is backed by Postgres-allocated memory. As such, the return value will become invalid the moment Postgres frees the varlena #[inline] pub unsafe fn text_to_rust_str<'a>( varlena: *const pg_sys::varlena, @@ -448,8 +444,7 @@ pub unsafe fn text_to_rust_str_unchecked<'a>(varlena: *const pg_sys::varlena) -> /// /// This function is unsafe because it blindly assumes the provided varlena pointer is non-null. /// -/// Note also that this function is zero-copy and the underlying Rust `&[u8]` slice is backed by Postgres-allocated -/// memory. As such, the return value will become invalid the moment Postgres frees the varlena +/// Note also that this function is zero-copy and the underlying Rust `&[u8]` slice is backed by Postgres-allocated memory. As such, the return value will become invalid the moment Postgres frees the varlena #[inline] pub unsafe fn varlena_to_byte_slice<'a>(varlena: *const pg_sys::varlena) -> &'a [u8] { let len = varsize_any_exhdr(varlena);