1use 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;
19const UNIX_OFFSET_US: i64 = -UNIX_TO_POSTGRES_USECS;
21const SUBMS_MINIMAL_STEP_NS: i64 = NS_PER_MS / (1 << 12) + 1;
23const GREGORIAN_EPOCH_JDATE: i64 = 2_299_161;
25
26#[derive(Debug, Default)]
28pub struct Randomness {
29 prng: Option<Prng>,
30 global: Option<Prng>,
32 previous_ns: i64,
33}
34
35thread_local! {
36 static LOCAL: RefCell<Randomness> = RefCell::new(Randomness::default());
38 static BYTES: Cell<Option<Prng>> = const { Cell::new(None) };
40}
41
42#[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
61fn 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
68fn 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
73pub(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
84fn 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
96fn 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
110fn drandom(call: &Call<'_>, _: &[Value]) -> Result<Value> {
112 Ok(Value::Float8(with_prng(call.session, Prng::double)))
113}
114
115fn 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#[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
151fn 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
158fn 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
164fn 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
172fn 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#[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
198fn 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
204fn 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
238fn 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#[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
278fn 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
288fn 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}