Skip to main content

manifoldb_graph/store/
id_gen.rs

1//! ID generation strategies for entities and edges.
2//!
3//! This module provides monotonically increasing ID generators that are
4//! thread-safe and persistent across restarts.
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use manifoldb_core::{EdgeId, EntityId};
9
10/// A monotonic ID generator.
11///
12/// Generates unique, monotonically increasing IDs. The generator is thread-safe
13/// and can be shared across threads. IDs start from 1 (0 is reserved for "no ID").
14///
15/// # Persistence
16///
17/// The generator can be initialized with the highest existing ID to resume
18/// after a restart. Use [`IdGenerator::with_start`] to set the starting value.
19///
20/// # Example
21///
22/// ```
23/// use manifoldb_graph::store::IdGenerator;
24///
25/// let gen = IdGenerator::new();
26/// let id1 = gen.next_entity_id();
27/// let id2 = gen.next_entity_id();
28/// assert!(id1.as_u64() < id2.as_u64());
29/// ```
30#[derive(Debug)]
31pub struct IdGenerator {
32    /// The next entity ID to assign.
33    next_entity_id: AtomicU64,
34    /// The next edge ID to assign.
35    next_edge_id: AtomicU64,
36}
37
38impl IdGenerator {
39    /// Create a new ID generator starting from 1.
40    #[must_use]
41    pub const fn new() -> Self {
42        Self { next_entity_id: AtomicU64::new(1), next_edge_id: AtomicU64::new(1) }
43    }
44
45    /// Create an ID generator starting from specific values.
46    ///
47    /// Use this to resume ID generation after loading existing data.
48    /// The provided values should be one greater than the highest existing IDs.
49    ///
50    /// # Arguments
51    ///
52    /// * `entity_start` - The first entity ID to generate
53    /// * `edge_start` - The first edge ID to generate
54    #[must_use]
55    pub const fn with_start(entity_start: u64, edge_start: u64) -> Self {
56        Self {
57            next_entity_id: AtomicU64::new(entity_start),
58            next_edge_id: AtomicU64::new(edge_start),
59        }
60    }
61
62    /// Generate the next entity ID.
63    ///
64    /// This operation is atomic and thread-safe.
65    pub fn next_entity_id(&self) -> EntityId {
66        let id = self.next_entity_id.fetch_add(1, Ordering::Relaxed);
67        EntityId::new(id)
68    }
69
70    /// Generate the next edge ID.
71    ///
72    /// This operation is atomic and thread-safe.
73    pub fn next_edge_id(&self) -> EdgeId {
74        let id = self.next_edge_id.fetch_add(1, Ordering::Relaxed);
75        EdgeId::new(id)
76    }
77
78    /// Get the current entity ID counter value (next ID to be assigned).
79    #[must_use]
80    pub fn current_entity_counter(&self) -> u64 {
81        self.next_entity_id.load(Ordering::Relaxed)
82    }
83
84    /// Get the current edge ID counter value (next ID to be assigned).
85    #[must_use]
86    pub fn current_edge_counter(&self) -> u64 {
87        self.next_edge_id.load(Ordering::Relaxed)
88    }
89
90    /// Reset the entity ID counter to a new value.
91    ///
92    /// This is primarily used for testing or when rebuilding indexes.
93    pub fn reset_entity_counter(&self, value: u64) {
94        self.next_entity_id.store(value, Ordering::Relaxed);
95    }
96
97    /// Reset the edge ID counter to a new value.
98    ///
99    /// This is primarily used for testing or when rebuilding indexes.
100    pub fn reset_edge_counter(&self, value: u64) {
101        self.next_edge_id.store(value, Ordering::Relaxed);
102    }
103}
104
105impl Default for IdGenerator {
106    fn default() -> Self {
107        Self::new()
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn new_generator_starts_at_one() {
117        let gen = IdGenerator::new();
118        assert_eq!(gen.next_entity_id().as_u64(), 1);
119        assert_eq!(gen.next_edge_id().as_u64(), 1);
120    }
121
122    #[test]
123    fn ids_are_monotonically_increasing() {
124        let gen = IdGenerator::new();
125        let ids: Vec<_> = (0..100).map(|_| gen.next_entity_id().as_u64()).collect();
126        for window in ids.windows(2) {
127            assert!(window[0] < window[1]);
128        }
129    }
130
131    #[test]
132    fn with_start_sets_initial_values() {
133        let gen = IdGenerator::with_start(100, 200);
134        assert_eq!(gen.next_entity_id().as_u64(), 100);
135        assert_eq!(gen.next_edge_id().as_u64(), 200);
136        assert_eq!(gen.next_entity_id().as_u64(), 101);
137        assert_eq!(gen.next_edge_id().as_u64(), 201);
138    }
139
140    #[test]
141    fn current_counter_reflects_next_id() {
142        let gen = IdGenerator::new();
143        assert_eq!(gen.current_entity_counter(), 1);
144        gen.next_entity_id();
145        assert_eq!(gen.current_entity_counter(), 2);
146    }
147
148    #[test]
149    fn reset_counter_works() {
150        let gen = IdGenerator::new();
151        gen.next_entity_id();
152        gen.next_entity_id();
153        gen.reset_entity_counter(1000);
154        assert_eq!(gen.next_entity_id().as_u64(), 1000);
155    }
156
157    #[test]
158    fn entity_and_edge_ids_are_independent() {
159        let gen = IdGenerator::new();
160        assert_eq!(gen.next_entity_id().as_u64(), 1);
161        assert_eq!(gen.next_entity_id().as_u64(), 2);
162        assert_eq!(gen.next_edge_id().as_u64(), 1);
163        assert_eq!(gen.next_entity_id().as_u64(), 3);
164        assert_eq!(gen.next_edge_id().as_u64(), 2);
165    }
166}