1//! The branch tree: a recording is not one line but a tree, because "run the
2//! bot from here" keeps the old future rather than truncating it
3//! (`research/replay-timeline.md`, "Branching"). One JSON file per
4//! recording, human-readable on purpose - like `mcp::states`'s `.state`
5//! files, this is exactly the kind of small debug-relevant file this project
6//! keeps as JSON rather than bincode.
7
8use std::path::Path;
9use std::time::{SystemTime, UNIX_EPOCH};
10use std::{fs, io};
11
12use serde::{Deserialize, Serialize};
13
14/// One branch: a run of frames sharing one bot lifetime, forked from its
15/// parent at `fork_frame` (the root has none: frame 0 is the recording's own
16/// header snapshot, shared by everything in the tree).
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
18pub struct BranchMeta {
19    pub id: String,
20    pub parent: Option<String>,
21    /// The absolute frame this branch diverged from its parent at. Zero for
22    /// the root.
23    pub fork_frame: u64,
24    pub created_at_unix: u64,
25    /// The last frame actually appended to this branch's own log so far -
26    /// kept current as recording proceeds, so [`Tree::owner_of`] can tell a
27    /// finished branch's span from one still being written to.
28    pub last_frame: u64,
29}
30
31#[derive(Debug, Clone, Default, Serialize, Deserialize)]
32pub struct Tree {
33    pub branches: Vec<BranchMeta>,
34}
35
36fn unix_now() -> u64 {
37    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)
38}
39
40impl Tree {
41    pub const ROOT: &'static str = "root";
42
43    /// A brand new recording: one branch, the root, forked at frame 0 (the
44    /// header) and with nothing recorded yet.
45    pub fn new() -> Self {
46        Self { branches: vec![BranchMeta { id: Self::ROOT.into(), parent: None, fork_frame: 0, created_at_unix: unix_now(), last_frame: 0 }] }
47    }
48
49    pub fn get(&self, id: &str) -> Option<&BranchMeta> {
50        self.branches.iter().find(|b| b.id == id)
51    }
52
53    pub fn get_mut(&mut self, id: &str) -> Option<&mut BranchMeta> {
54        self.branches.iter_mut().find(|b| b.id == id)
55    }
56
57    /// A fresh branch id forked from `parent` at `fork_frame`: readable, and
58    /// unique even if the same frame is branched from twice.
59    pub fn next_branch_id(&self, parent: &str, fork_frame: u64) -> String {
60        let mut n = 1u32;
61        loop {
62            let id = format!("{parent}-at-{fork_frame}-{n}");
63            if self.get(&id).is_none() {
64                return id;
65            }
66            n += 1;
67        }
68    }
69
70    /// Add a new branch forked from `parent` at `fork_frame`. Panics if
71    /// `parent` is not in this tree - the caller always just resolved it.
72    pub fn fork(&mut self, parent: &str, fork_frame: u64) -> BranchMeta {
73        assert!(self.get(parent).is_some(), "forking from a branch not in this tree: {parent}");
74        let meta = BranchMeta {
75            id: self.next_branch_id(parent, fork_frame),
76            parent: Some(parent.to_owned()),
77            fork_frame,
78            created_at_unix: unix_now(),
79            last_frame: fork_frame,
80        };
81        self.branches.push(meta.clone());
82        meta
83    }
84
85    /// `id`'s ancestry, root first and `id` last - the branches responsible,
86    /// between them, for every frame from 0 up to `id`'s own tip.
87    pub fn chain(&self, id: &str) -> Vec<&BranchMeta> {
88        let mut out = Vec::new();
89        let mut cur = self.get(id);
90        while let Some(b) = cur {
91            out.push(b);
92            cur = b.parent.as_deref().and_then(|p| self.get(p));
93        }
94        out.reverse();
95        out
96    }
97
98    /// The branch in `id`'s ancestry that owns absolute frame `frame`.
99    /// Checked deepest-first: each ancestor's own span is capped at the
100    /// point its CHILD (in this chain) forked from it, not at the
101    /// ancestor's own `last_frame` - an ancestor may have kept recording
102    /// long after that fork (the parent's own future, an unrelated
103    /// timeline), and none of that belongs to `id`'s history. A frame
104    /// exactly at a fork point resolves to the child, which always has its
105    /// own keyframe there (`crate::recorder`). `None` for a frame nobody in
106    /// the chain has reached yet.
107    pub fn owner_of(&self, id: &str, frame: u64) -> Option<&BranchMeta> {
108        let chain = self.chain(id);
109        let mut ceiling = chain.last()?.last_frame;
110        for b in chain.into_iter().rev() {
111            if frame >= b.fork_frame && frame <= ceiling {
112                return Some(b);
113            }
114            ceiling = b.fork_frame;
115        }
116        None
117    }
118
119    pub fn children_of<'a>(&'a self, id: &'a str) -> impl Iterator<Item = &'a BranchMeta> {
120        self.branches.iter().filter(move |b| b.parent.as_deref() == Some(id))
121    }
122
123    pub fn load(path: &Path) -> io::Result<Self> {
124        match fs::read(path) {
125            Ok(bytes) => serde_json::from_slice(&bytes).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)),
126            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(Self::new()),
127            Err(e) => Err(e),
128        }
129    }
130
131    /// Written atomically: a reader (the UI's own scrubber, or a probe) never
132    /// sees a half-written tree.
133    pub fn save(&self, path: &Path) -> io::Result<()> {
134        let json = serde_json::to_vec_pretty(self).expect("Tree always serializes");
135        let tmp = path.with_extension("json.part");
136        fs::write(&tmp, json)?;
137        fs::rename(tmp, path)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn a_fresh_tree_is_one_root_branch_at_frame_zero() {
147        let tree = Tree::new();
148        assert_eq!(tree.branches.len(), 1);
149        let root = tree.get(Tree::ROOT).expect("root");
150        assert_eq!(root.parent, None);
151        assert_eq!(root.fork_frame, 0);
152    }
153
154    #[test]
155    fn forking_makes_a_child_with_a_unique_id() {
156        let mut tree = Tree::new();
157        tree.get_mut(Tree::ROOT).unwrap().last_frame = 1000;
158        let a = tree.fork(Tree::ROOT, 500);
159        let b = tree.fork(Tree::ROOT, 500);
160        assert_ne!(a.id, b.id);
161        assert_eq!(a.parent.as_deref(), Some(Tree::ROOT));
162        assert_eq!(a.fork_frame, 500);
163    }
164
165    #[test]
166    fn chain_is_root_first_and_the_branch_itself_last() {
167        let mut tree = Tree::new();
168        tree.get_mut(Tree::ROOT).unwrap().last_frame = 1000;
169        let child = tree.fork(Tree::ROOT, 500);
170        tree.get_mut(&child.id).unwrap().last_frame = 900;
171        let grandchild = tree.fork(&child.id, 700);
172        let chain = tree.chain(&grandchild.id);
173        assert_eq!(chain.iter().map(|b| b.id.as_str()).collect::<Vec<_>>(), vec![Tree::ROOT, child.id.as_str(), grandchild.id.as_str()]);
174    }
175
176    #[test]
177    fn owner_of_resolves_along_the_chain_and_prefers_the_child_at_a_fork_point() {
178        let mut tree = Tree::new();
179        tree.get_mut(Tree::ROOT).unwrap().last_frame = 1000;
180        let child = tree.fork(Tree::ROOT, 500);
181        tree.get_mut(&child.id).unwrap().last_frame = 800;
182
183        assert_eq!(tree.owner_of(&child.id, 100).map(|b| b.id.as_str()), Some(Tree::ROOT));
184        // Exactly the fork point: the child owns it (its own explicit keyframe).
185        assert_eq!(tree.owner_of(&child.id, 500).map(|b| b.id.as_str()), Some(child.id.as_str()));
186        assert_eq!(tree.owner_of(&child.id, 700).map(|b| b.id.as_str()), Some(child.id.as_str()));
187        // Beyond the child's own tip: nobody owns it yet.
188        assert_eq!(tree.owner_of(&child.id, 900), None);
189        // Queried via the ROOT's own id, root still owns any frame within
190        // its own recorded span, even one a child later forked from - the
191        // child's existence never shrinks what root itself covers.
192        assert_eq!(tree.owner_of(Tree::ROOT, 700).map(|b| b.id.as_str()), Some(Tree::ROOT));
193    }
194
195    #[test]
196    fn branching_never_touches_the_parents_own_last_frame() {
197        // "The old future is kept, not truncated" - forking is read-only on
198        // the parent except for adding a child to the branch list.
199        let mut tree = Tree::new();
200        tree.get_mut(Tree::ROOT).unwrap().last_frame = 1000;
201        let before = tree.get(Tree::ROOT).unwrap().clone();
202        tree.fork(Tree::ROOT, 400);
203        assert_eq!(tree.get(Tree::ROOT).unwrap().last_frame, before.last_frame);
204    }
205
206    #[test]
207    fn a_tree_round_trips_through_a_file() {
208        let mut tree = Tree::new();
209        tree.get_mut(Tree::ROOT).unwrap().last_frame = 42;
210        tree.fork(Tree::ROOT, 10);
211        let dir = std::env::temp_dir().join(format!("jev-replay-tree-test-{}", std::process::id()));
212        let _ = fs::create_dir_all(&dir);
213        let path = dir.join("tree.json");
214        tree.save(&path).expect("save");
215        let back = Tree::load(&path).expect("load");
216        assert_eq!(back.branches, tree.branches);
217        let _ = fs::remove_dir_all(&dir);
218    }
219
220    #[test]
221    fn loading_a_missing_tree_file_gives_a_fresh_root() {
222        let path = std::env::temp_dir().join(format!("jev-replay-tree-missing-{}.json", std::process::id()));
223        let _ = fs::remove_file(&path);
224        let tree = Tree::load(&path).expect("a missing tree is a fresh one, not an error");
225        assert_eq!(tree.branches.len(), 1);
226    }
227}