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