tawasuyu/brahman

1//! Merkle Search Tree (MST).
2//!
3//! Estructura B-árbol probabilística sobre hashes, en la que el "nivel" de
4//! cada clave se deriva determinísticamente de su propio hash (cantidad de
5//! nibbles cero al inicio). Eso da dos propiedades clave:
6//!
7//! * **Independencia del orden de inserción.** El conjunto `{a, b, c}`
8//!   siempre produce el mismo árbol y el mismo `root_hash`, sin importar
9//!   en qué orden se insertaron las claves.
10//! * **Comparación logarítmica.** Dos repositorios pueden saber si tienen
11//!   el mismo conjunto de hashes con un único byte (`root_hash`); y, si
12//!   difieren, descender solo por las ramas con hashes distintos.
13//!
14//! Esta implementación es completa para insert/contains/iter y produce un
15//! `root_hash` Merkle correcto. La operación de `diff` mínima (delta de
16//! sincronización P2P) se construirá encima cuando exista `minga-p2p`.
17
18use crate::cas::ContentHash;
19use blake3::Hasher;
20use std::collections::HashMap;
21use std::sync::OnceLock;
22
23/// Resumen estructural de un nodo interno del MST: nivel al que viven
24/// sus claves, las claves a ese nivel, y el hash de cada uno de sus
25/// hijos (subárboles). Esto es lo que un peer transmite cuando otro le
26/// pregunta por la forma de un subárbol durante una sincronización
27/// recursiva: bandwidth proporcional a la divergencia, no al tamaño.
28#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
29pub struct NodeProbe {
30    pub level: u32,
31    pub keys: Vec<ContentHash>,
32    pub child_hashes: Vec<ContentHash>,
33}
34
35/// Hash canónico del subárbol vacío (el "neutro" del MST). Cualquier
36/// peer puede computarlo localmente sin tocar la red, lo que permite
37/// reconocer ramas vacías en el otro lado sin pedir un probe.
38pub fn empty_subtree_hash() -> ContentHash {
39    static H: OnceLock<ContentHash> = OnceLock::new();
40    *H.get_or_init(|| {
41        let mut h = Hasher::new();
42        h.update(b"E");
43        ContentHash(*h.finalize().as_bytes())
44    })
45}
46
47#[derive(Debug, Default, Clone, PartialEq, Eq)]
48pub struct Mst {
49    root: Subtree,
50}
51
52#[derive(Debug, Default, Clone, PartialEq, Eq)]
53enum Subtree {
54    #[default]
55    Empty,
56    Node(Box<NodeData>),
57}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
60struct NodeData {
61    level: u32,
62    keys: Vec<ContentHash>,
63    children: Vec<Subtree>,
64}
65
66/// Nivel determinístico de un hash: número de nibbles (4 bits) cero al
67/// inicio. Distribución geométrica con base 16, lo que da árbol balanceado
68/// en expectativa con profundidad logarítmica.
69fn level_of(h: &ContentHash) -> u32 {
70    let mut count = 0u32;
71    for &b in &h.0 {
72        if b == 0 {
73            count += 2;
74        } else if b < 0x10 {
75            count += 1;
76            break;
77        } else {
78            break;
79        }
80    }
81    count
82}
83
84impl Mst {
85    pub fn new() -> Self {
86        Self { root: Subtree::Empty }
87    }
88
89    /// Inserta `h`. Devuelve `true` si era una clave nueva.
90    pub fn insert(&mut self, h: ContentHash) -> bool {
91        let l = level_of(&h);
92        let root = std::mem::take(&mut self.root);
93        let (new_root, inserted) = insert_in(root, h, l);
94        self.root = new_root;
95        inserted
96    }
97
98    pub fn contains(&self, h: &ContentHash) -> bool {
99        contains_in(&self.root, h)
100    }
101
102    pub fn len(&self) -> usize {
103        len_of(&self.root)
104    }
105
106    pub fn is_empty(&self) -> bool {
107        matches!(self.root, Subtree::Empty)
108    }
109
110    /// Recorrido in-order: claves emitidas en orden ascendente por hash.
111    pub fn iter(&self) -> Iter<'_> {
112        let mut it = Iter { stack: Vec::new() };
113        it.descend_left(&self.root);
114        it
115    }
116
117    /// Hash Merkle del árbol completo. Dos MSTs con el mismo conjunto de
118    /// claves tienen el mismo `root_hash`, sin importar orden de inserción.
119    pub fn root_hash(&self) -> ContentHash {
120        subtree_hash(&self.root)
121    }
122
123    /// Construye un índice `subtree_hash -> NodeProbe` cubriendo cada
124    /// nodo interno del árbol. Sirve a un peer como tabla de respuestas
125    /// instantáneas a `ProbeReq`s del otro lado: dado un hash que el
126    /// peer recibió de nosotros (en un Hello o un ProbeRes previo),
127    /// podemos reconstituir su `NodeProbe` en `O(1)`.
128    pub fn build_probe_index(&self) -> HashMap<ContentHash, NodeProbe> {
129        let mut idx = HashMap::new();
130        index_subtree(&self.root, &mut idx);
131        idx
132    }
133
134    /// Diferencia simétrica entre `self` y `other`. Devuelve las claves
135    /// que están en `self` pero no en `other`, y viceversa.
136    ///
137    /// Aprovecha la estructura Merkle: cualquier subárbol cuya raíz
138    /// hashee igual entre ambos lados se descarta sin descender. Cuando
139    /// dos nodos comparten nivel y separadores, recurrimos en paralelo
140    /// sobre sus hijos — cada par idéntico se poda por hash. Cuando la
141    /// estructura diverge (niveles distintos o separadores distintos en
142    /// el mismo nivel), enumeramos las claves de ambos y hacemos merge
143    /// ordenado.
144    ///
145    /// El resultado siempre viene ordenado por hash ascendente, lo que
146    /// permite a un peer P2P hacer streaming de los bloques que faltan
147    /// en orden estable y deduplicar mientras los recibe.
148    pub fn diff(&self, other: &Mst) -> MstDiff {
149        let mut d = MstDiff::default();
150        diff_subtrees(&self.root, &other.root, &mut d.only_in_self, &mut d.only_in_other);
151        d
152    }
153}
154
155/// Resultado de comparar dos MSTs. `is_empty()` ⇔ ambos representan el
156/// mismo conjunto.
157#[derive(Debug, Default, Clone, PartialEq, Eq)]
158pub struct MstDiff {
159    pub only_in_self: Vec<ContentHash>,
160    pub only_in_other: Vec<ContentHash>,
161}
162
163impl MstDiff {
164    pub fn is_empty(&self) -> bool {
165        self.only_in_self.is_empty() && self.only_in_other.is_empty()
166    }
167
168    pub fn total(&self) -> usize {
169        self.only_in_self.len() + self.only_in_other.len()
170    }
171}
172
173fn contains_in(t: &Subtree, h: &ContentHash) -> bool {
174    match t {
175        Subtree::Empty => false,
176        Subtree::Node(n) => match n.keys.binary_search(h) {
177            Ok(_) => true,
178            Err(i) => contains_in(&n.children[i], h),
179        },
180    }
181}
182
183fn len_of(t: &Subtree) -> usize {
184    match t {
185        Subtree::Empty => 0,
186        Subtree::Node(n) => n.keys.len() + n.children.iter().map(len_of).sum::<usize>(),
187    }
188}
189
190fn subtree_hash(t: &Subtree) -> ContentHash {
191    let mut h = Hasher::new();
192    match t {
193        Subtree::Empty => {
194            h.update(b"E");
195        }
196        Subtree::Node(n) => {
197            h.update(b"N");
198            h.update(&n.level.to_le_bytes());
199            h.update(&(n.keys.len() as u64).to_le_bytes());
200            for k in &n.keys {
201                h.update(&k.0);
202            }
203            for c in &n.children {
204                h.update(&subtree_hash(c).0);
205            }
206        }
207    }
208    ContentHash(*h.finalize().as_bytes())
209}
210
211/// Inserta `h` (de nivel `l`) en el subárbol `t`. Devuelve el nuevo
212/// subárbol y si fue una inserción real (no duplicado).
213fn insert_in(t: Subtree, h: ContentHash, l: u32) -> (Subtree, bool) {
214    match t {
215        Subtree::Empty => {
216            let node = NodeData {
217                level: l,
218                keys: vec![h],
219                children: vec![Subtree::Empty, Subtree::Empty],
220            };
221            (Subtree::Node(Box::new(node)), true)
222        }
223        Subtree::Node(boxed) => {
224            let n = *boxed;
225            if l > n.level {
226                // Nueva clave de nivel mayor: parte el árbol actual y la
227                // promueve a nueva raíz.
228                let (left, right) = split_at(Subtree::Node(Box::new(n)), &h);
229                let new_root = NodeData {
230                    level: l,
231                    keys: vec![h],
232                    children: vec![left, right],
233                };
234                (Subtree::Node(Box::new(new_root)), true)
235            } else if l == n.level {
236                match n.keys.binary_search(&h) {
237                    Ok(_) => (Subtree::Node(Box::new(n)), false),
238                    Err(i) => {
239                        let NodeData { level, mut keys, mut children } = n;
240                        let middle = std::mem::replace(&mut children[i], Subtree::Empty);
241                        let (left, right) = split_at(middle, &h);
242                        keys.insert(i, h);
243                        children[i] = left;
244                        children.insert(i + 1, right);
245                        (
246                            Subtree::Node(Box::new(NodeData { level, keys, children })),
247                            true,
248                        )
249                    }
250                }
251            } else {
252                // l < n.level: la clave nueva pertenece a un subárbol bajo
253                // el separador correspondiente.
254                let i = match n.keys.binary_search(&h) {
255                    Ok(_) => unreachable!(
256                        "colisión: clave de nivel inferior coincide con separador de nivel superior"
257                    ),
258                    Err(i) => i,
259                };
260                let NodeData { level, keys, mut children } = n;
261                let child = std::mem::replace(&mut children[i], Subtree::Empty);
262                let (new_child, inserted) = insert_in(child, h, l);
263                children[i] = new_child;
264                (
265                    Subtree::Node(Box::new(NodeData { level, keys, children })),
266                    inserted,
267                )
268            }
269        }
270    }
271}
272
273/// Parte `t` en (claves < pivot, claves > pivot). Pre-condición: el nivel
274/// de cada subárbol involucrado es estrictamente menor que el del pivot
275/// (que vive arriba). El pivot mismo no aparece en el resultado.
276fn split_at(t: Subtree, pivot: &ContentHash) -> (Subtree, Subtree) {
277    match t {
278        Subtree::Empty => (Subtree::Empty, Subtree::Empty),
279        Subtree::Node(boxed) => {
280            let n = *boxed;
281            let i = match n.keys.binary_search(pivot) {
282                Ok(_) => unreachable!("pivot coincide con clave de nivel inferior"),
283                Err(i) => i,
284            };
285            let NodeData { level, keys, children } = n;
286
287            let mut left_keys = keys.clone();
288            left_keys.truncate(i);
289            let mut right_keys = keys;
290            right_keys.drain(..i);
291
292            let mut left_children: Vec<Subtree> = Vec::with_capacity(i + 1);
293            let mut right_children: Vec<Subtree> = Vec::with_capacity(level as usize + 1);
294
295            let mut iter = children.into_iter();
296            for _ in 0..i {
297                left_children.push(iter.next().expect("invariante: children > i"));
298            }
299            let middle = iter.next().expect("invariante: existe children[i]");
300            let (l_mid, r_mid) = split_at(middle, pivot);
301            left_children.push(l_mid);
302            right_children.push(r_mid);
303            for c in iter {
304                right_children.push(c);
305            }
306
307            let left = if left_keys.is_empty() {
308                left_children.pop().unwrap_or(Subtree::Empty)
309            } else {
310                Subtree::Node(Box::new(NodeData {
311                    level,
312                    keys: left_keys,
313                    children: left_children,
314                }))
315            };
316            let right = if right_keys.is_empty() {
317                right_children.pop().unwrap_or(Subtree::Empty)
318            } else {
319                Subtree::Node(Box::new(NodeData {
320                    level,
321                    keys: right_keys,
322                    children: right_children,
323                }))
324            };
325            (left, right)
326        }
327    }
328}
329
330fn index_subtree(t: &Subtree, idx: &mut HashMap<ContentHash, NodeProbe>) {
331    if let Subtree::Node(n) = t {
332        let child_hashes: Vec<ContentHash> = n.children.iter().map(subtree_hash).collect();
333        let probe = NodeProbe {
334            level: n.level,
335            keys: n.keys.clone(),
336            child_hashes,
337        };
338        idx.insert(subtree_hash(t), probe);
339        for c in &n.children {
340            index_subtree(c, idx);
341        }
342    }
343}
344
345fn diff_subtrees(
346    t1: &Subtree,
347    t2: &Subtree,
348    only_in_1: &mut Vec<ContentHash>,
349    only_in_2: &mut Vec<ContentHash>,
350) {
351    // Short-circuit por hash Merkle: si los dos subárboles colapsan al
352    // mismo hash de 32 bytes, representan el mismo conjunto. Una sola
353    // comparación poda toda la rama. Aplicado recursivamente, en árboles
354    // mayormente iguales el coste es proporcional a la divergencia, no al
355    // tamaño total.
356    if subtree_hash(t1) == subtree_hash(t2) {
357        return;
358    }
359    match (t1, t2) {
360        (Subtree::Empty, _) => collect_all(t2, only_in_2),
361        (_, Subtree::Empty) => collect_all(t1, only_in_1),
362        (Subtree::Node(n1), Subtree::Node(n2)) => {
363            if n1.level == n2.level && n1.keys == n2.keys {
364                // Mismo nivel y mismos separadores: los hijos se alinean
365                // posicionalmente. Recurrimos en paralelo — cada par
366                // idéntico se podará en su llamada por el hash de Merkle.
367                for (c1, c2) in n1.children.iter().zip(n2.children.iter()) {
368                    diff_subtrees(c1, c2, only_in_1, only_in_2);
369                }
370            } else {
371                // Estructura divergente. Enumeramos ambos lados ordenados
372                // y hacemos merge. Correcto pero sin más poda Merkle: una
373                // futura iteración con `split_at` por cada separador del
374                // nivel mayor recuperaría la poda en el caso desalineado.
375                let mut k1 = Vec::with_capacity(len_of(t1));
376                let mut k2 = Vec::with_capacity(len_of(t2));
377                collect_all(t1, &mut k1);
378                collect_all(t2, &mut k2);
379                merge_diff_sorted(&k1, &k2, only_in_1, only_in_2);
380            }
381        }
382    }
383}
384
385fn collect_all(t: &Subtree, out: &mut Vec<ContentHash>) {
386    if let Subtree::Node(n) = t {
387        for i in 0..n.keys.len() {
388            collect_all(&n.children[i], out);
389            out.push(n.keys[i]);
390        }
391        collect_all(&n.children[n.keys.len()], out);
392    }
393}
394
395fn merge_diff_sorted(
396    a: &[ContentHash],
397    b: &[ContentHash],
398    only_a: &mut Vec<ContentHash>,
399    only_b: &mut Vec<ContentHash>,
400) {
401    let mut i = 0;
402    let mut j = 0;
403    while i < a.len() && j < b.len() {
404        match a[i].cmp(&b[j]) {
405            std::cmp::Ordering::Less => {
406                only_a.push(a[i]);
407                i += 1;
408            }
409            std::cmp::Ordering::Greater => {
410                only_b.push(b[j]);
411                j += 1;
412            }
413            std::cmp::Ordering::Equal => {
414                i += 1;
415                j += 1;
416            }
417        }
418    }
419    only_a.extend_from_slice(&a[i..]);
420    only_b.extend_from_slice(&b[j..]);
421}
422
423pub struct Iter<'a> {
424    /// Cada frame es (nodo, próximo índice de clave a emitir). Cuando se
425    /// pushea un frame, ya descendimos por su hijo izquierdo (children[0]).
426    stack: Vec<(&'a NodeData, usize)>,
427}
428
429impl<'a> Iter<'a> {
430    fn descend_left(&mut self, t: &'a Subtree) {
431        let mut cur = t;
432        while let Subtree::Node(n) = cur {
433            self.stack.push((n.as_ref(), 0));
434            cur = &n.children[0];
435        }
436    }
437}
438
439impl<'a> Iterator for Iter<'a> {
440    type Item = &'a ContentHash;
441
442    fn next(&mut self) -> Option<Self::Item> {
443        loop {
444            let (node, ki) = {
445                let top = self.stack.last()?;
446                (top.0, top.1)
447            };
448            if ki < node.keys.len() {
449                self.stack.last_mut().unwrap().1 = ki + 1;
450                self.descend_left(&node.children[ki + 1]);
451                return Some(&node.keys[ki]);
452            } else {
453                self.stack.pop();
454            }
455        }
456    }
457}