rag_plusplus_core/trajectory/branch/
operations.rs

1//! Branch Operations and Types
2//!
3//! Core types for branch management in trajectory DAGs.
4
5use std::collections::HashMap;
6use std::time::{SystemTime, UNIX_EPOCH};
7use crate::trajectory::graph::NodeId;
8
9/// Get current Unix timestamp in seconds.
10#[inline]
11fn current_timestamp() -> i64 {
12    SystemTime::now()
13        .duration_since(UNIX_EPOCH)
14        .map(|d| d.as_secs() as i64)
15        .unwrap_or(0)
16}
17
18/// Unique identifier for a branch.
19pub type BranchId = u64;
20
21/// Status of a branch in the state machine.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
23pub enum BranchStatus {
24    /// Currently active branch (being explored)
25    Active,
26    /// Archived branch (not active but preserved)
27    Archived,
28    /// Merged into another branch
29    Merged,
30    /// Recovered from "lost" state
31    Recovered,
32    /// Orphaned (parent was deleted)
33    Orphaned,
34}
35
36impl Default for BranchStatus {
37    fn default() -> Self {
38        Self::Active
39    }
40}
41
42/// A fork point where branching occurred.
43///
44/// Represents a node in the DAG where multiple children exist,
45/// creating a decision point in the trajectory.
46#[derive(Debug, Clone)]
47pub struct ForkPoint {
48    /// The node ID where forking occurred
49    pub node_id: NodeId,
50    /// All child branches at this fork
51    pub children: Vec<BranchId>,
52    /// The selected/active child branch (if any)
53    pub selected_child: Option<BranchId>,
54    /// Timestamp when fork was created
55    pub created_at: i64,
56    /// Depth in the trajectory graph
57    pub depth: u32,
58}
59
60impl ForkPoint {
61    /// Create a new fork point.
62    pub fn new(node_id: NodeId, children: Vec<BranchId>, depth: u32) -> Self {
63        Self {
64            node_id,
65            children,
66            selected_child: None,
67            created_at: current_timestamp(),
68            depth,
69        }
70    }
71
72    /// Check if this fork has multiple paths (is actually a branch point).
73    #[inline]
74    pub fn is_branch_point(&self) -> bool {
75        self.children.len() > 1
76    }
77
78    /// Get the number of child branches.
79    #[inline]
80    pub fn child_count(&self) -> usize {
81        self.children.len()
82    }
83
84    /// Select a child branch as the active path.
85    pub fn select(&mut self, branch_id: BranchId) -> Result<(), BranchError> {
86        if self.children.contains(&branch_id) {
87            self.selected_child = Some(branch_id);
88            Ok(())
89        } else {
90            Err(BranchError::InvalidBranch(branch_id))
91        }
92    }
93}
94
95/// A branch in the trajectory state machine.
96///
97/// Represents a linear path through the DAG from a fork point to a leaf
98/// (or to another fork point). Branches can be split off into independent
99/// state machines and later merged back.
100#[derive(Debug, Clone)]
101pub struct Branch {
102    /// Unique identifier for this branch
103    pub id: BranchId,
104    /// Where this branch diverged from its parent
105    pub fork_point: NodeId,
106    /// Current tip (head) of the branch
107    pub head: NodeId,
108    /// Current status of the branch
109    pub status: BranchStatus,
110    /// Parent branch (if this is a child branch)
111    pub parent_branch: Option<BranchId>,
112    /// Child branches that forked from this one
113    pub child_branches: Vec<BranchId>,
114    /// Nodes that belong exclusively to this branch
115    pub nodes: Vec<NodeId>,
116    /// Metadata associated with this branch
117    pub metadata: HashMap<String, String>,
118    /// Timestamp when branch was created
119    pub created_at: i64,
120    /// Timestamp when branch was last updated
121    pub updated_at: i64,
122    /// Human-readable label (optional)
123    pub label: Option<String>,
124}
125
126impl Branch {
127    /// Create a new branch.
128    pub fn new(id: BranchId, fork_point: NodeId, head: NodeId) -> Self {
129        let now = current_timestamp();
130        Self {
131            id,
132            fork_point,
133            head,
134            status: BranchStatus::Active,
135            parent_branch: None,
136            child_branches: Vec::new(),
137            nodes: vec![head],
138            metadata: HashMap::new(),
139            created_at: now,
140            updated_at: now,
141            label: None,
142        }
143    }
144
145    /// Create a root branch (no fork point).
146    pub fn root(id: BranchId, root_node: NodeId) -> Self {
147        let now = current_timestamp();
148        Self {
149            id,
150            fork_point: root_node,
151            head: root_node,
152            status: BranchStatus::Active,
153            parent_branch: None,
154            child_branches: Vec::new(),
155            nodes: vec![root_node],
156            metadata: HashMap::new(),
157            created_at: now,
158            updated_at: now,
159            label: Some("main".to_string()),
160        }
161    }
162
163    /// Check if this branch is active.
164    #[inline]
165    pub fn is_active(&self) -> bool {
166        self.status == BranchStatus::Active
167    }
168
169    /// Check if this branch is the main/root branch.
170    #[inline]
171    pub fn is_root(&self) -> bool {
172        self.parent_branch.is_none()
173    }
174
175    /// Check if this branch has been merged.
176    #[inline]
177    pub fn is_merged(&self) -> bool {
178        self.status == BranchStatus::Merged
179    }
180
181    /// Archive this branch (preserve but mark inactive).
182    pub fn archive(&mut self) {
183        self.status = BranchStatus::Archived;
184        self.updated_at = current_timestamp();
185    }
186
187    /// Mark this branch as merged.
188    pub fn mark_merged(&mut self) {
189        self.status = BranchStatus::Merged;
190        self.updated_at = current_timestamp();
191    }
192
193    /// Recover this branch (reactivate from archived/orphaned).
194    pub fn recover(&mut self) {
195        self.status = BranchStatus::Recovered;
196        self.updated_at = current_timestamp();
197    }
198
199    /// Add a node to this branch.
200    pub fn add_node(&mut self, node_id: NodeId) {
201        self.nodes.push(node_id);
202        self.head = node_id;
203        self.updated_at = current_timestamp();
204    }
205
206    /// Get the length of this branch (number of nodes).
207    #[inline]
208    pub fn len(&self) -> usize {
209        self.nodes.len()
210    }
211
212    /// Check if branch is empty.
213    #[inline]
214    pub fn is_empty(&self) -> bool {
215        self.nodes.is_empty()
216    }
217
218    /// Set a label for this branch.
219    pub fn set_label(&mut self, label: impl Into<String>) {
220        self.label = Some(label.into());
221        self.updated_at = current_timestamp();
222    }
223
224    /// Add metadata to this branch.
225    pub fn add_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
226        self.metadata.insert(key.into(), value.into());
227        self.updated_at = current_timestamp();
228    }
229}
230
231/// Operations that can be performed on branches.
232///
233/// These operations form the core of the "lost branch" solution,
234/// ported from DLM's StateMachine operations.
235#[derive(Debug, Clone)]
236pub enum BranchOperation {
237    /// Split a branch at a node, creating a new independent branch.
238    /// Ported from DLM `StateMachine.split()`.
239    Split {
240        /// Source branch being split
241        from: BranchId,
242        /// New branch created by the split
243        to: BranchId,
244        /// Node where the split occurred
245        at: NodeId,
246        /// Timestamp of the operation
247        timestamp: i64,
248    },
249    /// Merge a branch back into another.
250    /// Ported from DLM `StateMachine.merge()`.
251    Merge {
252        /// Branch being merged (will become inactive)
253        from: BranchId,
254        /// Branch receiving the merge
255        into: BranchId,
256        /// Node where merge occurred
257        at: NodeId,
258        /// Timestamp of the operation
259        timestamp: i64,
260    },
261    /// Recover a "lost" branch, reactivating it.
262    Recover {
263        /// Branch being recovered
264        branch: BranchId,
265        /// How it was recovered
266        strategy: String,
267        /// Timestamp of the operation
268        timestamp: i64,
269    },
270    /// Traverse from one branch to another.
271    Traverse {
272        /// Branch being left
273        from: BranchId,
274        /// Branch being entered
275        to: BranchId,
276        /// Timestamp of the operation
277        timestamp: i64,
278    },
279    /// Archive a branch (preserve but deactivate).
280    Archive {
281        /// Branch being archived
282        branch: BranchId,
283        /// Reason for archiving
284        reason: Option<String>,
285        /// Timestamp of the operation
286        timestamp: i64,
287    },
288    /// Create a new branch at a fork point.
289    Fork {
290        /// Parent branch
291        parent: BranchId,
292        /// New branch created
293        child: BranchId,
294        /// Fork point node
295        at: NodeId,
296        /// Timestamp of the operation
297        timestamp: i64,
298    },
299}
300
301impl BranchOperation {
302    /// Create a new split operation.
303    pub fn split(from: BranchId, to: BranchId, at: NodeId) -> Self {
304        Self::Split {
305            from,
306            to,
307            at,
308            timestamp: current_timestamp(),
309        }
310    }
311
312    /// Create a new merge operation.
313    pub fn merge(from: BranchId, into: BranchId, at: NodeId) -> Self {
314        Self::Merge {
315            from,
316            into,
317            at,
318            timestamp: current_timestamp(),
319        }
320    }
321
322    /// Create a new recover operation.
323    pub fn recover(branch: BranchId, strategy: impl Into<String>) -> Self {
324        Self::Recover {
325            branch,
326            strategy: strategy.into(),
327            timestamp: current_timestamp(),
328        }
329    }
330
331    /// Create a new traverse operation.
332    pub fn traverse(from: BranchId, to: BranchId) -> Self {
333        Self::Traverse {
334            from,
335            to,
336            timestamp: current_timestamp(),
337        }
338    }
339
340    /// Create a new archive operation.
341    pub fn archive(branch: BranchId, reason: Option<String>) -> Self {
342        Self::Archive {
343            branch,
344            reason,
345            timestamp: current_timestamp(),
346        }
347    }
348
349    /// Create a new fork operation.
350    pub fn fork(parent: BranchId, child: BranchId, at: NodeId) -> Self {
351        Self::Fork {
352            parent,
353            child,
354            at,
355            timestamp: current_timestamp(),
356        }
357    }
358
359    /// Get the timestamp of this operation.
360    pub fn timestamp(&self) -> i64 {
361        match self {
362            Self::Split { timestamp, .. } => *timestamp,
363            Self::Merge { timestamp, .. } => *timestamp,
364            Self::Recover { timestamp, .. } => *timestamp,
365            Self::Traverse { timestamp, .. } => *timestamp,
366            Self::Archive { timestamp, .. } => *timestamp,
367            Self::Fork { timestamp, .. } => *timestamp,
368        }
369    }
370}
371
372/// Errors that can occur during branch operations.
373#[derive(Debug, Clone, PartialEq, Eq)]
374pub enum BranchError {
375    /// Branch with given ID was not found
376    BranchNotFound(BranchId),
377    /// Node with given ID was not found
378    NodeNotFound(NodeId),
379    /// Operation would create a cycle
380    CycleDetected,
381    /// Cannot split the root node
382    CannotSplitRoot,
383    /// Branch is already merged
384    AlreadyMerged(BranchId),
385    /// Invalid branch ID for operation
386    InvalidBranch(BranchId),
387    /// Cannot merge branch with itself
388    SelfMerge,
389    /// Branch has no parent
390    NoParent(BranchId),
391    /// Operation not allowed in current state
392    InvalidState(String),
393}
394
395impl std::fmt::Display for BranchError {
396    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
397        match self {
398            Self::BranchNotFound(id) => write!(f, "Branch {} not found", id),
399            Self::NodeNotFound(id) => write!(f, "Node {} not found", id),
400            Self::CycleDetected => write!(f, "Operation would create a cycle"),
401            Self::CannotSplitRoot => write!(f, "Cannot split the root node"),
402            Self::AlreadyMerged(id) => write!(f, "Branch {} is already merged", id),
403            Self::InvalidBranch(id) => write!(f, "Invalid branch ID: {}", id),
404            Self::SelfMerge => write!(f, "Cannot merge branch with itself"),
405            Self::NoParent(id) => write!(f, "Branch {} has no parent", id),
406            Self::InvalidState(msg) => write!(f, "Invalid state: {}", msg),
407        }
408    }
409}
410
411impl std::error::Error for BranchError {}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416
417    #[test]
418    fn test_branch_creation() {
419        let branch = Branch::new(1, 100, 101);
420        assert_eq!(branch.id, 1);
421        assert_eq!(branch.fork_point, 100);
422        assert_eq!(branch.head, 101);
423        assert!(branch.is_active());
424        // Note: parent_branch is None by default, so is_root() returns true
425        // This is expected for a newly created branch without explicit parent
426        assert!(branch.parent_branch.is_none());
427    }
428
429    #[test]
430    fn test_root_branch() {
431        let branch = Branch::root(0, 1);
432        assert!(branch.is_root());
433        assert_eq!(branch.label, Some("main".to_string()));
434    }
435
436    #[test]
437    fn test_branch_archive() {
438        let mut branch = Branch::new(1, 100, 101);
439        assert!(branch.is_active());
440        branch.archive();
441        assert!(!branch.is_active());
442        assert_eq!(branch.status, BranchStatus::Archived);
443    }
444
445    #[test]
446    fn test_branch_recover() {
447        let mut branch = Branch::new(1, 100, 101);
448        branch.archive();
449        branch.recover();
450        assert_eq!(branch.status, BranchStatus::Recovered);
451    }
452
453    #[test]
454    fn test_fork_point() {
455        let mut fork = ForkPoint::new(100, vec![1, 2, 3], 5);
456        assert!(fork.is_branch_point());
457        assert_eq!(fork.child_count(), 3);
458
459        assert!(fork.select(2).is_ok());
460        assert_eq!(fork.selected_child, Some(2));
461
462        assert!(fork.select(99).is_err());
463    }
464
465    #[test]
466    fn test_branch_operations() {
467        let split_op = BranchOperation::split(1, 2, 100);
468        assert!(split_op.timestamp() > 0);
469
470        let merge_op = BranchOperation::merge(2, 1, 100);
471        assert!(merge_op.timestamp() > 0);
472
473        let recover_op = BranchOperation::recover(1, "manual");
474        assert!(recover_op.timestamp() > 0);
475    }
476
477    #[test]
478    fn test_branch_add_node() {
479        let mut branch = Branch::new(1, 100, 101);
480        assert_eq!(branch.len(), 1);
481
482        branch.add_node(102);
483        assert_eq!(branch.len(), 2);
484        assert_eq!(branch.head, 102);
485        assert!(branch.nodes.contains(&102));
486    }
487
488    #[test]
489    fn test_branch_metadata() {
490        let mut branch = Branch::new(1, 100, 101);
491        branch.add_metadata("source", "regeneration");
492        branch.set_label("experiment-1");
493
494        assert_eq!(branch.metadata.get("source"), Some(&"regeneration".to_string()));
495        assert_eq!(branch.label, Some("experiment-1".to_string()));
496    }
497}
rag_plusplus_core/trajectory/branch/operations.rs

rag_plusplus_core/trajectory/branch/
operations.rs
