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