hyperactor_mesh/
shortuuid.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! This module implements a "short" (64-bit) UUID, used to assign
10//! names to various objects in meshes.
11
12use std::str::FromStr;
13use std::sync::LazyLock;
14
15use rand::RngCore;
16
17/// So-called ["Flickr base 58"](https://www.flickr.com/groups/api/discuss/72157616713786392/)
18/// as this alphabet was used in Flickr URLs. It has nice properties: 1) characters are all
19/// URL safe, and characters which are easily confused (e.g., I and l) are removed.
20const FLICKR_BASE_58: &str = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ";
21
22/// Precomputed character ordinals for the alphabet.
23static FLICKR_BASE_58_ORD: LazyLock<[Option<usize>; 256]> = LazyLock::new(|| {
24    let mut table = [None; 256];
25    for (i, c) in FLICKR_BASE_58.chars().enumerate() {
26        table[c as usize] = Some(i);
27    }
28    table
29});
30
31/// A short (64-bit) UUID. UUIDs can be generated with [`ShortUuid::generate`],
32/// displayed as short, URL-friendly, 12-character gobbleygook, and parsed accordingly.
33///
34/// ```
35/// # use hyperactor_mesh::shortuuid::ShortUuid;
36/// let uuid = ShortUuid::generate();
37/// println!("nice, short, URL friendly: {}", uuid);
38/// assert_eq!(uuid.to_string().parse::<ShortUuid>().unwrap(), uuid);
39/// ```
40///
41/// ShortUuids have a Base-58, alphanumeric URI-friendly representation.
42/// The characters "_" and "-" are ignored when decoding, and may be
43/// safely interspersed. By default, rendered UUIDs that begin with a
44/// numeric character is prefixed with "_".
45#[derive(PartialEq, Eq, Debug, Clone)]
46pub struct ShortUuid(u64);
47
48impl ShortUuid {
49    /// Generate a new UUID.
50    pub fn generate() -> ShortUuid {
51        ShortUuid(rand::thread_rng().next_u64())
52    }
53}
54impl std::fmt::Display for ShortUuid {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        let mut num = self.0;
57        let base = FLICKR_BASE_58.len() as u64;
58        let mut result = String::with_capacity(12);
59
60        for pos in 0..12 {
61            let remainder = (num % base) as usize;
62            num /= base;
63            let c = FLICKR_BASE_58.chars().nth(remainder).unwrap();
64            result.push(c);
65            // Make sure the first position is never a digit.
66            if pos == 11 && c >= '0' && c <= '9' {
67                result.push('_');
68            }
69        }
70        assert_eq!(num, 0);
71
72        let result = result.chars().rev().collect::<String>();
73        write!(f, "{}", result)
74    }
75}
76
77#[derive(Debug, thiserror::Error, PartialEq)]
78pub enum ParseShortUuidError {
79    #[error("invalid character '{0}' in ShortUuid")]
80    InvalidCharacter(char),
81}
82
83impl FromStr for ShortUuid {
84    type Err = ParseShortUuidError;
85
86    fn from_str(s: &str) -> Result<Self, Self::Err> {
87        let base = FLICKR_BASE_58.len() as u64;
88        let mut num = 0u64;
89
90        for c in s.chars() {
91            if c == '_' || c == '-' {
92                continue;
93            }
94            num *= base;
95            if let Some(pos) = FLICKR_BASE_58_ORD[c as usize] {
96                num += pos as u64;
97            } else {
98                return Err(ParseShortUuidError::InvalidCharacter(c));
99            }
100        }
101
102        Ok(ShortUuid(num))
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn test_basic() {
112        let cases = vec![
113            (0, "_111111111111"),
114            (1, "_111111111112"),
115            (2, "_111111111113"),
116            (3, "_111111111114"),
117            (1234657890119888222, "_13Se1Maqzryj"),
118            (58 * 58 - 1, "_1111111111ZZ"),
119            (u64::MAX, "_1JPwcyDCgEup"),
120            (58 * 58, "_111111111211"),
121        ];
122
123        for (num, display) in cases {
124            let uuid = ShortUuid(num);
125            assert_eq!(uuid.to_string(), display);
126            // Round-trip test:
127            assert_eq!(uuid.to_string().parse::<ShortUuid>().unwrap(), uuid);
128        }
129    }
130
131    #[test]
132    fn test_decode() {
133        let cases = vec![
134            ("__-_1111_11_111111", 0),
135            ("_111111111112", 1),
136            ("_111111111113", 2),
137            ("_111111111114-", 3),
138            ("13Se1-Maqzr-yj", 1234657890119888222),
139            ("1111111111ZZ", 58 * 58 - 1),
140            ("1JPwcy-----DCgEup", u64::MAX),
141            ("_111111111211", 58 * 58),
142        ];
143
144        for (display, num) in cases {
145            assert_eq!(display.parse::<ShortUuid>().unwrap(), ShortUuid(num));
146        }
147    }
148
149    #[test]
150    fn test_parse_error() {
151        let invalid_cases = vec![
152            ("11111111111O", 'O'),
153            ("11111111111I", 'I'),
154            ("11111111111l", 'l'),
155            ("11111111111@", '@'),
156        ];
157
158        for (input, invalid_char) in invalid_cases {
159            assert_eq!(
160                input.parse::<ShortUuid>().unwrap_err(),
161                ParseShortUuidError::InvalidCharacter(invalid_char),
162            )
163        }
164    }
165}