scirs2_core/collections/
rrb_vec.rs

1//! Persistent radix-balanced vector (RRB-tree).
2//!
3//! Implements an immutable persistent vector with:
4//! - O(log N) element access
5//! - O(log N) amortised `push_back`
6//! - O(log N) structural-sharing `update` (returns a new version while the
7//!   original remains intact)
8//!
9//! The branching factor is `BRANCHING = 32` (5-bit trie).
10//!
11//! # Structural Sharing
12//!
13//! All nodes are reference-counted via `std::rc::Rc`.  `update` copies only
14//! the O(log N) nodes on the path to the changed element; all other nodes are
15//! shared with the original tree.
16
17use std::rc::Rc;
18
19/// Branching factor of the trie (must be a power of two).
20const BRANCHING: usize = 32;
21
22/// log2(BRANCHING) — number of bits consumed per trie level.
23const BITS: usize = 5;
24
25/// Bit-mask for selecting a `BITS`-wide index at one trie level.
26const MASK: usize = BRANCHING - 1;
27
28// ─────────────────────────────────────────────────────────────────────────────
29// Internal node
30// ─────────────────────────────────────────────────────────────────────────────
31
32/// A node in the radix-balanced trie.
33#[derive(Clone)]
34enum RrbNode {
35    /// Leaf node holding up to `BRANCHING` elements.
36    Leaf(Vec<u64>),
37    /// Internal node with up to `BRANCHING` child subtrees.
38    Internal(Vec<Rc<RrbNode>>),
39}
40
41impl RrbNode {
42    /// Retrieve the element at `index` within a subtree of the given `height`.
43    ///
44    /// `height == 0` means this node is a [`RrbNode::Leaf`].
45    fn get(&self, index: usize, height: usize) -> Option<u64> {
46        match self {
47            RrbNode::Leaf(data) => data.get(index).copied(),
48            RrbNode::Internal(children) => {
49                let shift = height * BITS;
50                let child_idx = (index >> shift) & MASK;
51                let remainder = index & ((1 << shift) - 1);
52                children
53                    .get(child_idx)
54                    .and_then(|c| c.get(remainder, height - 1))
55            }
56        }
57    }
58
59    /// Return a new node (sharing structure) with `index` updated to `value`.
60    fn update(&self, index: usize, value: u64, height: usize) -> Option<Rc<RrbNode>> {
61        match self {
62            RrbNode::Leaf(data) => {
63                if index >= data.len() {
64                    return None;
65                }
66                let mut new_data = data.clone();
67                new_data[index] = value;
68                Some(Rc::new(RrbNode::Leaf(new_data)))
69            }
70            RrbNode::Internal(children) => {
71                let shift = height * BITS;
72                let child_idx = (index >> shift) & MASK;
73                let remainder = index & ((1 << shift) - 1);
74                let new_child = children
75                    .get(child_idx)
76                    .and_then(|c| c.update(remainder, value, height - 1))?;
77                let mut new_children = children.clone();
78                new_children[child_idx] = new_child;
79                Some(Rc::new(RrbNode::Internal(new_children)))
80            }
81        }
82    }
83
84    /// Collect all elements into `out` in order.
85    fn collect(&self, out: &mut Vec<u64>) {
86        match self {
87            RrbNode::Leaf(data) => out.extend_from_slice(data),
88            RrbNode::Internal(children) => {
89                for c in children {
90                    c.collect(out);
91                }
92            }
93        }
94    }
95}
96
97// ─────────────────────────────────────────────────────────────────────────────
98// RrbVec
99// ─────────────────────────────────────────────────────────────────────────────
100
101/// Persistent immutable vector backed by a 32-way radix-balanced trie.
102///
103/// All modifying operations return a *new* `RrbVec`; the original is
104/// unaffected.  Internal subtrees are reference-counted and shared between
105/// versions wherever possible.
106#[derive(Clone)]
107pub struct RrbVec {
108    root: Option<Rc<RrbNode>>,
109    /// Total number of elements stored.
110    length: usize,
111    /// Height of the trie (0 = root is a leaf, 1 = one level of internal
112    /// nodes above leaves, …).
113    height: usize,
114}
115
116impl Default for RrbVec {
117    fn default() -> Self {
118        Self::new()
119    }
120}
121
122impl std::fmt::Debug for RrbVec {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.debug_struct("RrbVec")
125            .field("length", &self.length)
126            .field("height", &self.height)
127            .finish()
128    }
129}
130
131impl RrbVec {
132    // ── constructors ──────────────────────────────────────────────────────
133
134    /// Create an empty vector.
135    pub fn new() -> Self {
136        RrbVec {
137            root: None,
138            length: 0,
139            height: 0,
140        }
141    }
142
143    /// Build a vector from a slice of `u64` values.
144    pub fn from_slice(data: &[u64]) -> Self {
145        let mut v = RrbVec::new();
146        for &x in data {
147            v = v.push_back(x);
148        }
149        v
150    }
151
152    // ── queries ───────────────────────────────────────────────────────────
153
154    /// Number of elements.
155    pub fn len(&self) -> usize {
156        self.length
157    }
158
159    /// True if the vector contains no elements.
160    pub fn is_empty(&self) -> bool {
161        self.length == 0
162    }
163
164    /// O(log N) element access.
165    pub fn get(&self, index: usize) -> Option<u64> {
166        if index >= self.length {
167            return None;
168        }
169        self.root.as_ref().and_then(|r| r.get(index, self.height))
170    }
171
172    // ── persistent mutations ──────────────────────────────────────────────
173
174    /// Return a new vector with `value` appended to the end.
175    ///
176    /// Amortised O(log N).
177    pub fn push_back(&self, value: u64) -> Self {
178        let new_root = match &self.root {
179            None => {
180                // First element: create a single-element leaf.
181                Rc::new(RrbNode::Leaf(vec![value]))
182            }
183            Some(root) => {
184                // Try to insert into the existing tree.
185                match push_node(root, self.length, self.height, value) {
186                    PushResult::Inserted(new_root) => new_root,
187                    PushResult::NeedsNewRoot(sibling) => {
188                        // The tree is full at the current height — grow by one level.
189                        Rc::new(RrbNode::Internal(vec![Rc::clone(root), sibling]))
190                    }
191                }
192            }
193        };
194        RrbVec {
195            root: Some(new_root),
196            length: self.length + 1,
197            height: if self.root.is_none() {
198                0
199            } else {
200                // Height may have grown inside push_node logic
201                new_height_after_push(self.length + 1, self.height)
202            },
203        }
204    }
205
206    /// Return a new vector with the element at `index` replaced by `value`.
207    ///
208    /// Returns `None` if `index` is out of bounds.  Runs in O(log N) and
209    /// shares structure with the original.
210    pub fn update(&self, index: usize, value: u64) -> Option<Self> {
211        if index >= self.length {
212            return None;
213        }
214        let new_root = self
215            .root
216            .as_ref()
217            .and_then(|r| r.update(index, value, self.height))?;
218        Some(RrbVec {
219            root: Some(new_root),
220            length: self.length,
221            height: self.height,
222        })
223    }
224
225    // ── iteration / conversion ────────────────────────────────────────────
226
227    /// Iterate over elements in order.
228    pub fn iter(&self) -> RrbVecIter<'_> {
229        RrbVecIter { vec: self, pos: 0 }
230    }
231
232    /// Collect all elements into a `Vec<u64>`.
233    pub fn to_vec(&self) -> Vec<u64> {
234        let mut out = Vec::with_capacity(self.length);
235        if let Some(root) = &self.root {
236            root.collect(&mut out);
237        }
238        out
239    }
240}
241
242// ─────────────────────────────────────────────────────────────────────────────
243// Iterator
244// ─────────────────────────────────────────────────────────────────────────────
245
246/// Iterator over an [`RrbVec`].
247pub struct RrbVecIter<'a> {
248    vec: &'a RrbVec,
249    pos: usize,
250}
251
252impl<'a> Iterator for RrbVecIter<'a> {
253    type Item = u64;
254
255    fn next(&mut self) -> Option<Self::Item> {
256        if self.pos >= self.vec.length {
257            return None;
258        }
259        let val = self.vec.get(self.pos);
260        self.pos += 1;
261        val
262    }
263
264    fn size_hint(&self) -> (usize, Option<usize>) {
265        let remaining = self.vec.length - self.pos;
266        (remaining, Some(remaining))
267    }
268}
269
270impl<'a> ExactSizeIterator for RrbVecIter<'a> {}
271
272// ─────────────────────────────────────────────────────────────────────────────
273// Internal push helpers
274// ─────────────────────────────────────────────────────────────────────────────
275
276enum PushResult {
277    /// The value was inserted; the returned node is the updated subtree root.
278    Inserted(Rc<RrbNode>),
279    /// The current subtree was full; a new sibling node carrying the value is
280    /// returned and must be attached at the parent level.
281    NeedsNewRoot(Rc<RrbNode>),
282}
283
284/// Capacity of a full subtree of the given height.
285fn subtree_capacity(height: usize) -> usize {
286    BRANCHING.pow((height + 1) as u32)
287}
288
289/// Recursively insert `value` into the subtree rooted at `node`.
290///
291/// `length` is the current number of elements in the subtree.
292/// `height` is the height of `node` (0 = leaf).
293fn push_node(node: &Rc<RrbNode>, length: usize, height: usize, value: u64) -> PushResult {
294    if height == 0 {
295        // Leaf node
296        match node.as_ref() {
297            RrbNode::Leaf(data) => {
298                if data.len() < BRANCHING {
299                    let mut new_data = data.clone();
300                    new_data.push(value);
301                    PushResult::Inserted(Rc::new(RrbNode::Leaf(new_data)))
302                } else {
303                    // Leaf is full — create a new sibling leaf
304                    PushResult::NeedsNewRoot(Rc::new(RrbNode::Leaf(vec![value])))
305                }
306            }
307            RrbNode::Internal(_) => unreachable!("height 0 must be a leaf"),
308        }
309    } else {
310        // Internal node
311        match node.as_ref() {
312            RrbNode::Internal(children) => {
313                let child_cap = subtree_capacity(height - 1);
314                let last_child_idx = if children.is_empty() {
315                    return PushResult::NeedsNewRoot(make_path(height, value));
316                } else {
317                    children.len() - 1
318                };
319                let last_child_len = length - last_child_idx * child_cap;
320
321                match push_node(&children[last_child_idx], last_child_len, height - 1, value) {
322                    PushResult::Inserted(new_child) => {
323                        let mut new_children = children.clone();
324                        new_children[last_child_idx] = new_child;
325                        PushResult::Inserted(Rc::new(RrbNode::Internal(new_children)))
326                    }
327                    PushResult::NeedsNewRoot(sibling) => {
328                        if children.len() < BRANCHING {
329                            let mut new_children = children.clone();
330                            new_children.push(sibling);
331                            PushResult::Inserted(Rc::new(RrbNode::Internal(new_children)))
332                        } else {
333                            // This internal node is also full
334                            PushResult::NeedsNewRoot(make_path(height, value))
335                        }
336                    }
337                }
338            }
339            RrbNode::Leaf(_) => unreachable!("height > 0 must be internal"),
340        }
341    }
342}
343
344/// Create a leftmost path of internal nodes down to a single-element leaf.
345fn make_path(height: usize, value: u64) -> Rc<RrbNode> {
346    if height == 0 {
347        Rc::new(RrbNode::Leaf(vec![value]))
348    } else {
349        Rc::new(RrbNode::Internal(vec![make_path(height - 1, value)]))
350    }
351}
352
353/// Determine the height of a balanced trie holding `n` elements.
354fn new_height_after_push(n: usize, old_height: usize) -> usize {
355    // A trie of height h can hold at most BRANCHING^(h+1) elements.
356    let mut h = old_height;
357    while subtree_capacity(h) < n {
358        h += 1;
359    }
360    // If the new element triggered a root split, the height increased.
361    h
362}
363
364// ─────────────────────────────────────────────────────────────────────────────
365// Tests
366// ─────────────────────────────────────────────────────────────────────────────
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn test_empty_vec() {
374        let v = RrbVec::new();
375        assert_eq!(v.len(), 0);
376        assert!(v.is_empty());
377        assert_eq!(v.get(0), None);
378    }
379
380    #[test]
381    fn test_push_back_sequential_0_to_99() {
382        let mut v = RrbVec::new();
383        for i in 0..100u64 {
384            v = v.push_back(i);
385        }
386        assert_eq!(v.len(), 100);
387        for i in 0..100u64 {
388            assert_eq!(v.get(i as usize), Some(i), "mismatch at index {}", i);
389        }
390    }
391
392    #[test]
393    fn test_get_after_push() {
394        let v = RrbVec::new().push_back(42).push_back(99).push_back(7);
395        assert_eq!(v.get(0), Some(42));
396        assert_eq!(v.get(1), Some(99));
397        assert_eq!(v.get(2), Some(7));
398        assert_eq!(v.get(3), None);
399    }
400
401    #[test]
402    fn test_update_does_not_modify_original() {
403        let v1 = RrbVec::from_slice(&[10, 20, 30]);
404        let v2 = v1.update(1, 999).expect("update failed");
405        // v1 must be unchanged
406        assert_eq!(v1.get(1), Some(20));
407        // v2 reflects the change
408        assert_eq!(v2.get(1), Some(999));
409        assert_eq!(v2.get(0), Some(10));
410        assert_eq!(v2.get(2), Some(30));
411    }
412
413    #[test]
414    fn test_iter_matches_to_vec() {
415        let data: Vec<u64> = (0..50).collect();
416        let v = RrbVec::from_slice(&data);
417        let from_iter: Vec<u64> = v.iter().collect();
418        assert_eq!(from_iter, v.to_vec());
419        assert_eq!(from_iter, data);
420    }
421
422    #[test]
423    fn test_from_slice_and_iter() {
424        let data = vec![100u64, 200, 300, 400, 500];
425        let v = RrbVec::from_slice(&data);
426        assert_eq!(v.len(), 5);
427        let out: Vec<u64> = v.iter().collect();
428        assert_eq!(out, data);
429    }
430
431    #[test]
432    fn test_large_vector_1000_elements() {
433        let mut v = RrbVec::new();
434        for i in 0..1000u64 {
435            v = v.push_back(i * 3 + 7);
436        }
437        assert_eq!(v.len(), 1000);
438        for i in 0..1000u64 {
439            let expected = i * 3 + 7;
440            assert_eq!(v.get(i as usize), Some(expected), "mismatch at {}", i);
441        }
442    }
443
444    #[test]
445    fn test_structural_sharing_update_creates_new_vec_original_unchanged() {
446        let original = RrbVec::from_slice(&[1, 2, 3, 4, 5]);
447        let updated = original.update(2, 99).expect("update failed");
448        // The updated version has the new value
449        assert_eq!(updated.get(2), Some(99));
450        // The original is entirely unaffected (structural sharing)
451        assert_eq!(original.get(2), Some(3));
452        // All other elements are shared/equal
453        for i in [0, 1, 3, 4] {
454            assert_eq!(original.get(i), updated.get(i), "element {} differs", i);
455        }
456    }
457}
scirs2_core/collections/rrb_vec.rs

scirs2_core/collections/
rrb_vec.rs
