radiate_gp/collections/trees/
node.rs

1use super::TreeIterator;
2use crate::{Arity, Factory, NodeStore, NodeType, Tree, node::Node};
3use radiate_core::genome::{Gene, Valid};
4#[cfg(feature = "serde")]
5use serde::{Deserialize, Serialize};
6use std::{fmt::Debug, hash::Hash};
7
8/// A node in a tree structure that represents a single element with optional children.
9///
10/// The [TreeNode] struct is a fundamental building block for tree-based genetic programming in Radiate.
11/// It represents a node in a tree that can have zero or more child nodes, forming a hierarchical structure.
12/// Each node has a value of type T and maintains an optional list of child nodes.
13///
14/// # Type Parameters
15/// * `T` - The type of value stored in the node. This type must implement `Clone`, `PartialEq`, and other traits
16///         required by the genetic programming operations.
17///
18/// # Fields
19/// * `value` - The actual value stored in the node
20/// * `arity` - Optional Arity that specifies how many children the node can have
21/// * `children` - Optional vector of child nodes
22///
23/// # Examples
24/// ```
25/// use radiate_gp::{collections::{TreeNode}, Arity, Node};
26///
27/// // Create a new node with value 42
28/// let node = TreeNode::new(42);
29///
30/// // Create a node with specific arity
31/// let node_with_arity = TreeNode::with_arity(42, Arity::Exact(2));
32/// let other_node_with_arity = TreeNode::from((42, Arity::Exact(2)));
33///
34/// assert_eq!(node_with_arity.arity(), other_node_with_arity.arity());
35///
36/// // Create a node with children
37/// let node_with_children = TreeNode::with_children(42, vec![
38///     TreeNode::new(1),
39///     TreeNode::new(2)
40/// ]);
41/// let other_node_with_children = TreeNode::from((42, vec![
42///     TreeNode::new(1),
43///     TreeNode::new(2),
44/// ]));
45/// ```
46///
47/// # Node Types and [Arity]
48/// The node's type and arity determine its behavior and validity:
49/// * `Leaf` nodes have no children (arity is [Arity::Zero])
50/// * `Vertex` nodes can have any number of children (arity is [Arity::Any])
51/// * `Root` nodes are the starting point of the tree and can have any number of children
52///
53/// # Tree Operations
54/// The struct provides several methods for tree manipulation:
55/// * `new()` - Creates a new node with no children
56/// * `with_arity()` - Creates a node with a specific arity
57/// * `with_children()` - Creates a node with a list of children
58/// * `add_child()` - Adds a child to the node
59/// * `attach()` - Attaches a child and returns self for method chaining
60/// * `detach()` - Removes a child at a specific index
61/// * `swap_subtrees()` - Swaps subtrees between two nodes
62///
63/// # Tree Traversal
64/// The struct implements the [TreeIterator] trait, providing three traversal methods:
65/// * `iter_pre_order()` - Traverses the tree in pre-order (root, then children)
66/// * `iter_post_order()` - Traverses the tree in post-order (children, then root)
67/// * `iter_breadth_first()` - Traverses the tree level by level
68///
69/// # Tree Properties
70/// The struct provides methods to query tree properties:
71/// * `is_leaf()` - Checks if the node has no children - must have [Arity::Zero]
72/// * `size()` - Returns the total number of nodes in the subtree
73/// * `height()` - Returns the height of the subtree
74///
75/// # Validity
76/// A node is considered valid based on its arity:
77/// * Nodes with [Arity::Zero] must have no children
78/// * Nodes with [variant@Arity::Exact] must have exactly n children
79/// * Nodes with [Arity::Any] can have any number of children
80///
81/// # Implementation Details
82/// The struct implements several traits:
83/// * `Node` - Provides common node behavior and access to value and type information
84/// * `Gene` - Enables genetic operations for the node making it compatible with genetic algorithms
85/// * `Valid` - Defines validity rules for the node
86/// * `Debug` - Provides debug formatting
87/// * `Clone`, `PartialEq` - Required for genetic programming operations
88/// * `Format` - Provides pretty-printing of the tree structure
89///
90/// # Evaluation
91/// When `T` implements the `Eval` trait, the node can be evaluated with input data:
92/// ```rust
93/// use radiate_gp::{Op, Eval, TreeNode};
94///
95/// let tree = TreeNode::new(Op::add())
96///     .attach(TreeNode::new(Op::constant(2.0)))
97///     .attach(TreeNode::new(Op::constant(3.0)));
98///
99/// let result = tree.eval(&[]); // Evaluates to 5.0
100/// ```
101#[derive(PartialEq)]
102#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
103pub struct TreeNode<T> {
104    value: T,
105    arity: Option<Arity>,
106    children: Option<Vec<TreeNode<T>>>,
107}
108
109impl<T> TreeNode<T> {
110    pub fn new(val: T) -> Self {
111        TreeNode {
112            value: val,
113            arity: None,
114            children: None,
115        }
116    }
117
118    pub fn with_arity(val: T, arity: Arity) -> Self {
119        TreeNode {
120            value: val,
121            arity: Some(arity),
122            children: None,
123        }
124    }
125
126    pub fn with_children<N>(val: T, children: Vec<N>) -> Self
127    where
128        N: Into<TreeNode<T>>,
129    {
130        TreeNode {
131            value: val,
132            arity: None,
133            children: Some(children.into_iter().map(|n| n.into()).collect()),
134        }
135    }
136
137    pub fn is_leaf(&self) -> bool {
138        self.children.is_none()
139    }
140
141    pub fn add_child(&mut self, child: impl Into<TreeNode<T>>) {
142        let node = child.into();
143        if let Some(children) = self.children.as_mut() {
144            children.push(node);
145        } else {
146            self.children = Some(vec![node]);
147        }
148    }
149
150    pub fn attach(mut self, other: impl Into<TreeNode<T>>) -> Self {
151        self.add_child(other);
152        self
153    }
154
155    pub fn detach(&mut self, index: usize) -> Option<TreeNode<T>> {
156        if let Some(children) = self.children.as_mut() {
157            if index < children.len() {
158                return Some(children.remove(index));
159            }
160        }
161
162        None
163    }
164
165    pub fn children(&self) -> Option<&[TreeNode<T>]> {
166        self.children.as_ref().map(|children| children.as_slice())
167    }
168
169    pub fn children_mut(&mut self) -> Option<&mut Vec<TreeNode<T>>> {
170        self.children.as_mut()
171    }
172
173    pub fn take_children(&mut self) -> Option<Vec<TreeNode<T>>> {
174        self.children.take()
175    }
176
177    #[inline]
178    pub fn size(&self) -> usize {
179        if let Some(children) = self.children.as_ref() {
180            children.iter().fold(1, |acc, child| acc + child.size())
181        } else {
182            1
183        }
184    }
185
186    #[inline]
187    pub fn height(&self) -> usize {
188        if let Some(children) = self.children.as_ref() {
189            1 + children
190                .iter()
191                .map(|child| child.height())
192                .max()
193                .unwrap_or(0)
194        } else {
195            0
196        }
197    }
198
199    #[inline]
200    pub fn get_mut(&mut self, index: usize) -> Option<&mut TreeNode<T>> {
201        let mut cur = 0;
202        Self::get_mut_preorder(self, index, &mut cur)
203    }
204
205    #[inline]
206    fn get_mut_preorder<'a>(
207        node: &'a mut TreeNode<T>,
208        target: usize,
209        cur: &mut usize,
210    ) -> Option<&'a mut TreeNode<T>> {
211        if *cur == target {
212            return Some(node);
213        }
214
215        if let Some(children) = node.children_mut() {
216            for child in children {
217                *cur += 1;
218                if let Some(found) = Self::get_mut_preorder(child, target, cur) {
219                    return Some(found);
220                }
221            }
222        }
223
224        None
225    }
226}
227
228impl<T> Node for TreeNode<T> {
229    type Value = T;
230
231    fn value(&self) -> &Self::Value {
232        &self.value
233    }
234
235    fn value_mut(&mut self) -> &mut Self::Value {
236        &mut self.value
237    }
238
239    fn node_type(&self) -> NodeType {
240        if self.children.is_some() {
241            NodeType::Vertex
242        } else {
243            NodeType::Leaf
244        }
245    }
246
247    fn arity(&self) -> Arity {
248        if let Some(arity) = self.arity {
249            arity
250        } else if let Some(children) = self.children.as_ref() {
251            Arity::Exact(children.len())
252        } else {
253            match self.node_type() {
254                NodeType::Leaf => Arity::Zero,
255                NodeType::Vertex => Arity::Any,
256                NodeType::Root => Arity::Any,
257                _ => Arity::Zero,
258            }
259        }
260    }
261}
262
263impl<T> Gene for TreeNode<T>
264where
265    T: Clone + PartialEq,
266{
267    type Allele = T;
268
269    fn allele(&self) -> &Self::Allele {
270        &self.value
271    }
272
273    fn allele_mut(&mut self) -> &mut Self::Allele {
274        &mut self.value
275    }
276
277    fn new_instance(&self) -> Self {
278        TreeNode {
279            value: self.value.clone(),
280            arity: self.arity,
281            children: self.children.as_ref().map(|children| {
282                children
283                    .iter()
284                    .map(|child| child.new_instance())
285                    .collect::<Vec<TreeNode<T>>>()
286            }),
287        }
288    }
289
290    fn with_allele(&self, allele: &Self::Allele) -> Self {
291        TreeNode {
292            value: allele.clone(),
293            arity: self.arity,
294            children: self.children.as_ref().map(|children| children.to_vec()),
295        }
296    }
297}
298
299impl<T> Valid for TreeNode<T> {
300    fn is_valid(&self) -> bool {
301        for node in self.iter_breadth_first() {
302            match node.arity() {
303                Arity::Zero => return node.children.is_none(),
304                Arity::Exact(n) => {
305                    if node.children.is_none() || node.children.as_ref().unwrap().len() != n {
306                        return false;
307                    }
308                }
309                Arity::Any => {}
310            }
311        }
312
313        true
314    }
315}
316
317impl<T> Factory<(usize, Option<NodeStore<T>>), Option<TreeNode<T>>> for TreeNode<T>
318where
319    T: Clone + Default,
320{
321    fn new_instance(&self, (index, store): (usize, Option<NodeStore<T>>)) -> Option<TreeNode<T>> {
322        store
323            .map(|store| Tree::with_depth(index, store).take_root())
324            .flatten()
325    }
326}
327
328impl<T: Clone> Clone for TreeNode<T> {
329    fn clone(&self) -> Self {
330        TreeNode {
331            value: self.value.clone(),
332            arity: self.arity,
333            children: self.children.as_ref().map(|children| children.to_vec()),
334        }
335    }
336}
337
338impl<T: Default> Default for TreeNode<T> {
339    fn default() -> Self {
340        TreeNode {
341            value: T::default(),
342            arity: None,
343            children: None,
344        }
345    }
346}
347
348impl<T: Hash> Hash for TreeNode<T> {
349    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
350        self.value.hash(state);
351        self.arity.hash(state);
352        if let Some(children) = self.children.as_ref() {
353            for child in children {
354                child.hash(state);
355            }
356        }
357    }
358}
359
360impl<T: Debug> Debug for TreeNode<T> {
361    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
362        write!(
363            f,
364            "{:>10?} :: {:<10} {:<12} C: {:?}",
365            format!("{:?}", self.node_type())[..3].to_owned(),
366            self.arity(),
367            format!("{:?}", self.value).to_owned(),
368            match &self.children {
369                Some(children) => children.len(),
370                None => 0,
371            },
372        )
373    }
374}
375
376impl<T> From<(T, Arity)> for TreeNode<T> {
377    fn from(value: (T, Arity)) -> Self {
378        TreeNode::with_arity(value.0, value.1)
379    }
380}
381
382impl<T> From<(T, Vec<TreeNode<T>>)> for TreeNode<T> {
383    fn from(value: (T, Vec<TreeNode<T>>)) -> Self {
384        TreeNode::with_children(value.0, value.1)
385    }
386}
387
388macro_rules! impl_from {
389    ($($t:ty),+) => {
390        $(
391            impl From<$t> for TreeNode<$t> {
392                fn from(value: $t) -> Self {
393                    TreeNode::new(value)
394                }
395            }
396        )+
397    };
398}
399
400impl_from!(
401    u8,
402    u16,
403    u32,
404    u64,
405    u128,
406    i8,
407    i16,
408    i32,
409    i64,
410    i128,
411    f32,
412    f64,
413    String,
414    bool,
415    char,
416    usize,
417    isize,
418    &'static str,
419    ()
420);
421
422impl<T> From<TreeNode<T>> for Vec<TreeNode<T>> {
423    fn from(node: TreeNode<T>) -> Self {
424        vec![node]
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use crate::{Arity, Eval, NodeType, Op};
432
433    #[test]
434    fn test_node_creation() {
435        let node = TreeNode::new(42);
436
437        assert_eq!(node.value(), &42);
438        assert!(node.is_leaf());
439        assert_eq!(node.arity(), Arity::Zero);
440        assert_eq!(node.node_type(), NodeType::Leaf);
441
442        let node = TreeNode::with_arity(42, Arity::Exact(2));
443
444        assert_eq!(node.value(), &42);
445        assert!(node.is_leaf());
446        assert_eq!(node.arity(), Arity::Exact(2));
447        assert_eq!(node.node_type(), NodeType::Leaf);
448
449        let node = TreeNode::with_children(42, vec![TreeNode::new(1), TreeNode::new(2)]);
450
451        assert_eq!(node.value(), &42);
452        assert!(!node.is_leaf());
453        assert_eq!(node.arity(), Arity::Exact(2));
454        assert_eq!(node.node_type(), NodeType::Vertex);
455    }
456
457    #[test]
458    fn test_node_manipulation() {
459        let mut node = TreeNode::new(42);
460        node.add_child(TreeNode::new(1));
461
462        assert!(!node.is_leaf());
463        assert_eq!(node.children().unwrap().len(), 1);
464        assert_eq!(node.children().unwrap()[0].value(), &1);
465
466        let node = TreeNode::new(42)
467            .attach(TreeNode::new(1))
468            .attach(TreeNode::new(2));
469
470        assert_eq!(node.children().unwrap().len(), 2);
471        assert_eq!(node.children().unwrap()[0].value(), &1);
472        assert_eq!(node.children().unwrap()[1].value(), &2);
473
474        let mut node = TreeNode::with_children(42, vec![TreeNode::new(1), TreeNode::new(2)]);
475        let detached = node.detach(0);
476
477        assert!(detached.is_some());
478        assert_eq!(detached.unwrap().value(), &1);
479        assert_eq!(node.children().unwrap().len(), 1);
480        assert_eq!(node.children().unwrap()[0].value(), &2);
481
482        assert!(node.detach(5).is_none());
483    }
484
485    #[test]
486    fn test_tree_properties() {
487        let node = TreeNode::new(42).attach(TreeNode::new(1)).attach(
488            TreeNode::new(2)
489                .attach(TreeNode::new(3))
490                .attach(TreeNode::new(4)),
491        );
492
493        assert_eq!(node.size(), 5);
494        assert_eq!(node.height(), 2);
495
496        let leaf = TreeNode::new(42);
497
498        assert_eq!(leaf.size(), 1);
499        assert_eq!(leaf.height(), 0);
500    }
501
502    #[test]
503    fn test_tree_traversal() {
504        // Create a tree:
505        //       42
506        //      /  \
507        //     1    2
508        //         / \
509        //        3   4
510        let node = TreeNode::new(42).attach(TreeNode::new(1)).attach(
511            TreeNode::new(2)
512                .attach(TreeNode::new(3))
513                .attach(TreeNode::new(4)),
514        );
515
516        let pre_order: Vec<i32> = node.iter_pre_order().map(|n| *n.value()).collect();
517        assert_eq!(pre_order, vec![42, 1, 2, 3, 4]);
518
519        let post_order: Vec<i32> = node.iter_post_order().map(|n| *n.value()).collect();
520        assert_eq!(post_order, vec![1, 3, 4, 2, 42]);
521
522        let bfs: Vec<i32> = node.iter_breadth_first().map(|n| *n.value()).collect();
523        assert_eq!(bfs, vec![42, 1, 2, 3, 4]);
524    }
525
526    #[test]
527    fn test_node_validity() {
528        let node = TreeNode::with_arity(42, Arity::Zero);
529        assert!(node.is_valid());
530
531        let mut node = TreeNode::with_arity(42, Arity::Zero);
532        node.add_child(TreeNode::new(1));
533        assert!(!node.is_valid());
534
535        let node = TreeNode::with_arity(42, Arity::Exact(2))
536            .attach(TreeNode::new(1))
537            .attach(TreeNode::new(2));
538        assert!(node.is_valid());
539
540        let mut node = TreeNode::with_arity(42, Arity::Exact(2));
541        node.add_child(TreeNode::new(1));
542        assert!(!node.is_valid());
543
544        let node = TreeNode::with_arity(42, Arity::Any)
545            .attach(TreeNode::new(1))
546            .attach(TreeNode::new(2))
547            .attach(TreeNode::new(3));
548        assert!(node.is_valid());
549    }
550
551    #[test]
552    fn test_node_evaluation() {
553        // Test simple arithmetic expression: (1 + 2) * 3
554        let node = TreeNode::new(Op::mul())
555            .attach(
556                TreeNode::new(Op::add())
557                    .attach(TreeNode::new(Op::constant(1.0)))
558                    .attach(TreeNode::new(Op::constant(2.0))),
559            )
560            .attach(TreeNode::new(Op::constant(3.0)));
561
562        let result = node.eval(&[]);
563        assert_eq!(result, 9.0);
564
565        // Test expression with variables: (x + 2) * 3
566        let node = TreeNode::new(Op::mul())
567            .attach(
568                TreeNode::new(Op::add())
569                    .attach(TreeNode::new(Op::var(0)))
570                    .attach(TreeNode::new(Op::constant(2.0))),
571            )
572            .attach(TreeNode::new(Op::constant(3.0)));
573
574        assert_eq!(node.eval(&[1.0]), 9.0);
575        assert_eq!(node.eval(&[2.0]), 12.0);
576        assert_eq!(node.eval(&[3.0]), 15.0);
577    }
578
579    #[test]
580    fn test_cloning_and_equality() {
581        let node1 = TreeNode::new(42)
582            .attach(TreeNode::new(1))
583            .attach(TreeNode::new(2));
584
585        let node2 = node1.clone();
586        assert_eq!(node1, node2);
587
588        let node3 = TreeNode::new(43)
589            .attach(TreeNode::new(1))
590            .attach(TreeNode::new(2));
591        assert_ne!(node1, node3);
592
593        let node4 = TreeNode::new(42).attach(TreeNode::new(1));
594        assert_ne!(node1, node4);
595    }
596
597    #[test]
598    #[cfg(feature = "serde")]
599    fn test_node_can_serialize() {
600        let root = TreeNode::new(42)
601            .attach(TreeNode::new(1))
602            .attach(
603                TreeNode::new(2)
604                    .attach(TreeNode::new(3))
605                    .attach(TreeNode::new(4)),
606            )
607            .attach(TreeNode::with_arity(3, Arity::Exact(2)));
608
609        let serialized = serde_json::to_string(&root).unwrap();
610        let deserialized: TreeNode<i32> = serde_json::from_str(&serialized).unwrap();
611
612        assert_eq!(root, deserialized);
613    }
614}
radiate_gp/collections/trees/node.rs

radiate_gp/collections/trees/
node.rs
