Skip to main content

cufflink_pgcodec/
numeric.rs

1//! Exact Postgres `NUMERIC` values: decimal text in, the binary wire format
2//! (base-10000 digits, weight, sign, display scale) out, and back, with
3//! `NaN` and the infinities. No precision is lost in either direction.
4
5use 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;
20/// The wire format counts base-10000 digits in an `i16`.
21const 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        /// Base-10000 digits, without leading or trailing zero digits.
33        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    /// The value when it is an integer that fits `i64`.
59    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}