nuid/
lib.rs

1// Copyright 2018 Ivan Porto Carrero
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15
16use std::fmt::{self, Display};
17use std::ops::Deref;
18use std::str;
19use std::sync::Mutex;
20
21use rand::distributions::Alphanumeric;
22use rand::rngs::OsRng;
23use rand::thread_rng;
24use rand::Rng;
25
26const BASE: usize = 62;
27const ALPHABET: [u8; BASE] = *b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
28
29const PRE_LEN: usize = 12;
30const MAX_SEQ: u64 = 839_299_365_868_340_224; // (BASE ^ remaining bytes 22 - 12) == 62^10
31const MIN_INC: u64 = 33;
32const MAX_INC: u64 = 333;
33
34/// The number of bytes/characters of a NUID.
35pub const TOTAL_LEN: usize = 22;
36
37static GLOBAL_NUID: Mutex<NUID> = Mutex::new(NUID::new());
38
39/// Generate the next `NUID` string from the global locked `NUID` instance.
40#[allow(clippy::missing_panics_doc)]
41#[must_use]
42pub fn next() -> NUIDStr {
43    GLOBAL_NUID.lock().unwrap().next()
44}
45
46/// NUID needs to be very fast to generate and truly unique, all while being entropy pool friendly.
47/// We will use 12 bytes of crypto generated data (entropy draining), and 10 bytes of sequential data
48/// that is started at a pseudo random number and increments with a pseudo-random increment.
49/// Total is 22 bytes of base 62 ascii text :)
50pub struct NUID {
51    pre: [u8; PRE_LEN],
52    seq: u64,
53    inc: u64,
54}
55
56/// An `NUID` string.
57///
58/// Use [`NUIDStr::as_str`], [`NUIDStr::into_bytes`], the [`Deref`] implementation or
59/// [`ToString`] to access the string.
60pub struct NUIDStr(
61    // INVARIANT: this buffer must always contain a valid utf-8 string
62    [u8; TOTAL_LEN],
63);
64
65impl Default for NUID {
66    fn default() -> Self {
67        Self::new()
68    }
69}
70
71impl NUID {
72    /// generate a new `NUID` and properly initialize the prefix, sequential start, and sequential increment.
73    #[must_use]
74    pub const fn new() -> Self {
75        Self {
76            pre: [0; PRE_LEN],
77            // the first call to `next` will cause the prefix and sequential to be regenerated
78            seq: MAX_SEQ,
79            inc: 0,
80        }
81    }
82
83    pub fn randomize_prefix(&mut self) {
84        let rng = OsRng;
85        for (i, n) in rng.sample_iter(&Alphanumeric).take(PRE_LEN).enumerate() {
86            self.pre[i] = ALPHABET[n as usize % BASE];
87        }
88    }
89
90    /// Generate the next `NUID` string.
91    #[allow(clippy::should_implement_trait)]
92    #[must_use]
93    pub fn next(&mut self) -> NUIDStr {
94        let mut buffer = [0u8; TOTAL_LEN];
95
96        self.seq += self.inc;
97        if self.seq >= MAX_SEQ {
98            self.randomize_prefix();
99            self.reset_sequential();
100        }
101        #[allow(clippy::cast_possible_truncation)]
102        let seq = self.seq as usize;
103
104        for (i, n) in self.pre.iter().enumerate() {
105            buffer[i] = *n;
106        }
107
108        let mut l = seq;
109        for i in (PRE_LEN..TOTAL_LEN).rev() {
110            buffer[i] = ALPHABET[l % BASE];
111            l /= BASE;
112        }
113
114        // `buffer` has been filled with base62 data, which is always valid utf-8
115        NUIDStr(buffer)
116    }
117
118    fn reset_sequential(&mut self) {
119        let mut rng = thread_rng();
120        self.seq = rng.gen_range(0..MAX_SEQ);
121        self.inc = rng.gen_range(MIN_INC..MAX_INC);
122    }
123}
124
125impl NUIDStr {
126    /// Get a reference to the inner string
127    #[must_use]
128    pub fn as_str(&self) -> &str {
129        // SAFETY: the invariant guarantees the buffer to always contain utf-8 characters
130        unsafe { str::from_utf8_unchecked(&self.0) }
131    }
132
133    /// Extract the inner buffer
134    #[must_use]
135    pub fn into_bytes(self) -> [u8; TOTAL_LEN] {
136        self.0
137    }
138}
139
140impl Display for NUIDStr {
141    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142        f.write_str(self.as_str())
143    }
144}
145
146impl Deref for NUIDStr {
147    type Target = str;
148
149    /// Get a reference to the inner string
150    fn deref(&self) -> &Self::Target {
151        self.as_str()
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use std::collections::HashSet;
159
160    #[test]
161    fn alphabet_size() {
162        assert_eq!(ALPHABET.len(), BASE);
163    }
164
165    #[test]
166    fn global_nuid_init() {
167        assert_eq!(GLOBAL_NUID.lock().unwrap().pre.len(), PRE_LEN);
168        assert_ne!(GLOBAL_NUID.lock().unwrap().seq, 0);
169    }
170
171    #[test]
172    fn nuid_rollover() {
173        let mut n = NUID::new();
174        n.seq = MAX_SEQ;
175        let old = n.pre.to_vec();
176        let _ = n.next();
177        assert_ne!(n.pre.to_vec(), old);
178
179        let mut n = NUID::new();
180        n.seq = 1;
181        let old = n.pre.to_vec();
182        let _ = n.next();
183        assert_eq!(n.pre.to_vec(), old);
184    }
185
186    #[test]
187    fn nuid_len() {
188        let id = next();
189        assert_eq!(id.len(), TOTAL_LEN);
190    }
191
192    #[test]
193    fn proper_prefix() {
194        let mut min: u8 = 255;
195        let mut max: u8 = 0;
196
197        for nn in &ALPHABET {
198            let n = *nn;
199            if n < min {
200                min = n;
201            }
202            if n > max {
203                max = n;
204            }
205        }
206
207        for _ in 0..100_000 {
208            let nuid = NUID::new();
209            for j in 0..PRE_LEN {
210                assert!(nuid.pre[j] >= min || nuid.pre[j] <= max);
211            }
212        }
213    }
214
215    #[test]
216    fn unique() {
217        let mut set = HashSet::new();
218        for _ in 0..10_000_000 {
219            assert!(set.insert(next().to_string()));
220        }
221    }
222}