mergle/
lib.rs

1#![no_std]
2
3#[cfg(test)]
4#[macro_use]
5extern crate std;
6
7extern crate alloc;
8
9#[cfg(test)]
10mod test;
11
12use alloc::collections::btree_map::BTreeMap;
13use alloc::rc::Rc;
14use alloc::vec::Vec;
15use core::cell::RefCell;
16use core::cmp::Ordering;
17
18pub use bromberg_sl2::{BrombergHashable, HashMatrix, I};
19
20#[derive(Clone)]
21pub enum PrefixDiff<T> {
22    LessThan,
23    PrefixOf(T),
24    Equal,
25    PrefixedBy(T),
26    GreaterThan,
27}
28
29#[derive(Clone)]
30pub struct MergleNode<T> {
31    pub elem: T,
32    elem_hash: HashMatrix,
33    height: usize,
34    hash: HashMatrix,
35    pub left: Option<Rc<MergleNode<T>>>,
36    pub right: Option<Rc<MergleNode<T>>>,
37}
38
39type MergleDiff<T> = PrefixDiff<Rc<MergleNode<T>>>;
40
41type MemoizedDiffTable<T> = BTreeMap<(HashMatrix, HashMatrix), MergleDiff<T>>;
42
43#[derive(Default)]
44pub struct MemTableRef<T>(Rc<RefCell<MemoizedDiffTable<T>>>);
45
46impl<T> Clone for MemTableRef<T> {
47    fn clone(&self) -> Self {
48        MemTableRef(Rc::clone(&self.0))
49    }
50}
51
52impl<T: Clone> MemTableRef<T> {
53    pub fn new() -> Self {
54        MemTableRef(Rc::new(RefCell::new(BTreeMap::new())))
55    }
56
57    fn insert(&self, a: HashMatrix, b: HashMatrix, r: PrefixDiff<Rc<MergleNode<T>>>) {
58        let mut table = self.0.borrow_mut();
59        if a > b {
60            table.insert((b, a), r.reverse());
61        } else {
62            table.insert((a, b), r.clone());
63        }
64    }
65
66    fn lookup(&self, a: HashMatrix, b: HashMatrix) -> Option<PrefixDiff<Rc<MergleNode<T>>>> {
67        if a == b {
68            Some(PrefixDiff::Equal)
69        } else {
70            let table = self.0.borrow();
71            if a > b {
72                table.get(&(b, a)).map(|r| r.reverse())
73            } else {
74                table.get(&(a, b)).cloned()
75            }
76        }
77    }
78}
79
80impl<T: Clone> PrefixDiff<T> {
81    fn reverse(&self) -> Self {
82        match self {
83            PrefixDiff::LessThan => PrefixDiff::GreaterThan,
84            PrefixDiff::PrefixOf(t) => PrefixDiff::PrefixedBy(t.clone()),
85            PrefixDiff::Equal => PrefixDiff::Equal,
86            PrefixDiff::PrefixedBy(t) => PrefixDiff::PrefixOf(t.clone()),
87            PrefixDiff::GreaterThan => PrefixDiff::LessThan,
88        }
89    }
90
91    fn from_ord(o: Ordering) -> Self {
92        match o {
93            Ordering::Less => PrefixDiff::LessThan,
94            Ordering::Equal => PrefixDiff::Equal,
95            Ordering::Greater => PrefixDiff::GreaterThan,
96        }
97    }
98}
99
100impl<T: BrombergHashable + Ord + Clone> MergleNode<T> {
101    fn new(
102        elem: T,
103        elem_hash: HashMatrix,
104        left: Option<Rc<Self>>,
105        right: Option<Rc<Self>>,
106    ) -> Self {
107        MergleNode {
108            elem,
109            elem_hash,
110            height: usize::max(Self::height(&left), Self::height(&right)) + 1,
111            hash: Self::hash(&left) * elem_hash * Self::hash(&right),
112            left,
113            right,
114        }
115    }
116
117    fn create_node_map(self: &Rc<Self>, map: &mut BTreeMap<HashMatrix, Rc<Self>>) {
118        if let alloc::collections::btree_map::Entry::Vacant(e) = map.entry(self.hash) {
119            e.insert(Rc::clone(self));
120        } else {
121            return;
122        }
123
124        if let Some(ln) = &self.left {
125            MergleNode::create_node_map(ln, map)
126        }
127
128        if let Some(rn) = &self.right {
129            MergleNode::create_node_map(rn, map)
130        }
131    }
132
133    fn singleton(e: T, h: HashMatrix) -> Self {
134        Self::new(e, h, None, None)
135    }
136
137    fn replace_left(&self, subtree: Option<Rc<Self>>) -> Self {
138        Self::new(
139            self.elem.clone(),
140            self.elem_hash,
141            subtree,
142            self.right.clone(),
143        )
144    }
145
146    fn replace_right(&self, subtree: Option<Rc<Self>>) -> Self {
147        Self::new(
148            self.elem.clone(),
149            self.elem_hash,
150            self.left.clone(),
151            subtree,
152        )
153    }
154
155    fn height(t: &Option<Rc<Self>>) -> usize {
156        match t {
157            None => 0,
158            Some(p) => p.height,
159        }
160    }
161
162    fn balance(&self) -> isize {
163        (Self::height(&self.left) as isize) - (Self::height(&self.right) as isize)
164    }
165
166    fn hash(t: &Option<Rc<Self>>) -> HashMatrix {
167        match t {
168            None => I,
169            Some(p) => p.hash,
170        }
171    }
172
173    fn rotate_left(&self) -> Self {
174        let right = self.right.as_ref().unwrap();
175        right.replace_left(Some(Rc::new(self.replace_right(right.left.clone()))))
176    }
177
178    fn rotate_right(&self) -> Self {
179        let left = self.left.as_ref().unwrap();
180        left.replace_right(Some(Rc::new(self.replace_left(left.right.clone()))))
181    }
182
183    fn rebalance(self) -> Self {
184        let b = self.balance();
185        let res = if b > 1 {
186            // left is too heavy.
187            let left = self.left.as_ref().unwrap();
188            if left.balance() < 0 {
189                self.replace_left(Some(Rc::new(left.rotate_left())))
190                    .rotate_right()
191            } else {
192                self.rotate_right()
193            }
194        } else if b < -1 {
195            let right = self.right.as_ref().unwrap();
196            // right is too heavy.
197            if right.balance() > 0 {
198                self.replace_right(Some(Rc::new(right.rotate_right())))
199                    .rotate_left()
200            } else {
201                self.rotate_left()
202            }
203        } else {
204            self.clone()
205        };
206        debug_assert!(isize::abs(Self::balance(&res)) < 2);
207        res
208    }
209
210    fn pop_right(&self) -> (T, HashMatrix, Option<Rc<Self>>) {
211        match &self.right {
212            None => (self.elem.clone(), self.elem_hash, self.left.clone()),
213            Some(t) => {
214                let (v, h, r_res) = t.pop_right();
215                let candidate_res = self.replace_right(r_res);
216                (v, h, Some(Rc::new(candidate_res.rebalance())))
217            }
218        }
219    }
220
221    fn push_left(&self, insertion: T, hash: HashMatrix) -> Self {
222        let left = match &self.left {
223            None => Self::singleton(insertion, hash),
224            Some(t) => t.push_left(insertion, hash),
225        };
226        self.replace_left(Some(Rc::new(left))).rebalance()
227    }
228
229    fn join_left_with_insert(left: &Rc<Self>, t: T, h: HashMatrix, right: &Rc<Self>) -> Self {
230        let t_prime = if Self::height(&right.left) > left.height + 1 {
231            Self::join_left_with_insert(left, t, h, right.left.as_ref().unwrap())
232        } else {
233            Self::new(t, h, Some(left.clone()), right.left.clone())
234        };
235        right.replace_left(Some(Rc::new(t_prime))).rebalance()
236    }
237
238    fn join_right_with_insert(left: &Rc<Self>, t: T, h: HashMatrix, right: &Rc<Self>) -> Self {
239        let t_prime = if Self::height(&left.right) > right.height + 1 {
240            Self::join_right_with_insert(left.right.as_ref().unwrap(), t, h, right)
241        } else {
242            Self::new(t, h, left.right.clone(), Some(right.clone()))
243        };
244
245        left.replace_right(Some(Rc::new(t_prime))).rebalance()
246    }
247
248    fn join_with_insert(
249        left: &Rc<Self>,
250        insertion: T,
251        elem_hash: HashMatrix,
252        right: &Rc<Self>,
253    ) -> Self {
254        let balance = (left.height as isize) - (right.height as isize);
255        if balance > 1 {
256            // left-weighted
257            Self::join_right_with_insert(left, insertion, elem_hash, right)
258        } else if balance < -1 {
259            // right-weighted
260            Self::join_left_with_insert(left, insertion, elem_hash, right)
261        } else {
262            Self::new(
263                insertion,
264                elem_hash,
265                Some(left.clone()),
266                Some(right.clone()),
267            )
268        }
269    }
270
271    fn join(left: &Rc<Self>, right: &Rc<Self>) -> Self {
272        match left.pop_right() {
273            (v, h, None) => right.push_left(v, h),
274            (v, h, Some(new_left)) => Self::join_with_insert(&new_left, v, h, right),
275        }
276    }
277
278    fn elem_plus_right(&self) -> Self {
279        match &self.right {
280            None => MergleNode::singleton(self.elem.clone(), self.elem_hash),
281            Some(r) => r.push_left(self.elem.clone(), self.elem_hash),
282        }
283    }
284
285    fn prefix_diff(&self, other: &Self, table: &MemTableRef<T>) -> PrefixDiff<Rc<Self>> {
286        if let Some(res) = table.lookup(self.hash, other.hash) {
287            return res;
288        }
289        if self.hash == other.hash {
290            return PrefixDiff::Equal;
291        }
292
293        let res = match (self.height, other.height) {
294            (1, 1) => PrefixDiff::from_ord(self.elem.cmp(&other.elem)),
295            (a, b) if a >= b => {
296                let left_subtree = self.left.as_ref().unwrap();
297                match left_subtree.prefix_diff(other, table) {
298                    PrefixDiff::LessThan => PrefixDiff::LessThan,
299                    PrefixDiff::PrefixOf(b_suffix) => {
300                        MergleNode::prefix_diff(&self.elem_plus_right(), &b_suffix, table)
301                    }
302                    PrefixDiff::Equal => PrefixDiff::PrefixedBy(Rc::new(self.elem_plus_right())),
303                    PrefixDiff::PrefixedBy(a_suffix) => PrefixDiff::PrefixedBy(Rc::new(
304                        MergleNode::join(&a_suffix, &Rc::new(self.elem_plus_right())),
305                    )),
306                    PrefixDiff::GreaterThan => PrefixDiff::GreaterThan,
307                }
308            }
309            (_, _) => MergleNode::prefix_diff(other, self, table).reverse(),
310        };
311
312        table.insert(self.hash, other.hash, res.clone());
313        res
314    }
315}
316
317impl<T: BrombergHashable> BrombergHashable for MergleNode<T> {
318    fn bromberg_hash(&self) -> HashMatrix {
319        self.hash
320    }
321}
322
323#[derive(Clone)]
324pub struct Mergle<T> {
325    root: Rc<MergleNode<T>>,
326    table: MemTableRef<T>,
327}
328
329impl<T: BrombergHashable + Clone + Ord> Mergle<T> {
330    pub fn singleton(t: T, table: &MemTableRef<T>) -> Mergle<T> {
331        let hash = t.bromberg_hash();
332        Mergle {
333            root: Rc::new(MergleNode::singleton(t, hash)),
334            table: table.clone(),
335        }
336    }
337
338    #[must_use]
339    pub fn merge(&self, other: &Self) -> Self {
340        Mergle {
341            root: Rc::new(MergleNode::join(&self.root, &other.root)),
342            table: self.table.clone(),
343        }
344    }
345
346    pub fn iter(&self) -> impl Iterator<Item = T> + '_ {
347        use alloc::vec;
348        let stack = vec![StackElem::Node(self.root.clone())];
349        Iter { stack }
350    }
351
352    #[must_use]
353    pub fn pop(&self) -> (T, Option<Mergle<T>>) {
354        let (v, _, n) = self.root.pop_right();
355        (
356            v,
357            n.map(|r| Mergle {
358                root: r,
359                table: self.table.clone(),
360            }),
361        )
362    }
363
364    pub fn node_map(&self) -> BTreeMap<HashMatrix, Rc<MergleNode<T>>> {
365        let mut map = BTreeMap::new();
366        MergleNode::create_node_map(&self.root, &mut map);
367        map
368    }
369}
370
371impl<T: BrombergHashable + Clone + Ord> Mergle<T> {
372    #[must_use]
373    pub fn prefix_cmp(&self, other: &Self) -> PrefixDiff<Mergle<T>> {
374        match self.root.prefix_diff(&*other.root, &self.table) {
375            PrefixDiff::LessThan => PrefixDiff::LessThan,
376            PrefixDiff::PrefixOf(node) => PrefixDiff::PrefixOf(Mergle {
377                root: node.clone(),
378                table: self.table.clone(),
379            }),
380            PrefixDiff::Equal => PrefixDiff::Equal,
381            PrefixDiff::PrefixedBy(node) => PrefixDiff::PrefixedBy(Mergle {
382                root: node.clone(),
383                table: self.table.clone(),
384            }),
385            PrefixDiff::GreaterThan => PrefixDiff::GreaterThan,
386        }
387    }
388}
389
390impl<T: BrombergHashable + Clone + Ord> Ord for Mergle<T> {
391    fn cmp(&self, other: &Self) -> Ordering {
392        match self.prefix_cmp(other) {
393            PrefixDiff::LessThan | PrefixDiff::PrefixOf(_) => Ordering::Less,
394            PrefixDiff::Equal => Ordering::Equal,
395            PrefixDiff::PrefixedBy(_) | PrefixDiff::GreaterThan => Ordering::Greater,
396        }
397    }
398}
399
400impl<T: BrombergHashable + Clone + Ord> PartialOrd for Mergle<T> {
401    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
402        Some(self.cmp(other))
403    }
404}
405
406impl<T: BrombergHashable + Clone + Ord> PartialEq for Mergle<T> {
407    fn eq(&self, other: &Self) -> bool {
408        self.cmp(other) == Ordering::Equal
409    }
410}
411
412impl<T: BrombergHashable + Clone + Ord> Eq for Mergle<T> {}
413
414impl<T: BrombergHashable> BrombergHashable for Mergle<T> {
415    fn bromberg_hash(&self) -> HashMatrix {
416        self.root.bromberg_hash()
417    }
418}
419
420pub struct Iter<T> {
421    stack: Vec<StackElem<T>>,
422}
423
424enum StackElem<T> {
425    // Union type for subtrees and individual elements
426    Elem(T),
427    Node(Rc<MergleNode<T>>),
428}
429
430impl<T: BrombergHashable + Clone + Ord> Iterator for Iter<T> {
431    type Item = T;
432    fn next(&mut self) -> Option<Self::Item> {
433        let mut result = None;
434        while let Some(stack_elem) = self.stack.pop() {
435            match stack_elem {
436                StackElem::Elem(elem) => {
437                    // yield single element
438                    result = Some(elem.clone());
439                    break;
440                }
441                StackElem::Node(node) => {
442                    // unexplored elements are pushed to the stack in the order:
443                    // right, center, left
444                    if let Some(right_tree) = &node.right {
445                        self.stack.push(StackElem::Node(right_tree.clone()));
446                    };
447
448                    if let Some(left_tree) = &node.left {
449                        self.stack.push(StackElem::Elem(node.elem.clone()));
450                        self.stack.push(StackElem::Node(left_tree.clone()));
451                    } else {
452                        // If there are no more left subtrees, yield the central element
453                        result = Some(node.elem.clone());
454                        break;
455                    }
456                }
457            }
458        }
459        result
460    }
461}
mergle/lib.rs

mergle/
lib.rs
