burn_std/
id.rs

1//! # Unique Identifiers
2use crate::rand::gen_random;
3
4/// Simple ID generator.
5pub struct IdGenerator {}
6
7impl IdGenerator {
8    /// Generates a new ID.
9    pub fn generate() -> u64 {
10        // Generate a random u64 (18,446,744,073,709,551,615 combinations)
11        let random_bytes: [u8; 8] = gen_random();
12        u64::from_le_bytes(random_bytes)
13    }
14}
15
16pub use cubecl_common::stream_id::StreamId;
17
18use core::hash::{BuildHasher, Hasher};
19
20use alloc::str::FromStr;
21use data_encoding::BASE32_DNSSEC;
22use serde::{Deserialize, Serialize};
23
24// Hashbrown changed its default hasher in 0.15, but there are some issues
25// https://github.com/rust-lang/hashbrown/issues/577
26// Also, `param_serde_deserialize_legacy_uuid` doesn't pass with the default hasher.
27type DefaultHashBuilder = core::hash::BuildHasherDefault<ahash::AHasher>;
28
29/// Unique ID for a parameter of a module.
30#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Serialize, Deserialize)]
31pub struct ParamId {
32    value: u64,
33}
34
35impl From<u64> for ParamId {
36    fn from(value: u64) -> Self {
37        Self { value }
38    }
39}
40
41impl Default for ParamId {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl ParamId {
48    /// Create a new parameter ID.
49    pub fn new() -> Self {
50        Self {
51            value: IdGenerator::generate(),
52        }
53    }
54
55    /// Gets the internal value of the id.
56    pub fn val(&self) -> u64 {
57        self.value
58    }
59}
60
61impl FromStr for ParamId {
62    type Err = &'static str; // Or a custom error type if preferred
63
64    /// Construct a param id from str.
65    ///
66    /// Preserves compatibility with previous formats (6 bytes, 16-byte uuid).
67    ///
68    /// # Returns
69    /// A `Result` containing the `ParamId` when valid.
70    fn from_str(encoded: &str) -> Result<Self, Self::Err> {
71        let u64_id: Option<u64> = match BASE32_DNSSEC.decode(encoded.as_bytes()) {
72            Ok(bytes) => {
73                let mut buffer = [0u8; 8];
74                buffer[..bytes.len()].copy_from_slice(&bytes);
75                Some(u64::from_le_bytes(buffer))
76            }
77            Err(_) => match uuid::Uuid::try_parse(encoded) {
78                // Backward compatibility with uuid parameter identifiers
79                Ok(id) => {
80                    // Hash the 128-bit uuid to 64-bit
81                    // Though not *theoretically* unique, the probability of a collision should be extremely low
82                    let mut hasher = DefaultHashBuilder::default().build_hasher();
83                    // let mut hasher = DefaultHasher::new();
84                    hasher.write(id.as_bytes());
85                    Some(hasher.finish())
86                }
87                Err(_) => None,
88            },
89        };
90        u64_id.map(Self::from).ok_or("Invalid id.")
91    }
92}
93
94impl core::fmt::Display for ParamId {
95    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
96        let encoded = BASE32_DNSSEC.encode(&self.value.to_le_bytes());
97        f.write_str(&encoded)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    use alloc::collections::BTreeSet;
106    use alloc::string::ToString;
107
108    #[cfg(feature = "std")]
109    use dashmap::DashSet; //Concurrent HashMap
110    #[cfg(feature = "std")]
111    use std::{sync::Arc, thread};
112
113    #[test]
114    fn uniqueness_test() {
115        const IDS_CNT: usize = 10_000;
116
117        let mut set: BTreeSet<u64> = BTreeSet::new();
118
119        for _i in 0..IDS_CNT {
120            assert!(set.insert(IdGenerator::generate()));
121        }
122
123        assert_eq!(set.len(), IDS_CNT);
124    }
125
126    #[cfg(feature = "std")]
127    #[test]
128    fn thread_safety_test() {
129        const NUM_THREADS: usize = 10;
130        const NUM_REPEATS: usize = 1_000;
131        const EXPECTED_TOTAL_IDS: usize = NUM_THREADS * NUM_REPEATS;
132
133        let set: Arc<DashSet<u64>> = Arc::new(DashSet::new());
134
135        let mut handles = vec![];
136
137        for _ in 0..NUM_THREADS {
138            let set = set.clone();
139
140            let handle = thread::spawn(move || {
141                for _i in 0..NUM_REPEATS {
142                    assert!(set.insert(IdGenerator::generate()));
143                }
144            });
145            handles.push(handle);
146        }
147
148        for handle in handles {
149            handle.join().unwrap();
150        }
151        assert_eq!(set.len(), EXPECTED_TOTAL_IDS);
152    }
153
154    #[test]
155    fn param_serde_try_deserialize() {
156        let val = ParamId::from(123456u64);
157        let deserialized = ParamId::from_str(&val.to_string()).unwrap();
158        assert_eq!(val, deserialized);
159
160        assert_eq!(ParamId::from_str("invalid_id"), Err("Invalid id."));
161    }
162
163    #[test]
164    fn param_serde_deserialize() {
165        let val = ParamId::from(123456u64);
166        let deserialized = ParamId::from_str(&val.to_string()).unwrap();
167        assert_eq!(val, deserialized);
168    }
169
170    #[test]
171    fn param_serde_deserialize_legacy() {
172        let legacy_val = [45u8; 6];
173        let param_id = ParamId::from_str(&BASE32_DNSSEC.encode(&legacy_val)).unwrap();
174        assert_eq!(param_id.val().to_le_bytes()[0..6], legacy_val);
175        assert_eq!(param_id.val().to_le_bytes()[6..], [0, 0]);
176    }
177
178    #[test]
179    fn param_serde_deserialize_legacy_uuid() {
180        // Ensure support for legacy uuid deserialization and make sure it results in the same output
181        let legacy_id = "30b82c23-788d-4d63-a743-ada258d5f13c";
182        let param_id1 = ParamId::from_str(legacy_id).unwrap();
183        let param_id2 = ParamId::from_str(legacy_id).unwrap();
184        assert_eq!(param_id1, param_id2);
185    }
186
187    #[test]
188    #[should_panic = "Invalid id."]
189    fn param_serde_deserialize_invalid_id() {
190        let invalid_uuid = "30b82c23-788d-4d63-ada258d5f13c";
191        let _ = ParamId::from_str(invalid_uuid).unwrap();
192    }
193}
burn_std/id.rs

burn_std/
id.rs
