Skip to main content

rupg_func/
random.rs

1//! The random functions of `pseudorandomfuncs.c` and the UUID functions of `uuid.c`: `setseed`, `random` of each type, `random_normal`, `gen_random_uuid`, `uuidv7`, `uuid_extract_version` and `uuid_extract_timestamp`, with the casts between `uuid` and `bytea`. The generator is that of `pg_prng.c`, so the numbers after `setseed` are the numbers of PostgreSQL. A generator with no seed takes its seed from [`Session::strong_random`], as `pg_prng_strong_seed` does.
2
3use std::cell::{Cell, RefCell};
4
5use rupg_common::{Error, Result, SqlState};
6use rupg_types::prng::Prng;
7use rupg_types::{
8    DATE_INFINITY, DATE_NEGATIVE_INFINITY, Interval, POSTGRES_EPOCH_JDATE, TIMESTAMP_INFINITY,
9    TIMESTAMP_NEGATIVE_INFINITY, UNIX_TO_POSTGRES_USECS, Value,
10};
11
12use crate::calendar::format_g;
13use crate::datetime::timestamp_plus;
14use crate::{Call, Kernel, Session, bad_value};
15
16const NS_PER_MS: i64 = 1_000_000;
17const NS_PER_US: i64 = 1000;
18const US_PER_MS: i64 = 1000;
19/// The microseconds from 1970-01-01 to 2000-01-01, `PG_UNIX_EPOCH_OFFSET_US`.
20const UNIX_OFFSET_US: i64 = -UNIX_TO_POSTGRES_USECS;
21/// `SUBMS_MINIMAL_STEP_NS`: the smallest step of the clock of `uuidv7`, with the 12 bits of the fraction of a millisecond that Linux gives.
22const SUBMS_MINIMAL_STEP_NS: i64 = NS_PER_MS / (1 << 12) + 1;
23/// `GREGORIAN_EPOCH_JDATE`: the Julian day of 1582-10-15, the epoch of UUID version 1.
24const GREGORIAN_EPOCH_JDATE: i64 = 2_299_161;
25
26/// The random state of a session: the generator of `random`, and the last time of `uuidv7` in nanoseconds since 1970-01-01.
27#[derive(Debug, Default)]
28pub struct Randomness {
29    prng: Option<Prng>,
30    /// `pg_global_prng_state`, the generator of `array_shuffle` and `array_sample`, which `setseed` does not change.
31    global: Option<Prng>,
32    previous_ns: i64,
33}
34
35thread_local! {
36    /// The state of a session with no server.
37    static LOCAL: RefCell<Randomness> = RefCell::new(Randomness::default());
38    /// The generator of the random bytes of a session with no server.
39    static BYTES: Cell<Option<Prng>> = const { Cell::new(None) };
40}
41
42/// Random bytes from a generator with a seed from the clock, as `initialize_prng` seeds the generator when the system gives no random bytes. A session with no server uses these bytes.
43#[allow(clippy::cast_sign_loss)]
44pub(crate) fn clock_random(clock: i64, buf: &mut [u8]) {
45    BYTES.with(|cell| {
46        let mut prng = cell.get().unwrap_or_else(|| Prng::seed(clock as u64));
47        for chunk in buf.chunks_mut(8) {
48            chunk.copy_from_slice(&prng.next_u64().to_ne_bytes()[..chunk.len()]);
49        }
50        cell.set(Some(prng));
51    });
52}
53
54fn with_state<T>(session: &dyn Session, f: impl FnOnce(&mut Randomness) -> T) -> T {
55    match session.randomness() {
56        Some(cell) => f(&mut cell.borrow_mut()),
57        None => LOCAL.with(|cell| f(&mut cell.borrow_mut())),
58    }
59}
60
61/// `pg_prng_strong_seed`.
62fn strong_seed(session: &dyn Session) -> Prng {
63    let mut bytes = [0; 16];
64    session.strong_random(&mut bytes);
65    Prng::from_bytes(bytes)
66}
67
68/// `initialize_prng` and then `f` with the generator.
69fn with_prng<T>(session: &dyn Session, f: impl FnOnce(&mut Prng) -> T) -> T {
70    with_state(session, |state| f(state.prng.get_or_insert_with(|| strong_seed(session))))
71}
72
73/// `pg_prng_uint64_range` with `pg_global_prng_state`.
74pub(crate) fn global_range(session: &dyn Session, rmin: u64, rmax: u64) -> u64 {
75    with_state(session, |state| {
76        state.global.get_or_insert_with(|| strong_seed(session)).uint64_range(rmin, rmax)
77    })
78}
79
80fn parameter_error(message: &str) -> Error {
81    Error::new(SqlState::INVALID_PARAMETER_VALUE, message)
82}
83
84/// `CHECK_RANGE_BOUNDS`.
85fn check_bounds<T: PartialOrd>(rmin: T, rmax: T) -> Result<()> {
86    if rmin > rmax {
87        return Err(parameter_error("lower bound must be less than or equal to upper bound"));
88    }
89    Ok(())
90}
91
92fn finite_bounds(finite: bool) -> Result<()> {
93    if finite { Ok(()) } else { Err(parameter_error("lower and upper bounds must be finite")) }
94}
95
96/// `setseed`.
97fn setseed(call: &Call<'_>, args: &[Value]) -> Result<Value> {
98    let [Value::Float8(seed)] = args else { return Err(bad_value()) };
99    let seed = *seed;
100    if !(-1.0..=1.0).contains(&seed) {
101        return Err(parameter_error(&format!(
102            "setseed parameter {} is out of allowed range [-1,1]",
103            format_g(seed)
104        )));
105    }
106    with_state(call.session, |state| state.prng = Some(Prng::fseed(seed)));
107    Ok(Value::Text(String::new()))
108}
109
110/// `drandom`: a value from 0 up to but not including 1.
111fn drandom(call: &Call<'_>, _: &[Value]) -> Result<Value> {
112    Ok(Value::Float8(with_prng(call.session, Prng::double)))
113}
114
115/// `drandom_normal`.
116fn drandom_normal(call: &Call<'_>, args: &[Value]) -> Result<Value> {
117    let [Value::Float8(mean), Value::Float8(stddev)] = args else { return Err(bad_value()) };
118    let z = with_prng(call.session, Prng::double_normal);
119    Ok(Value::Float8(stddev * z + mean))
120}
121
122/// `int4random`, `int8random`, `date_random`, `timestamp_random` and `timestamptz_random`.
123#[allow(clippy::cast_possible_truncation)]
124fn int_random(call: &Call<'_>, args: &[Value]) -> Result<Value> {
125    let range = |rmin: i64, rmax: i64| -> Result<i64> {
126        check_bounds(rmin, rmax)?;
127        Ok(with_prng(call.session, |prng| prng.int64_range(rmin, rmax)))
128    };
129    Ok(match args {
130        [Value::Int4(a), Value::Int4(b)] => Value::Int4(range((*a).into(), (*b).into())? as i32),
131        [Value::Int8(a), Value::Int8(b)] => Value::Int8(range(*a, *b)?),
132        [Value::Date(a), Value::Date(b)] => {
133            check_bounds(a, b)?;
134            finite_bounds(*a != DATE_NEGATIVE_INFINITY && *b != DATE_INFINITY)?;
135            Value::Date(range((*a).into(), (*b).into())? as i32)
136        }
137        [Value::Timestamp(a), Value::Timestamp(b)]
138        | [Value::TimestampTz(a), Value::TimestampTz(b)] => {
139            check_bounds(a, b)?;
140            finite_bounds(*a != TIMESTAMP_NEGATIVE_INFINITY && *b != TIMESTAMP_INFINITY)?;
141            let result = range(*a, *b)?;
142            match args[0] {
143                Value::Timestamp(_) => Value::Timestamp(result),
144                _ => Value::TimestampTz(result),
145            }
146        }
147        _ => return Err(bad_value()),
148    })
149}
150
151/// `numeric_random`.
152fn numeric_random(call: &Call<'_>, args: &[Value]) -> Result<Value> {
153    let [Value::Numeric(rmin), Value::Numeric(rmax)] = args else { return Err(bad_value()) };
154    let result = with_prng(call.session, |prng| rupg_types::Numeric::random(prng, rmin, rmax))?;
155    Ok(Value::Numeric(result))
156}
157
158/// `uuid_set_version`: the version in the top four bits of byte 6, and the variant of RFC 9562 in the top two bits of byte 8.
159fn set_version(uuid: &mut [u8; 16], version: u8) {
160    uuid[6] = (uuid[6] & 0x0f) | (version << 4);
161    uuid[8] = (uuid[8] & 0x3f) | 0x80;
162}
163
164/// `gen_random_uuid`: a UUID of version 4.
165fn gen_random_uuid(call: &Call<'_>, _: &[Value]) -> Result<Value> {
166    let mut uuid = [0; 16];
167    call.session.strong_random(&mut uuid);
168    set_version(&mut uuid, 4);
169    Ok(Value::Uuid(uuid))
170}
171
172/// `get_real_time_ns_ascending`: the time in nanoseconds since 1970-01-01, at least `SUBMS_MINIMAL_STEP_NS` after the last time of the session.
173fn time_ns_ascending(session: &dyn Session) -> i64 {
174    let mut ns = (session.clock() + UNIX_OFFSET_US) * NS_PER_US;
175    with_state(session, |state| {
176        if state.previous_ns + SUBMS_MINIMAL_STEP_NS >= ns {
177            ns = state.previous_ns + SUBMS_MINIMAL_STEP_NS;
178        }
179        state.previous_ns = ns;
180    });
181    ns
182}
183
184/// `generate_uuidv7`: 48 bits of milliseconds since 1970-01-01, 12 bits of the fraction of the millisecond, and random bits.
185#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
186fn generate_uuidv7(session: &dyn Session, unix_ts_ms: i64, sub_ms: i64) -> Value {
187    let mut uuid = [0; 16];
188    let ms = (unix_ts_ms as u64).to_be_bytes();
189    uuid[..6].copy_from_slice(&ms[2..]);
190    let precision = (sub_ms * (1 << 12) / NS_PER_MS) as u32;
191    uuid[6] = (precision >> 8) as u8;
192    uuid[7] = precision as u8;
193    session.strong_random(&mut uuid[8..]);
194    set_version(&mut uuid, 7);
195    Value::Uuid(uuid)
196}
197
198/// `uuidv7`.
199fn uuidv7(call: &Call<'_>, _: &[Value]) -> Result<Value> {
200    let ns = time_ns_ascending(call.session);
201    Ok(generate_uuidv7(call.session, ns / NS_PER_MS, ns % NS_PER_MS))
202}
203
204/// `uuidv7_interval`: a UUID of version 7 with the time now moved by an interval.
205fn uuidv7_interval(call: &Call<'_>, args: &[Value]) -> Result<Value> {
206    let [Value::Interval(shift)] = args else { return Err(bad_value()) };
207    if *shift == Interval::INFINITY || *shift == Interval::NEGATIVE_INFINITY {
208        return Err(Error::new(
209            SqlState::DATETIME_VALUE_OUT_OF_RANGE,
210            "interval out of range for UUID version 7",
211        )
212        .with_detail("UUID version 7 does not support infinite intervals."));
213    }
214    let ns = time_ns_ascending(call.session);
215    let zone = call.session.zone()?;
216    let ts = timestamp_plus(ns / NS_PER_US - UNIX_OFFSET_US, shift, Some(zone.as_ref()))?;
217    let max = ((1i64 << 48) - 1) * US_PER_MS - UNIX_OFFSET_US;
218    if ts < -UNIX_OFFSET_US || ts > max {
219        return Err(Error::new(
220            SqlState::DATETIME_VALUE_OUT_OF_RANGE,
221            "timestamp out of range for UUID version 7",
222        )
223        .with_detail(
224            "UUID version 7 supports timestamps from 1970-01-01 to approximately year 10889.",
225        ));
226    }
227    let us = ts + UNIX_OFFSET_US;
228    Ok(generate_uuidv7(call.session, us / US_PER_MS, (us % US_PER_MS) * NS_PER_US + ns % NS_PER_US))
229}
230
231fn uuid_arg(args: &[Value]) -> Result<&[u8; 16]> {
232    match args {
233        [Value::Uuid(uuid)] => Ok(uuid),
234        _ => Err(bad_value()),
235    }
236}
237
238/// `uuid_extract_version`: null for a UUID that is not of the variant of RFC 9562.
239fn uuid_extract_version(_: &Call<'_>, args: &[Value]) -> Result<Value> {
240    let uuid = uuid_arg(args)?;
241    if uuid[8] & 0xc0 != 0x80 {
242        return Ok(Value::Null);
243    }
244    Ok(Value::Int2(i16::from(uuid[6] >> 4)))
245}
246
247/// `uuid_extract_timestamp`: the time of a UUID of version 1 or 7, and null for the other UUIDs.
248#[allow(clippy::cast_possible_wrap)]
249fn uuid_extract_timestamp(_: &Call<'_>, args: &[Value]) -> Result<Value> {
250    let uuid = uuid_arg(args)?;
251    if uuid[8] & 0xc0 != 0x80 {
252        return Ok(Value::Null);
253    }
254    let byte = |at: usize, shift: u32| u64::from(uuid[at]) << shift;
255    Ok(match uuid[6] >> 4 {
256        1 => {
257            let tms = byte(0, 24)
258                + byte(1, 16)
259                + byte(2, 8)
260                + byte(3, 0)
261                + byte(4, 40)
262                + byte(5, 32)
263                + (u64::from(uuid[6] & 0xf) << 56)
264                + byte(7, 48);
265            let offset =
266                (i64::from(POSTGRES_EPOCH_JDATE) - GREGORIAN_EPOCH_JDATE) * 86_400 * 1_000_000;
267            Value::TimestampTz((tms / 10) as i64 - offset)
268        }
269        7 => {
270            let tms =
271                byte(5, 0) + byte(4, 8) + byte(3, 16) + byte(2, 24) + byte(1, 32) + byte(0, 40);
272            Value::TimestampTz((tms * US_PER_MS as u64) as i64 - UNIX_OFFSET_US)
273        }
274        _ => Value::Null,
275    })
276}
277
278/// `bytea_uuid`: the cast of 16 bytes to a UUID.
279fn bytea_uuid(_: &Call<'_>, args: &[Value]) -> Result<Value> {
280    let [Value::Bytea(bytes)] = args else { return Err(bad_value()) };
281    let uuid = <[u8; 16]>::try_from(bytes.as_slice()).map_err(|_| {
282        Error::new(SqlState::INVALID_BINARY_REPRESENTATION, "invalid input length for type uuid")
283            .with_detail(format!("Expected 16 bytes, got {}.", bytes.len()))
284    })?;
285    Ok(Value::Uuid(uuid))
286}
287
288/// `uuid_bytea`: the 16 bytes of a UUID.
289fn uuid_bytea(_: &Call<'_>, args: &[Value]) -> Result<Value> {
290    Ok(Value::Bytea(uuid_arg(args)?.to_vec()))
291}
292
293pub(crate) fn by_src(src: &str) -> Option<Kernel> {
294    Some(match src {
295        "setseed" => setseed,
296        "drandom" => drandom,
297        "drandom_normal" => drandom_normal,
298        "int4random" | "int8random" | "date_random" | "timestamp_random" | "timestamptz_random" => {
299            int_random
300        }
301        "numeric_random" => numeric_random,
302        "gen_random_uuid" => gen_random_uuid,
303        "uuidv7" => uuidv7,
304        "uuidv7_interval" => uuidv7_interval,
305        "uuid_extract_version" => uuid_extract_version,
306        "uuid_extract_timestamp" => uuid_extract_timestamp,
307        "bytea_uuid" => bytea_uuid,
308        "uuid_bytea" => uuid_bytea,
309        _ => return None,
310    })
311}