subms_treap/features/
merge_split.rs

1//! Sequence-builder treap with explicit `split(key) -> (left, right)`
2//! and `merge(left, right)` operations.
3//!
4//! Both run in `O(log N)` expected time under the standard treap
5//! rotation invariant (max-heap on priorities, BST on keys). The
6//! pair is the textbook implicit-treap toolkit and is the right shape
7//! for piecewise sequence construction: build two halves separately,
8//! `merge` to splice; or use `split` to chop a range out of a sorted
9//! stream.
10//!
11//! `merge` requires the BST-on-key invariant: every key in `left`
12//! strictly less than every key in `right`. Violating that
13//! precondition panics in debug builds and yields a malformed treap
14//! in release. The standard pairing is `split(t, k) -> (lo, hi)`
15//! then `merge(lo, hi) -> t'` - which is the round-trip identity
16//! the test suite asserts.
17//!
18//! The base `Treap` in `lib.rs` uses an arena layout that doesn't
19//! cheaply support cross-tree splice. This feature lives on its own
20//! pointer-based node type so split/merge stay zero-copy on the
21//! detached subtree.
22
23use std::cmp::Ordering;
24
25type Link<K, V> = Option<Box<Node<K, V>>>;
26
27struct Node<K, V> {
28    key: K,
29    value: V,
30    priority: u64,
31    left: Link<K, V>,
32    right: Link<K, V>,
33}
34
35pub struct SplittableTreap<K, V> {
36    root: Link<K, V>,
37    len: usize,
38    rng_state: u64,
39}
40
41impl<K: Ord, V> SplittableTreap<K, V> {
42    pub fn new(seed: u64) -> Self {
43        Self {
44            root: None,
45            len: 0,
46            rng_state: seed | 1,
47        }
48    }
49
50    pub fn len(&self) -> usize {
51        self.len
52    }
53    pub fn is_empty(&self) -> bool {
54        self.len == 0
55    }
56
57    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
58        let priority = self.next_priority();
59        let (new_root, replaced) = ins(self.root.take(), key, value, priority);
60        self.root = new_root;
61        if replaced.is_none() {
62            self.len += 1;
63        }
64        replaced
65    }
66
67    pub fn get(&self, key: &K) -> Option<&V> {
68        let mut cur = self.root.as_deref();
69        while let Some(node) = cur {
70            match key.cmp(&node.key) {
71                Ordering::Less => cur = node.left.as_deref(),
72                Ordering::Greater => cur = node.right.as_deref(),
73                Ordering::Equal => return Some(&node.value),
74            }
75        }
76        None
77    }
78
79    /// Consume `self` and split into `(left, right)` where every key
80    /// in `left` is strictly less than `pivot` and every key in
81    /// `right` is greater-than-or-equal-to `pivot`.
82    pub fn split(mut self, pivot: &K) -> (Self, Self) {
83        let (l, r) = split_node(self.root.take(), pivot);
84        let l_len = count(&l);
85        let r_len = count(&r);
86        (
87            Self {
88                root: l,
89                len: l_len,
90                rng_state: self.rng_state,
91            },
92            Self {
93                root: r,
94                len: r_len,
95                rng_state: self.rng_state.wrapping_add(1),
96            },
97        )
98    }
99
100    /// Consume `left` and `right` and produce a single treap. Every
101    /// key in `left` must be strictly less than every key in `right`
102    /// or the resulting BST invariant is violated.
103    pub fn merge(left: Self, right: Self) -> Self {
104        if let (Some(l_max), Some(r_min)) = (max_key(&left.root), min_key(&right.root)) {
105            debug_assert!(
106                l_max < r_min,
107                "SplittableTreap::merge precondition violated (left max >= right min)"
108            );
109        }
110        let rng_state = left.rng_state.wrapping_add(right.rng_state) | 1;
111        let len = left.len + right.len;
112        Self {
113            root: merge_nodes(left.root, right.root),
114            len,
115            rng_state,
116        }
117    }
118
119    pub fn collect_in_order(&self) -> Vec<(&K, &V)> {
120        let mut out = Vec::with_capacity(self.len);
121        in_order(self.root.as_deref(), &mut out);
122        out
123    }
124
125    fn next_priority(&mut self) -> u64 {
126        self.rng_state = self
127            .rng_state
128            .wrapping_mul(6364136223846793005)
129            .wrapping_add(1442695040888963407);
130        // SplitMix64 finalizer - decorrelate the priority from the key so
131        // the tree keeps its expected O(log n) height. Mirrors the base
132        // Treap fix.
133        let mut z = self.rng_state;
134        z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
135        z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb);
136        z ^ (z >> 31)
137    }
138}
139
140fn ins<K: Ord, V>(link: Link<K, V>, key: K, value: V, priority: u64) -> (Link<K, V>, Option<V>) {
141    match link {
142        None => (
143            Some(Box::new(Node {
144                key,
145                value,
146                priority,
147                left: None,
148                right: None,
149            })),
150            None,
151        ),
152        Some(mut node) => match key.cmp(&node.key) {
153            Ordering::Equal => {
154                let old = std::mem::replace(&mut node.value, value);
155                (Some(node), Some(old))
156            }
157            Ordering::Less => {
158                let (new_left, replaced) = ins(node.left.take(), key, value, priority);
159                node.left = new_left;
160                let need_rotate =
161                    node.left.as_ref().map(|l| l.priority).unwrap_or(0) > node.priority;
162                let rebuilt = if need_rotate {
163                    rotate_right(node)
164                } else {
165                    node
166                };
167                (Some(rebuilt), replaced)
168            }
169            Ordering::Greater => {
170                let (new_right, replaced) = ins(node.right.take(), key, value, priority);
171                node.right = new_right;
172                let need_rotate =
173                    node.right.as_ref().map(|r| r.priority).unwrap_or(0) > node.priority;
174                let rebuilt = if need_rotate { rotate_left(node) } else { node };
175                (Some(rebuilt), replaced)
176            }
177        },
178    }
179}
180
181fn split_node<K: Ord, V>(link: Link<K, V>, pivot: &K) -> (Link<K, V>, Link<K, V>) {
182    match link {
183        None => (None, None),
184        Some(mut node) => {
185            if &node.key < pivot {
186                let right = node.right.take();
187                let (lo_r, hi) = split_node(right, pivot);
188                node.right = lo_r;
189                (Some(node), hi)
190            } else {
191                let left = node.left.take();
192                let (lo, hi_l) = split_node(left, pivot);
193                node.left = hi_l;
194                (lo, Some(node))
195            }
196        }
197    }
198}
199
200fn merge_nodes<K, V>(left: Link<K, V>, right: Link<K, V>) -> Link<K, V> {
201    match (left, right) {
202        (None, r) => r,
203        (l, None) => l,
204        (Some(mut l), Some(mut r)) => {
205            if l.priority > r.priority {
206                let l_right = l.right.take();
207                l.right = merge_nodes(l_right, Some(r));
208                Some(l)
209            } else {
210                let r_left = r.left.take();
211                r.left = merge_nodes(Some(l), r_left);
212                Some(r)
213            }
214        }
215    }
216}
217
218fn rotate_right<K, V>(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
219    let mut left = node.left.take().expect("rotate_right needs left child");
220    node.left = left.right.take();
221    left.right = Some(node);
222    left
223}
224
225fn rotate_left<K, V>(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
226    let mut right = node.right.take().expect("rotate_left needs right child");
227    node.right = right.left.take();
228    right.left = Some(node);
229    right
230}
231
232fn count<K, V>(link: &Link<K, V>) -> usize {
233    match link {
234        None => 0,
235        Some(node) => 1 + count(&node.left) + count(&node.right),
236    }
237}
238
239fn min_key<K, V>(link: &Link<K, V>) -> Option<&K> {
240    let mut cur = link.as_deref()?;
241    while let Some(l) = cur.left.as_deref() {
242        cur = l;
243    }
244    Some(&cur.key)
245}
246
247fn max_key<K, V>(link: &Link<K, V>) -> Option<&K> {
248    let mut cur = link.as_deref()?;
249    while let Some(r) = cur.right.as_deref() {
250        cur = r;
251    }
252    Some(&cur.key)
253}
254
255fn in_order<'a, K, V>(link: Option<&'a Node<K, V>>, out: &mut Vec<(&'a K, &'a V)>) {
256    if let Some(node) = link {
257        in_order(node.left.as_deref(), out);
258        out.push((&node.key, &node.value));
259        in_order(node.right.as_deref(), out);
260    }
261}
262
263#[cfg(test)]
264#[path = "merge_split_tests.rs"]
265mod tests;
subms_treap/features/merge_split.rs

subms_treap/features/
merge_split.rs
