gioser

tawasuyu/brahman

Rust 296TOML 105JSON 14Markdown 73751 Definitions

⎇ e3980d005f2096694d54da0f5f3fc76fac8ae914 crates/modules/yachay/yachay-core/src/notebook.rs Rust · 349 lines · 11.4 KB History Who touched what raw

feat(yachay): notebooks reproducibles — yachay-core + demo e3980d005 · sergio ·

1//! El notebook — celdas en orden de presentación + un DAG de dependencias.
2//!
3//! Un notebook tiene dos estructuras a la vez: el **orden de
4//! presentación** (la lista de celdas tal como se leen) y el **DAG de
5//! dependencias** (qué celda necesita el resultado de cuál). La
6//! ejecución sigue el DAG; el digest Merkle certifica que dos corridas
7//! del mismo notebook producen lo mismo.
8
9use std::collections::{BTreeMap, BTreeSet, VecDeque};
10
11use serde::{Deserialize, Serialize};
12
13use crate::cell::{Cell, CellId, CellKind, CellState};
14
15/// Un notebook reproducible.
16#[derive(Debug, Clone, Default, Serialize, Deserialize)]
17pub struct Notebook {
18    /// Celdas en orden de presentación.
19    cells: Vec<Cell>,
20    next_id: CellId,
21}
22
23impl Notebook {
24    pub fn new() -> Self {
25        Self { cells: Vec::new(), next_id: 1 }
26    }
27
28    /// Añade una celda al final, sin dependencias y en estado `Stale`.
29    /// Devuelve su id.
30    pub fn push(&mut self, kind: CellKind, source: impl Into<String>) -> CellId {
31        let id = self.next_id;
32        self.next_id += 1;
33        self.cells.push(Cell {
34            id,
35            kind,
36            source: source.into(),
37            depends_on: Vec::new(),
38            state: CellState::Stale,
39        });
40        id
41    }
42
43    pub fn len(&self) -> usize {
44        self.cells.len()
45    }
46
47    pub fn is_empty(&self) -> bool {
48        self.cells.is_empty()
49    }
50
51    /// Celdas en orden de presentación.
52    pub fn cells(&self) -> &[Cell] {
53        &self.cells
54    }
55
56    pub fn cell(&self, id: CellId) -> Option<&Cell> {
57        self.cells.iter().find(|c| c.id == id)
58    }
59
60    fn cell_mut(&mut self, id: CellId) -> Option<&mut Cell> {
61        self.cells.iter_mut().find(|c| c.id == id)
62    }
63
64    /// `true` si `a` depende —directa o transitivamente— de `b`.
65    fn depends_transitively(&self, a: CellId, b: CellId) -> bool {
66        let mut seen: BTreeSet<CellId> = BTreeSet::new();
67        let mut queue: VecDeque<CellId> = VecDeque::from([a]);
68        while let Some(cur) = queue.pop_front() {
69            let Some(cell) = self.cell(cur) else { continue };
70            for &dep in &cell.depends_on {
71                if dep == b {
72                    return true;
73                }
74                if seen.insert(dep) {
75                    queue.push_back(dep);
76                }
77            }
78        }
79        false
80    }
81
82    /// Declara que `cell` depende de `dep`. Rechaza (devuelve `false`)
83    /// si alguna celda no existe o si la arista crearía un ciclo.
84    pub fn add_dependency(&mut self, cell: CellId, dep: CellId) -> bool {
85        if cell == dep || self.cell(cell).is_none() || self.cell(dep).is_none() {
86            return false;
87        }
88        // Si `dep` ya depende de `cell`, esta arista cerraría un ciclo.
89        if self.depends_transitively(dep, cell) {
90            return false;
91        }
92        let c = self.cell_mut(cell).expect("verificado");
93        if !c.depends_on.contains(&dep) {
94            c.depends_on.push(dep);
95        }
96        true
97    }
98
99    /// Reemplaza la fuente de una celda: la marca `Stale` y propaga la
100    /// obsolescencia a todas sus dependientes. Devuelve los ids marcados
101    /// (sin contar la celda misma). `false` si la celda no existe.
102    pub fn set_source(&mut self, id: CellId, source: impl Into<String>) -> bool {
103        let Some(c) = self.cell_mut(id) else {
104            return false;
105        };
106        c.source = source.into();
107        c.state = CellState::Stale;
108        self.propagate_stale(id);
109        true
110    }
111
112    /// Marca el estado de una celda. `false` si no existe.
113    pub fn set_state(&mut self, id: CellId, state: CellState) -> bool {
114        match self.cell_mut(id) {
115            Some(c) => {
116                c.state = state;
117                true
118            }
119            None => false,
120        }
121    }
122
123    /// Dependientes directos de `id`.
124    pub fn dependents(&self, id: CellId) -> Vec<CellId> {
125        self.cells
126            .iter()
127            .filter(|c| c.depends_on.contains(&id))
128            .map(|c| c.id)
129            .collect()
130    }
131
132    /// Marca `Stale` a todo dependiente transitivo de `id`. Devuelve los
133    /// ids afectados.
134    pub fn propagate_stale(&mut self, id: CellId) -> Vec<CellId> {
135        let mut affected: Vec<CellId> = Vec::new();
136        let mut seen: BTreeSet<CellId> = BTreeSet::from([id]);
137        let mut queue: VecDeque<CellId> = VecDeque::from([id]);
138        while let Some(cur) = queue.pop_front() {
139            for child in self.dependents(cur) {
140                if seen.insert(child) {
141                    if let Some(c) = self.cell_mut(child) {
142                        c.state = CellState::Stale;
143                    }
144                    affected.push(child);
145                    queue.push_back(child);
146                }
147            }
148        }
149        affected
150    }
151
152    /// Orden topológico de ejecución (dependencias antes que
153    /// dependientes). `None` si el DAG tiene un ciclo.
154    pub fn execution_order(&self) -> Option<Vec<CellId>> {
155        let mut indeg: BTreeMap<CellId, usize> =
156            self.cells.iter().map(|c| (c.id, 0usize)).collect();
157        for c in &self.cells {
158            for &dep in &c.depends_on {
159                if self.cell(dep).is_some() {
160                    *indeg.get_mut(&c.id).unwrap() += 1;
161                }
162            }
163        }
164        let mut queue: VecDeque<CellId> =
165            indeg.iter().filter(|(_, &d)| d == 0).map(|(&k, _)| k).collect();
166        let mut order: Vec<CellId> = Vec::with_capacity(self.cells.len());
167        while let Some(u) = queue.pop_front() {
168            order.push(u);
169            for child in self.dependents(u) {
170                if let Some(d) = indeg.get_mut(&child) {
171                    *d -= 1;
172                    if *d == 0 {
173                        queue.push_back(child);
174                    }
175                }
176            }
177        }
178        (order.len() == self.cells.len()).then_some(order)
179    }
180
181    /// Digest Merkle de cada celda: `blake3(content_hash ‖ digests de las
182    /// dependencias)`. Captura la celda y todo su linaje — dos notebooks
183    /// con los mismos digests producen, reproduciblemente, lo mismo.
184    /// `None` si hay un ciclo.
185    fn all_digests(&self) -> Option<BTreeMap<CellId, [u8; 32]>> {
186        let order = self.execution_order()?;
187        let mut digests: BTreeMap<CellId, [u8; 32]> = BTreeMap::new();
188        for id in order {
189            let cell = self.cell(id).expect("del orden");
190            let mut h = blake3::Hasher::new();
191            h.update(&cell.content_hash());
192            // Dependencias ordenadas → el digest no depende del orden de
193            // declaración.
194            let mut deps = cell.depends_on.clone();
195            deps.sort_unstable();
196            for dep in deps {
197                if let Some(d) = digests.get(&dep) {
198                    h.update(d);
199                }
200            }
201            digests.insert(id, *h.finalize().as_bytes());
202        }
203        Some(digests)
204    }
205
206    /// Digest reproducible de una celda concreta.
207    pub fn digest(&self, id: CellId) -> Option<[u8; 32]> {
208        self.all_digests()?.get(&id).copied()
209    }
210
211    /// Digest reproducible del notebook entero: `blake3` de los digests
212    /// de todas las celdas en orden de id. Dos notebooks con el mismo
213    /// digest son reproduciblemente equivalentes. `None` si hay ciclo.
214    pub fn notebook_digest(&self) -> Option<[u8; 32]> {
215        let digests = self.all_digests()?;
216        let mut h = blake3::Hasher::new();
217        for d in digests.values() {
218            h.update(d);
219        }
220        Some(*h.finalize().as_bytes())
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    fn code(nb: &mut Notebook, src: &str) -> CellId {
229        nb.push(CellKind::Code { language: "rust".into() }, src)
230    }
231
232    /// Notebook a → b → c (cada uno depende del anterior).
233    fn chain() -> (Notebook, CellId, CellId, CellId) {
234        let mut nb = Notebook::new();
235        let a = code(&mut nb, "let x = 1;");
236        let b = code(&mut nb, "let y = x + 1;");
237        let c = code(&mut nb, "println!(\"{y}\");");
238        nb.add_dependency(b, a);
239        nb.add_dependency(c, b);
240        (nb, a, b, c)
241    }
242
243    #[test]
244    fn push_keeps_display_order() {
245        let (nb, a, b, c) = chain();
246        let ids: Vec<_> = nb.cells().iter().map(|x| x.id).collect();
247        assert_eq!(ids, vec![a, b, c]);
248    }
249
250    #[test]
251    fn execution_order_respects_dependencies() {
252        let (nb, a, b, c) = chain();
253        assert_eq!(nb.execution_order(), Some(vec![a, b, c]));
254    }
255
256    #[test]
257    fn add_dependency_rejects_cycles() {
258        let (mut nb, a, _b, c) = chain();
259        // a depender de c cerraría el ciclo a→b→c→a.
260        assert!(!nb.add_dependency(a, c));
261        assert!(nb.execution_order().is_some());
262    }
263
264    #[test]
265    fn add_dependency_rejects_self_and_missing() {
266        let (mut nb, a, ..) = chain();
267        assert!(!nb.add_dependency(a, a));
268        assert!(!nb.add_dependency(a, 999));
269    }
270
271    #[test]
272    fn editing_a_cell_propagates_staleness() {
273        let (mut nb, a, b, c) = chain();
274        for id in [a, b, c] {
275            nb.set_state(id, CellState::Fresh);
276        }
277        nb.set_source(a, "let x = 42;");
278        // La celda editada y sus descendientes quedan Stale.
279        assert_eq!(nb.cell(a).unwrap().state, CellState::Stale);
280        assert_eq!(nb.cell(b).unwrap().state, CellState::Stale);
281        assert_eq!(nb.cell(c).unwrap().state, CellState::Stale);
282    }
283
284    #[test]
285    fn editing_a_leaf_does_not_stale_its_ancestors() {
286        let (mut nb, a, b, c) = chain();
287        for id in [a, b, c] {
288            nb.set_state(id, CellState::Fresh);
289        }
290        nb.set_source(c, "println!(\"fin\");");
291        assert_eq!(nb.cell(a).unwrap().state, CellState::Fresh);
292        assert_eq!(nb.cell(b).unwrap().state, CellState::Fresh);
293        assert_eq!(nb.cell(c).unwrap().state, CellState::Stale);
294    }
295
296    #[test]
297    fn notebook_digest_is_stable_across_calls() {
298        let (nb, ..) = chain();
299        assert_eq!(nb.notebook_digest(), nb.notebook_digest());
300    }
301
302    #[test]
303    fn editing_a_source_changes_the_digest() {
304        let (mut nb, a, ..) = chain();
305        let before = nb.notebook_digest();
306        nb.set_source(a, "let x = 999;");
307        assert_ne!(before, nb.notebook_digest());
308    }
309
310    #[test]
311    fn cell_digest_reflects_upstream_changes() {
312        // Cambiar `a` cambia el digest de `c` (su descendiente).
313        let (mut nb, a, _b, c) = chain();
314        let c_before = nb.digest(c);
315        nb.set_source(a, "let x = 7;");
316        assert_ne!(c_before, nb.digest(c));
317    }
318
319    #[test]
320    fn dependency_order_does_not_affect_digest() {
321        // Dos celdas con las mismas dos dependencias, declaradas en
322        // distinto orden, dan el mismo digest.
323        let mut x = Notebook::new();
324        let xa = code(&mut x, "a");
325        let xb = code(&mut x, "b");
326        let xc = code(&mut x, "c");
327        x.add_dependency(xc, xa);
328        x.add_dependency(xc, xb);
329
330        let mut y = Notebook::new();
331        let ya = code(&mut y, "a");
332        let yb = code(&mut y, "b");
333        let yc = code(&mut y, "c");
334        y.add_dependency(yc, yb);
335        y.add_dependency(yc, ya);
336
337        assert_eq!(x.digest(xc), y.digest(yc));
338    }
339
340    #[test]
341    fn embed_cells_carry_their_module() {
342        let mut nb = Notebook::new();
343        let id = nb.push(CellKind::Embed { module: "dominium".into() }, "preset: caos");
344        assert!(matches!(
345            &nb.cell(id).unwrap().kind,
346            CellKind::Embed { module } if module == "dominium"
347        ));
348    }
349}