1use bytes::{BufMut, BytesMut};
6use std::error::Error;
7use std::fmt;
8use std::str::FromStr;
9use tokio_postgres::types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
10
11const SIGN_POSITIVE: u16 = 0x0000;
12const SIGN_NEGATIVE: u16 = 0x4000;
13const SIGN_NAN: u16 = 0xC000;
14const SIGN_INFINITY: u16 = 0xD000;
15const SIGN_NEG_INFINITY: u16 = 0xF000;
16const DSCALE_MAX: usize = 0x3FFF;
17const WEIGHT_MAX: i32 = i16::MAX as i32;
18const WEIGHT_MIN: i32 = i16::MIN as i32;
19const BASE: i32 = 10_000;
20const DIGITS_MAX: usize = i16::MAX as usize;
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum PgNumeric {
25 NaN,
26 Infinity,
27 NegInfinity,
28 Finite {
29 negative: bool,
30 weight: i16,
31 dscale: u16,
32 digits: Vec<i16>,
34 },
35}
36
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct NumericError(String);
39
40impl fmt::Display for NumericError {
41 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42 f.write_str(&self.0)
43 }
44}
45
46impl Error for NumericError {}
47
48impl PgNumeric {
49 fn zero(dscale: u16) -> Self {
50 PgNumeric::Finite {
51 negative: false,
52 weight: 0,
53 dscale,
54 digits: Vec::new(),
55 }
56 }
57
58 pub fn to_i64(&self) -> Option<i64> {
60 let text = self.to_string();
61 let (int, frac) = text.split_once('.').unwrap_or((&text, ""));
62 frac.bytes()
63 .all(|b| b == b'0')
64 .then(|| int.parse().ok())
65 .flatten()
66 }
67}
68
69impl FromStr for PgNumeric {
70 type Err = NumericError;
71
72 fn from_str(text: &str) -> Result<Self, NumericError> {
73 let invalid = || NumericError(format!("{text:?} is not a decimal number"));
74 let trimmed = text.trim();
75 match trimmed.to_ascii_lowercase().as_str() {
76 "nan" => return Ok(PgNumeric::NaN),
77 "infinity" | "+infinity" | "inf" | "+inf" => return Ok(PgNumeric::Infinity),
78 "-infinity" | "-inf" => return Ok(PgNumeric::NegInfinity),
79 _ => {}
80 }
81 let (negative, unsigned) = match trimmed.as_bytes().first() {
82 Some(b'-') => (true, &trimmed[1..]),
83 Some(b'+') => (false, &trimmed[1..]),
84 _ => (false, trimmed),
85 };
86 let (mantissa, exponent) = match unsigned.find(['e', 'E']) {
87 Some(at) => (
88 &unsigned[..at],
89 unsigned[at + 1..].parse::<i32>().map_err(|_| invalid())?,
90 ),
91 None => (unsigned, 0),
92 };
93 let (int_part, frac_part) = mantissa.split_once('.').unwrap_or((mantissa, ""));
94 let all_digits = |s: &str| s.bytes().all(|b| b.is_ascii_digit());
95 if int_part.len() + frac_part.len() == 0 || !all_digits(int_part) || !all_digits(frac_part)
96 {
97 return Err(invalid());
98 }
99 let dscale = usize::try_from(
100 i64::try_from(frac_part.len()).map_err(|_| invalid())? - i64::from(exponent),
101 )
102 .unwrap_or(0);
103 if dscale > DSCALE_MAX {
104 return Err(NumericError(format!(
105 "{text:?} has more than {DSCALE_MAX} fractional digits"
106 )));
107 }
108 let dscale = dscale as u16;
109 let decimal: Vec<u8> = int_part
110 .bytes()
111 .chain(frac_part.bytes())
112 .map(|b| b - b'0')
113 .collect();
114 let leading = decimal.iter().take_while(|d| **d == 0).count();
115 let significant = &decimal[leading..];
116 let trailing = significant.iter().rev().take_while(|d| **d == 0).count();
117 let significant = &significant[..significant.len() - trailing];
118 if significant.is_empty() {
119 return Ok(PgNumeric::zero(dscale));
120 }
121 let point = i64::try_from(int_part.len()).map_err(|_| invalid())? + i64::from(exponent)
122 - i64::try_from(leading).map_err(|_| invalid())?;
123 let pad_front = (4 - point.rem_euclid(4)) % 4;
124 let weight = (point + pad_front) / 4 - 1;
125 if !(i64::from(WEIGHT_MIN)..=i64::from(WEIGHT_MAX)).contains(&weight) {
126 return Err(NumericError(format!(
127 "{text:?} is outside the numeric range"
128 )));
129 }
130 if (pad_front as usize + significant.len()).div_ceil(4) > DIGITS_MAX {
131 return Err(NumericError(format!(
132 "{text:?} has more than {} significant digits",
133 DIGITS_MAX * 4
134 )));
135 }
136 let padded: Vec<u8> = std::iter::repeat_n(0, pad_front as usize)
137 .chain(significant.iter().copied())
138 .collect();
139 let digits = padded
140 .chunks(4)
141 .map(|chunk| {
142 chunk
143 .iter()
144 .chain(std::iter::repeat(&0))
145 .take(4)
146 .fold(0i16, |acc, d| acc * 10 + i16::from(*d))
147 })
148 .collect();
149 Ok(PgNumeric::Finite {
150 negative,
151 weight: weight as i16,
152 dscale,
153 digits,
154 })
155 }
156}
157
158impl fmt::Display for PgNumeric {
159 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
160 let (negative, weight, dscale, digits) = match self {
161 PgNumeric::NaN => return f.write_str("NaN"),
162 PgNumeric::Infinity => return f.write_str("Infinity"),
163 PgNumeric::NegInfinity => return f.write_str("-Infinity"),
164 PgNumeric::Finite {
165 negative,
166 weight,
167 dscale,
168 digits,
169 } => (*negative, i32::from(*weight), usize::from(*dscale), digits),
170 };
171 let group = |i: i32| {
172 usize::try_from(i)
173 .ok()
174 .and_then(|i| digits.get(i))
175 .copied()
176 .unwrap_or(0)
177 };
178 if negative {
179 f.write_str("-")?;
180 }
181 match weight {
182 w if w < 0 => f.write_str("0")?,
183 w => {
184 write!(f, "{}", group(0))?;
185 (1..=w).try_for_each(|i| write!(f, "{:04}", group(i)))?;
186 }
187 }
188 if dscale > 0 {
189 let fraction: String = (weight + 1..)
190 .flat_map(|i| format!("{:04}", group(i)).into_bytes())
191 .take(dscale)
192 .map(char::from)
193 .collect();
194 write!(f, ".{fraction}")?;
195 }
196 Ok(())
197 }
198}
199
200impl ToSql for PgNumeric {
201 fn to_sql(&self, _: &Type, out: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
202 let (sign, weight, dscale, digits): (u16, i16, u16, &[i16]) = match self {
203 PgNumeric::NaN => (SIGN_NAN, 0, 0, &[]),
204 PgNumeric::Infinity => (SIGN_INFINITY, 0, 0, &[]),
205 PgNumeric::NegInfinity => (SIGN_NEG_INFINITY, 0, 0, &[]),
206 PgNumeric::Finite {
207 negative,
208 weight,
209 dscale,
210 digits,
211 } => (
212 if *negative {
213 SIGN_NEGATIVE
214 } else {
215 SIGN_POSITIVE
216 },
217 *weight,
218 *dscale,
219 digits,
220 ),
221 };
222 out.put_i16(i16::try_from(digits.len())?);
223 out.put_i16(weight);
224 out.put_u16(sign);
225 out.put_u16(dscale);
226 digits.iter().for_each(|d| out.put_i16(*d));
227 Ok(IsNull::No)
228 }
229
230 fn accepts(ty: &Type) -> bool {
231 *ty == Type::NUMERIC
232 }
233
234 to_sql_checked!();
235}
236
237impl<'a> FromSql<'a> for PgNumeric {
238 fn from_sql(_: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
239 let word = |i: usize| -> Result<[u8; 2], NumericError> {
240 raw.get(i * 2..i * 2 + 2)
241 .and_then(|b| b.try_into().ok())
242 .ok_or_else(|| NumericError("truncated numeric value".into()))
243 };
244 let ndigits = usize::try_from(i16::from_be_bytes(word(0)?))
245 .map_err(|_| NumericError("negative numeric digit count".into()))?;
246 let weight = i16::from_be_bytes(word(1)?);
247 let sign = u16::from_be_bytes(word(2)?);
248 let dscale = u16::from_be_bytes(word(3)?);
249 let digits = (4..4 + ndigits)
250 .map(|i| {
251 let digit = i16::from_be_bytes(word(i)?);
252 (0..BASE as i16)
253 .contains(&digit)
254 .then_some(digit)
255 .ok_or_else(|| NumericError(format!("numeric digit {digit} out of range")))
256 })
257 .collect::<Result<Vec<_>, _>>()?;
258 Ok(match sign {
259 SIGN_NAN => PgNumeric::NaN,
260 SIGN_INFINITY => PgNumeric::Infinity,
261 SIGN_NEG_INFINITY => PgNumeric::NegInfinity,
262 SIGN_POSITIVE | SIGN_NEGATIVE => {
263 let leading = digits.iter().take_while(|d| **d == 0).count();
264 let trailing = digits[leading..]
265 .iter()
266 .rev()
267 .take_while(|d| **d == 0)
268 .count();
269 let digits = digits[leading..digits.len() - trailing].to_vec();
270 match digits.is_empty() {
271 true => PgNumeric::zero(dscale),
272 false => PgNumeric::Finite {
273 negative: sign == SIGN_NEGATIVE,
274 weight: weight - leading as i16,
275 dscale,
276 digits,
277 },
278 }
279 }
280 other => return Err(NumericError(format!("unknown numeric sign {other:#06x}")).into()),
281 })
282 }
283
284 fn accepts(ty: &Type) -> bool {
285 *ty == Type::NUMERIC
286 }
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292
293 fn numeric(text: &str) -> PgNumeric {
294 text.parse().expect("valid numeric")
295 }
296
297 fn wire(value: &PgNumeric) -> Vec<u8> {
298 let mut out = BytesMut::new();
299 value.to_sql(&Type::NUMERIC, &mut out).expect("encode");
300 out.to_vec()
301 }
302
303 fn round_trip(text: &str) -> String {
304 PgNumeric::from_sql(&Type::NUMERIC, &wire(&numeric(text)))
305 .expect("decode")
306 .to_string()
307 }
308
309 #[test]
310 fn vectors_encode_to_the_postgres_layout() {
311 let words =
312 |words: &[i16]| -> Vec<u8> { words.iter().flat_map(|w| w.to_be_bytes()).collect() };
313 assert_eq!(wire(&numeric("0")), words(&[0, 0, 0, 0]));
314 assert_eq!(wire(&numeric("123.45")), words(&[2, 0, 0, 2, 123, 4500]));
315 assert_eq!(wire(&numeric("-1")), words(&[1, 0, 0x4000, 0, 1]));
316 assert_eq!(wire(&numeric("1e-20")), words(&[1, -5, 0, 20, 1]));
317 assert_eq!(wire(&numeric("10000")), words(&[1, 1, 0, 0, 1]));
318 assert_eq!(wire(&numeric("NaN")), words(&[0, 0, 0xC000u16 as i16, 0]));
319 assert_eq!(
320 wire(&numeric("Infinity")),
321 words(&[0, 0, 0xD000u16 as i16, 0])
322 );
323 assert_eq!(
324 wire(&numeric("-Infinity")),
325 words(&[0, 0, 0xF000u16 as i16, 0])
326 );
327 }
328
329 #[test]
330 fn values_survive_a_round_trip_exactly() {
331 for (input, canonical) in [
332 ("0", "0"),
333 ("-0.00", "0.00"),
334 ("1e-20", "0.00000000000000000001"),
335 (
336 "123456789012345678901234567890.123",
337 "123456789012345678901234567890.123",
338 ),
339 ("1.50", "1.50"),
340 ("0.001", "0.001"),
341 ("-42", "-42"),
342 ("+7.0", "7.0"),
343 ("1.5e3", "1500"),
344 ("12.5E-1", "1.25"),
345 ("0000.1000", "0.1000"),
346 ("99999999", "99999999"),
347 (".5", "0.5"),
348 ("5.", "5"),
349 ("nan", "NaN"),
350 ("-inf", "-Infinity"),
351 ("Infinity", "Infinity"),
352 ] {
353 assert_eq!(round_trip(input), canonical, "{input}");
354 assert_eq!(numeric(input).to_string(), canonical, "{input}");
355 }
356 }
357
358 #[test]
359 fn decimal_text_round_trips_for_generated_values() {
360 let mut seed = 0x9E37_79B9_7F4A_7C15u64;
361 let mut next = || {
362 seed ^= seed << 13;
363 seed ^= seed >> 7;
364 seed ^= seed << 17;
365 seed
366 };
367 for _ in 0..2_000 {
368 let int_len = (next() % 25) as usize;
369 let frac_len = (next() % 25) as usize;
370 let digit = |n: u64| char::from(b'0' + (n % 10) as u8);
371 let int: String = (0..int_len).map(|_| digit(next())).collect();
372 let frac: String = (0..frac_len).map(|_| digit(next())).collect();
373 let sign = if next() % 2 == 0 { "-" } else { "" };
374 let int = int.trim_start_matches('0');
375 let int = if int.is_empty() { "0" } else { int };
376 let text = match frac_len {
377 0 => format!("{sign}{int}"),
378 _ => format!("{sign}{int}.{frac}"),
379 };
380 let zero = text
381 .trim_start_matches('-')
382 .bytes()
383 .all(|b| matches!(b, b'0' | b'.'));
384 let expected = if zero {
385 text.trim_start_matches('-').to_string()
386 } else {
387 text.clone()
388 };
389 assert_eq!(round_trip(&text), expected, "{text}");
390 }
391 }
392
393 #[test]
394 fn malformed_text_is_refused() {
395 for bad in [
396 "",
397 "-",
398 ".",
399 "1.2.3",
400 "1e",
401 "abc",
402 "1,5",
403 "e5",
404 "1e99999999999",
405 "--1",
406 ] {
407 assert!(bad.parse::<PgNumeric>().is_err(), "{bad:?}");
408 }
409 assert!(format!("1e-{}", DSCALE_MAX + 1)
410 .parse::<PgNumeric>()
411 .is_err());
412 assert!("1e200000".parse::<PgNumeric>().is_err());
413 let widest = "9".repeat(131_068);
414 assert!(widest.parse::<PgNumeric>().is_ok());
415 let too_many = format!("{widest}.{}", "9".repeat(8));
416 assert!(too_many
417 .parse::<PgNumeric>()
418 .is_err_and(|e| e.to_string().contains("significant digits")));
419 }
420
421 #[test]
422 fn corrupt_wire_values_are_errors() {
423 assert!(PgNumeric::from_sql(&Type::NUMERIC, &[0, 1]).is_err());
424 assert!(PgNumeric::from_sql(&Type::NUMERIC, &[0, 1, 0, 0, 0, 0, 0, 0]).is_err());
425 assert!(
426 PgNumeric::from_sql(&Type::NUMERIC, &[0, 1, 0, 0, 0, 0, 0, 0, 0x27, 0x10]).is_err()
427 );
428 assert!(PgNumeric::from_sql(&Type::NUMERIC, &[0, 0, 0, 0, 0x12, 0x34, 0, 0]).is_err());
429 }
430
431 #[test]
432 fn integers_are_recognised() {
433 assert_eq!(numeric("42").to_i64(), Some(42));
434 assert_eq!(numeric("-42.000").to_i64(), Some(-42));
435 assert_eq!(numeric("42.5").to_i64(), None);
436 assert_eq!(numeric("1e30").to_i64(), None);
437 assert_eq!(numeric("NaN").to_i64(), None);
438 }
439}