1#[cfg(feature = "ppl")]
25use fugue::{addr, ChoiceValue, Trace};
26use rand::Rng;
27use serde::{Deserialize, Serialize};
28use std::fmt;
29
30use crate::error::GenomeError;
31use crate::genome::bounds::MultiBounds;
32use crate::genome::traits::EvolutionaryGenome;
33
34#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
36#[serde(bound = "")]
37pub enum TreeNode<T: Terminal, F: Function> {
38 Terminal(T),
40 Function(F, Vec<TreeNode<T, F>>),
42}
43
44impl<T: Terminal, F: Function> TreeNode<T, F> {
45 pub fn terminal(value: T) -> Self {
47 Self::Terminal(value)
48 }
49
50 pub fn function(func: F, children: Vec<Self>) -> Self {
52 Self::Function(func, children)
53 }
54
55 pub fn is_terminal(&self) -> bool {
57 matches!(self, Self::Terminal(_))
58 }
59
60 pub fn is_function(&self) -> bool {
62 matches!(self, Self::Function(_, _))
63 }
64
65 pub fn depth(&self) -> usize {
71 let mut max_depth = 0;
72 let mut stack: Vec<(&Self, usize)> = vec![(self, 1)];
74 while let Some((node, d)) = stack.pop() {
75 if d > max_depth {
76 max_depth = d;
77 }
78 if let Self::Function(_, children) = node {
79 for child in children {
80 stack.push((child, d + 1));
81 }
82 }
83 }
84 max_depth
85 }
86
87 pub fn size(&self) -> usize {
92 let mut count = 0;
93 let mut stack: Vec<&Self> = vec![self];
94 while let Some(node) = stack.pop() {
95 count += 1;
96 if let Self::Function(_, children) = node {
97 for child in children {
98 stack.push(child);
99 }
100 }
101 }
102 count
103 }
104
105 pub fn positions(&self) -> Vec<Vec<usize>> {
107 let mut positions = Vec::new();
108 self.collect_positions(&[], &mut positions);
109 positions
110 }
111
112 fn collect_positions(&self, path: &[usize], positions: &mut Vec<Vec<usize>>) {
115 let mut stack: Vec<(&Self, Vec<usize>)> = vec![(self, path.to_vec())];
118 while let Some((node, node_path)) = stack.pop() {
119 positions.push(node_path.clone());
120 if let Self::Function(_, children) = node {
121 for (i, child) in children.iter().enumerate().rev() {
122 let mut child_path = node_path.clone();
123 child_path.push(i);
124 stack.push((child, child_path));
125 }
126 }
127 }
128 }
129
130 pub fn get_subtree(&self, path: &[usize]) -> Option<&Self> {
132 if path.is_empty() {
133 return Some(self);
134 }
135
136 if let Self::Function(_, children) = self {
137 let idx = path[0];
138 if idx < children.len() {
139 children[idx].get_subtree(&path[1..])
140 } else {
141 None
142 }
143 } else {
144 None
145 }
146 }
147
148 pub fn get_subtree_mut(&mut self, path: &[usize]) -> Option<&mut Self> {
150 if path.is_empty() {
151 return Some(self);
152 }
153
154 if let Self::Function(_, children) = self {
155 let idx = path[0];
156 if idx < children.len() {
157 children[idx].get_subtree_mut(&path[1..])
158 } else {
159 None
160 }
161 } else {
162 None
163 }
164 }
165
166 pub fn replace_subtree(&mut self, path: &[usize], new_subtree: Self) -> bool {
168 if path.is_empty() {
169 *self = new_subtree;
170 return true;
171 }
172
173 if let Self::Function(_, children) = self {
174 let idx = path[0];
175 if idx < children.len() {
176 if path.len() == 1 {
177 children[idx] = new_subtree;
178 true
179 } else {
180 children[idx].replace_subtree(&path[1..], new_subtree)
181 }
182 } else {
183 false
184 }
185 } else {
186 false
187 }
188 }
189
190 pub fn terminal_positions(&self) -> Vec<Vec<usize>> {
192 let mut positions = Vec::new();
193 self.collect_terminal_positions(&[], &mut positions);
194 positions
195 }
196
197 fn collect_terminal_positions(&self, path: &[usize], positions: &mut Vec<Vec<usize>>) {
200 let mut stack: Vec<(&Self, Vec<usize>)> = vec![(self, path.to_vec())];
201 while let Some((node, node_path)) = stack.pop() {
202 match node {
203 Self::Terminal(_) => positions.push(node_path),
204 Self::Function(_, children) => {
205 for (i, child) in children.iter().enumerate().rev() {
206 let mut child_path = node_path.clone();
207 child_path.push(i);
208 stack.push((child, child_path));
209 }
210 }
211 }
212 }
213 }
214
215 pub fn function_positions(&self) -> Vec<Vec<usize>> {
217 let mut positions = Vec::new();
218 self.collect_function_positions(&[], &mut positions);
219 positions
220 }
221
222 fn collect_function_positions(&self, path: &[usize], positions: &mut Vec<Vec<usize>>) {
225 let mut stack: Vec<(&Self, Vec<usize>)> = vec![(self, path.to_vec())];
226 while let Some((node, node_path)) = stack.pop() {
227 if let Self::Function(_, children) = node {
228 positions.push(node_path.clone());
229 for (i, child) in children.iter().enumerate().rev() {
230 let mut child_path = node_path.clone();
231 child_path.push(i);
232 stack.push((child, child_path));
233 }
234 }
235 }
236 }
237}
238
239pub trait Terminal:
241 Clone + Send + Sync + PartialEq + fmt::Debug + Serialize + for<'de> Deserialize<'de> + 'static
242{
243 fn random<R: Rng>(rng: &mut R) -> Self;
245
246 fn terminals() -> &'static [Self];
248
249 fn evaluate(&self, variables: &[f64]) -> f64;
251
252 fn to_string(&self) -> String;
254
255 fn encode(&self) -> (f64, f64);
266
267 fn decode(type_code: f64, payload: f64) -> Self;
272}
273
274pub trait Function:
276 Clone + Send + Sync + PartialEq + fmt::Debug + Serialize + for<'de> Deserialize<'de> + 'static
277{
278 fn arity(&self) -> usize;
280
281 fn random<R: Rng>(rng: &mut R) -> Self;
283
284 fn functions() -> &'static [Self];
286
287 fn apply(&self, args: &[f64]) -> f64;
289
290 fn to_string(&self) -> String;
292}
293
294#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
296pub enum ArithmeticTerminal {
297 Variable(usize),
299 Constant(f64),
301 Erc(f64),
303}
304
305impl Terminal for ArithmeticTerminal {
306 fn random<R: Rng>(rng: &mut R) -> Self {
307 let choice: u8 = rng.gen_range(0..3);
308 match choice {
309 0 => Self::Variable(rng.gen_range(0..10)),
310 1 => Self::Constant(rng.gen_range(-10.0..10.0)),
311 _ => Self::Erc(rng.gen_range(-1.0..1.0)),
312 }
313 }
314
315 fn terminals() -> &'static [Self] {
316 &[]
318 }
319
320 fn evaluate(&self, variables: &[f64]) -> f64 {
321 match self {
322 Self::Variable(i) => variables.get(*i).copied().unwrap_or(0.0),
323 Self::Constant(c) | Self::Erc(c) => *c,
324 }
325 }
326
327 fn to_string(&self) -> String {
328 match self {
329 Self::Variable(i) => format!("x{}", i),
330 Self::Constant(c) | Self::Erc(c) => format!("{:.4}", c),
331 }
332 }
333
334 fn encode(&self) -> (f64, f64) {
335 match self {
337 Self::Variable(i) => (0.0, *i as f64),
338 Self::Constant(c) => (1.0, *c),
339 Self::Erc(c) => (2.0, *c),
340 }
341 }
342
343 fn decode(type_code: f64, payload: f64) -> Self {
344 match type_code.round() as i64 {
345 0 => Self::Variable(payload.max(0.0) as usize),
346 1 => Self::Constant(payload),
347 2 => Self::Erc(payload),
348 _ => Self::Constant(payload),
351 }
352 }
353}
354
355#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
357pub enum ArithmeticFunction {
358 Add,
360 Sub,
362 Mul,
364 Div,
366 Sin,
368 Cos,
370 Exp,
372 Log,
374 Sqrt,
376 Pow,
378 Neg,
380 Abs,
382}
383
384impl Function for ArithmeticFunction {
385 fn arity(&self) -> usize {
386 match self {
387 Self::Add | Self::Sub | Self::Mul | Self::Div | Self::Pow => 2,
388 Self::Sin | Self::Cos | Self::Exp | Self::Log | Self::Sqrt | Self::Neg | Self::Abs => 1,
389 }
390 }
391
392 fn random<R: Rng>(rng: &mut R) -> Self {
393 let funcs = Self::functions();
394 funcs[rng.gen_range(0..funcs.len())].clone()
395 }
396
397 fn functions() -> &'static [Self] {
398 &[
405 Self::Add,
406 Self::Sub,
407 Self::Mul,
408 Self::Div,
409 Self::Sin,
410 Self::Cos,
411 Self::Exp,
412 Self::Log,
413 Self::Sqrt,
414 Self::Pow,
415 Self::Neg,
416 Self::Abs,
417 ]
418 }
419
420 fn apply(&self, args: &[f64]) -> f64 {
421 match self {
422 Self::Add => args.get(0).unwrap_or(&0.0) + args.get(1).unwrap_or(&0.0),
423 Self::Sub => args.get(0).unwrap_or(&0.0) - args.get(1).unwrap_or(&0.0),
424 Self::Mul => args.get(0).unwrap_or(&1.0) * args.get(1).unwrap_or(&1.0),
425 Self::Div => {
426 let a = args.get(0).unwrap_or(&0.0);
427 let b = args.get(1).unwrap_or(&1.0);
428 if b.abs() < 1e-10 {
429 1.0 } else {
431 a / b
432 }
433 }
434 Self::Sin => args.get(0).unwrap_or(&0.0).sin(),
435 Self::Cos => args.get(0).unwrap_or(&0.0).cos(),
436 Self::Exp => {
437 let x = args.get(0).unwrap_or(&0.0);
438 if *x > 700.0 {
439 f64::MAX } else {
441 x.exp()
442 }
443 }
444 Self::Log => {
445 let x = args.get(0).unwrap_or(&1.0);
446 if *x <= 0.0 {
447 0.0 } else {
449 x.ln()
450 }
451 }
452 Self::Sqrt => {
453 let x = args.get(0).unwrap_or(&0.0);
454 if *x < 0.0 {
455 (-x).sqrt() } else {
457 x.sqrt()
458 }
459 }
460 Self::Pow => {
461 let base = args.get(0).unwrap_or(&1.0);
462 let exp = args.get(1).unwrap_or(&1.0);
463 if base.abs() < 1e-10 && *exp < 0.0 {
470 0.0
471 } else {
472 let result = base.powf(*exp);
473 if result.is_nan() {
474 1.0
476 } else {
477 result.clamp(-1e10, 1e10)
479 }
480 }
481 }
482 Self::Neg => -args.get(0).unwrap_or(&0.0),
483 Self::Abs => args.get(0).unwrap_or(&0.0).abs(),
484 }
485 }
486
487 fn to_string(&self) -> String {
488 match self {
489 Self::Add => "+".to_string(),
490 Self::Sub => "-".to_string(),
491 Self::Mul => "*".to_string(),
492 Self::Div => "/".to_string(),
493 Self::Sin => "sin".to_string(),
494 Self::Cos => "cos".to_string(),
495 Self::Exp => "exp".to_string(),
496 Self::Log => "log".to_string(),
497 Self::Sqrt => "sqrt".to_string(),
498 Self::Pow => "pow".to_string(),
499 Self::Neg => "neg".to_string(),
500 Self::Abs => "abs".to_string(),
501 }
502 }
503}
504
505#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
507#[serde(bound = "")]
508pub struct TreeGenome<T: Terminal = ArithmeticTerminal, F: Function = ArithmeticFunction> {
509 pub root: TreeNode<T, F>,
511 pub max_depth: usize,
513}
514
515impl<T: Terminal, F: Function> TreeGenome<T, F> {
516 pub fn new(root: TreeNode<T, F>, max_depth: usize) -> Self {
518 Self { root, max_depth }
519 }
520
521 pub fn depth(&self) -> usize {
523 self.root.depth()
524 }
525
526 pub fn size(&self) -> usize {
528 self.root.size()
529 }
530
531 pub fn evaluate(&self, variables: &[f64]) -> f64 {
537 enum Task<'a, T: Terminal, F: Function> {
540 Eval(&'a TreeNode<T, F>),
541 Apply(&'a F, usize),
542 }
543
544 let mut tasks: Vec<Task<T, F>> = vec![Task::Eval(&self.root)];
545 let mut values: Vec<f64> = Vec::new();
546
547 while let Some(task) = tasks.pop() {
548 match task {
549 Task::Eval(node) => match node {
550 TreeNode::Terminal(t) => values.push(t.evaluate(variables)),
551 TreeNode::Function(f, children) => {
552 tasks.push(Task::Apply(f, children.len()));
556 for child in children.iter().rev() {
557 tasks.push(Task::Eval(child));
558 }
559 }
560 },
561 Task::Apply(f, arity) => {
562 let start = values.len() - arity;
563 let args = values.split_off(start);
564 values.push(f.apply(&args));
565 }
566 }
567 }
568
569 values.pop().unwrap_or(0.0)
570 }
571
572 pub fn dismantle(self) {
578 drop_node_iteratively(self.root);
579 }
580
581 pub fn generate_full<R: Rng>(rng: &mut R, depth: usize, max_depth: usize) -> Self {
583 let root = Self::generate_full_node(rng, depth, 0);
584 Self { root, max_depth }
585 }
586
587 fn generate_full_node<R: Rng>(
588 rng: &mut R,
589 target_depth: usize,
590 current_depth: usize,
591 ) -> TreeNode<T, F> {
592 if current_depth >= target_depth {
593 TreeNode::Terminal(T::random(rng))
594 } else {
595 let func = F::random(rng);
596 let arity = func.arity();
597 let children: Vec<TreeNode<T, F>> = (0..arity)
598 .map(|_| Self::generate_full_node(rng, target_depth, current_depth + 1))
599 .collect();
600 TreeNode::Function(func, children)
601 }
602 }
603
604 pub fn generate_grow<R: Rng>(rng: &mut R, max_depth: usize, terminal_prob: f64) -> Self {
606 let root = Self::generate_grow_node(rng, max_depth, 0, terminal_prob);
607 Self { root, max_depth }
608 }
609
610 fn generate_grow_node<R: Rng>(
611 rng: &mut R,
612 max_depth: usize,
613 current_depth: usize,
614 terminal_prob: f64,
615 ) -> TreeNode<T, F> {
616 if current_depth >= max_depth {
617 TreeNode::Terminal(T::random(rng))
618 } else if rng.gen::<f64>() < terminal_prob {
619 TreeNode::Terminal(T::random(rng))
620 } else {
621 let func = F::random(rng);
622 let arity = func.arity();
623 let children: Vec<TreeNode<T, F>> = (0..arity)
624 .map(|_| Self::generate_grow_node(rng, max_depth, current_depth + 1, terminal_prob))
625 .collect();
626 TreeNode::Function(func, children)
627 }
628 }
629
630 pub fn generate_ramped_half_and_half<R: Rng>(
632 rng: &mut R,
633 min_depth: usize,
634 max_depth: usize,
635 ) -> Self {
636 let depth = rng.gen_range(min_depth..=max_depth);
637 if rng.gen() {
638 Self::generate_full(rng, depth, max_depth)
639 } else {
640 Self::generate_grow(rng, depth, 0.3)
641 }
642 }
643
644 pub fn generate_with_depth<R: Rng>(rng: &mut R, max_depth: usize) -> Self {
652 let max_depth = max_depth.max(1);
653 let min_depth = 2.min(max_depth);
654 Self::generate_ramped_half_and_half(rng, min_depth, max_depth)
655 }
656
657 pub fn to_sexpr(&self) -> String {
659 self.node_to_sexpr(&self.root)
660 }
661
662 fn node_to_sexpr(&self, node: &TreeNode<T, F>) -> String {
663 match node {
664 TreeNode::Terminal(t) => t.to_string(),
665 TreeNode::Function(f, children) => {
666 let child_strs: Vec<String> =
667 children.iter().map(|c| self.node_to_sexpr(c)).collect();
668 format!("({} {})", f.to_string(), child_strs.join(" "))
669 }
670 }
671 }
672
673 pub fn random_position<R: Rng>(&self, rng: &mut R) -> Vec<usize> {
675 let positions = self.root.positions();
676 positions[rng.gen_range(0..positions.len())].clone()
677 }
678
679 pub fn random_terminal_position<R: Rng>(&self, rng: &mut R) -> Option<Vec<usize>> {
681 let positions = self.root.terminal_positions();
682 if positions.is_empty() {
683 None
684 } else {
685 Some(positions[rng.gen_range(0..positions.len())].clone())
686 }
687 }
688
689 pub fn random_function_position<R: Rng>(&self, rng: &mut R) -> Option<Vec<usize>> {
691 let positions = self.root.function_positions();
692 if positions.is_empty() {
693 None
694 } else {
695 Some(positions[rng.gen_range(0..positions.len())].clone())
696 }
697 }
698}
699
700impl<T: Terminal, F: Function> EvolutionaryGenome for TreeGenome<T, F> {
701 type Allele = TreeNode<T, F>;
702 type Phenotype = Self;
703
704 fn decode(&self) -> Self::Phenotype {
705 self.clone()
706 }
707
708 fn dimension(&self) -> usize {
709 self.size()
710 }
711
712 fn generate<R: Rng>(rng: &mut R, bounds: &MultiBounds) -> Self {
719 let max_depth = bounds.dimension().clamp(3, 10);
720 Self::generate_with_depth(rng, max_depth)
721 }
722
723 fn distance(&self, other: &Self) -> f64 {
724 let size_diff = (self.size() as f64 - other.size() as f64).abs();
726 let depth_diff = (self.depth() as f64 - other.depth() as f64).abs();
727 size_diff + depth_diff
728 }
729
730 fn try_distance(&self, other: &Self) -> Result<f64, GenomeError> {
731 Ok(self.distance(other))
733 }
734}
735
736#[cfg(feature = "ppl")]
737impl<T: Terminal, F: Function> crate::genome::trace_genome::TraceGenome for TreeGenome<T, F> {
738 fn to_trace(&self) -> Trace {
739 let mut trace = Trace::default();
740 let mut index = 0;
741 self.node_to_trace(&self.root, &mut trace, &mut index);
742 trace.insert_choice(
744 addr!("tree_max_depth"),
745 ChoiceValue::Usize(self.max_depth),
746 0.0,
747 );
748 trace.insert_choice(addr!("tree_size"), ChoiceValue::Usize(index), 0.0);
749 trace
750 }
751
752 fn from_trace(trace: &Trace) -> Result<Self, GenomeError> {
753 let max_depth = trace
754 .get_usize(&addr!("tree_max_depth"))
755 .ok_or_else(|| GenomeError::MissingAddress("tree_max_depth".to_string()))?;
756
757 let mut index = 0;
758 let root = Self::node_from_trace(trace, &mut index)?;
759 Ok(Self { root, max_depth })
760 }
761
762 fn trace_prefix() -> &'static str {
763 "tree"
764 }
765}
766
767#[cfg(feature = "ppl")]
768impl<T: Terminal, F: Function> TreeGenome<T, F> {
769 fn node_to_trace(&self, node: &TreeNode<T, F>, trace: &mut Trace, index: &mut usize) {
770 let current_index = *index;
771 *index += 1;
772
773 match node {
774 TreeNode::Terminal(t) => {
775 trace.insert_choice(
777 addr!("tree_is_terminal", current_index),
778 ChoiceValue::Bool(true),
779 0.0,
780 );
781 let (term_type, term_val) = Self::encode_terminal(t);
784 trace.insert_choice(
785 addr!("tree_term_type", current_index),
786 ChoiceValue::F64(term_type),
787 0.0,
788 );
789 trace.insert_choice(
790 addr!("tree_term_val", current_index),
791 ChoiceValue::F64(term_val),
792 0.0,
793 );
794 }
795 TreeNode::Function(f, children) => {
796 trace.insert_choice(
798 addr!("tree_is_terminal", current_index),
799 ChoiceValue::Bool(false),
800 0.0,
801 );
802 let func_idx = Self::encode_function(f);
804 trace.insert_choice(
805 addr!("tree_func_idx", current_index),
806 ChoiceValue::Usize(func_idx),
807 0.0,
808 );
809 trace.insert_choice(
810 addr!("tree_arity", current_index),
811 ChoiceValue::Usize(children.len()),
812 0.0,
813 );
814 for child in children {
816 self.node_to_trace(child, trace, index);
817 }
818 }
819 }
820 }
821
822 fn node_from_trace(trace: &Trace, index: &mut usize) -> Result<TreeNode<T, F>, GenomeError> {
823 let current_index = *index;
824 *index += 1;
825
826 let is_terminal = trace
827 .get_bool(&addr!("tree_is_terminal", current_index))
828 .ok_or_else(|| {
829 GenomeError::MissingAddress(format!("tree_is_terminal#{}", current_index))
830 })?;
831
832 if is_terminal {
833 let term_type = trace
834 .get_f64(&addr!("tree_term_type", current_index))
835 .ok_or_else(|| {
836 GenomeError::MissingAddress(format!("tree_term_type#{}", current_index))
837 })?;
838 let term_val = trace
839 .get_f64(&addr!("tree_term_val", current_index))
840 .ok_or_else(|| {
841 GenomeError::MissingAddress(format!("tree_term_val#{}", current_index))
842 })?;
843
844 let terminal = Self::decode_terminal(term_type, term_val)?;
845 Ok(TreeNode::Terminal(terminal))
846 } else {
847 let func_idx = trace
848 .get_usize(&addr!("tree_func_idx", current_index))
849 .ok_or_else(|| {
850 GenomeError::MissingAddress(format!("tree_func_idx#{}", current_index))
851 })?;
852 let arity = trace
853 .get_usize(&addr!("tree_arity", current_index))
854 .ok_or_else(|| {
855 GenomeError::MissingAddress(format!("tree_arity#{}", current_index))
856 })?;
857
858 let func = Self::decode_function(func_idx)?;
859 let mut children = Vec::with_capacity(arity);
860 for _ in 0..arity {
861 children.push(Self::node_from_trace(trace, index)?);
862 }
863 Ok(TreeNode::Function(func, children))
864 }
865 }
866
867 fn encode_terminal(terminal: &T) -> (f64, f64) {
871 terminal.encode()
872 }
873
874 fn decode_terminal(term_type: f64, term_val: f64) -> Result<T, GenomeError> {
875 Ok(T::decode(term_type, term_val))
876 }
877
878 fn encode_function(func: &F) -> usize {
889 match F::functions()
890 .iter()
891 .position(|candidate| candidate == func)
892 {
893 Some(idx) => idx,
894 None => {
895 debug_assert!(
896 false,
897 "encode_function: function not present in F::functions(); a \
898 custom Function impl must enumerate every variant it can \
899 produce so trace round-trips stay lossless (EV-04)"
900 );
901 0
902 }
903 }
904 }
905
906 fn decode_function(func_idx: usize) -> Result<F, GenomeError> {
907 let funcs = F::functions();
908 funcs.get(func_idx).cloned().ok_or_else(|| {
909 GenomeError::InvalidStructure(format!(
910 "Function index {} out of range ({} functions available)",
911 func_idx,
912 funcs.len()
913 ))
914 })
915 }
916}
917
918impl<T: Terminal, F: Function> fmt::Display for TreeGenome<T, F> {
919 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
920 write!(f, "{}", self.to_sexpr())
921 }
922}
923
924pub trait TreeGenomeType: EvolutionaryGenome {
926 type Term: Terminal;
928 type Func: Function;
930
931 fn root(&self) -> &TreeNode<Self::Term, Self::Func>;
933
934 fn root_mut(&mut self) -> &mut TreeNode<Self::Term, Self::Func>;
936
937 fn max_depth(&self) -> usize;
939
940 fn from_root(root: TreeNode<Self::Term, Self::Func>, max_depth: usize) -> Self;
942}
943
944impl<T: Terminal, F: Function> TreeGenomeType for TreeGenome<T, F> {
945 type Term = T;
946 type Func = F;
947
948 fn root(&self) -> &TreeNode<T, F> {
949 &self.root
950 }
951
952 fn root_mut(&mut self) -> &mut TreeNode<T, F> {
953 &mut self.root
954 }
955
956 fn max_depth(&self) -> usize {
957 self.max_depth
958 }
959
960 fn from_root(root: TreeNode<T, F>, max_depth: usize) -> Self {
961 Self { root, max_depth }
962 }
963}
964
965impl<T: Terminal, F: Function> Drop for TreeNode<T, F> {
966 fn drop(&mut self) {
978 let mut stack: Vec<TreeNode<T, F>> = match self {
980 TreeNode::Terminal(_) => return,
981 TreeNode::Function(_, children) => std::mem::take(children),
982 };
983 while let Some(mut node) = stack.pop() {
984 if let TreeNode::Function(_, grandchildren) = &mut node {
985 stack.append(&mut std::mem::take(grandchildren));
988 }
989 }
992 }
993}
994
995pub fn drop_node_iteratively<T: Terminal, F: Function>(node: TreeNode<T, F>) {
1003 drop(node);
1004}
1005
1006#[cfg(test)]
1007mod tests {
1008 use super::*;
1009
1010 #[test]
1011 fn test_tree_node_terminal() {
1012 let node: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1013 TreeNode::terminal(ArithmeticTerminal::Variable(0));
1014 assert!(node.is_terminal());
1015 assert!(!node.is_function());
1016 assert_eq!(node.depth(), 1);
1017 assert_eq!(node.size(), 1);
1018 }
1019
1020 #[test]
1021 fn test_tree_node_function() {
1022 let left = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1023 let right = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1024 let node = TreeNode::function(ArithmeticFunction::Add, vec![left, right]);
1025
1026 assert!(!node.is_terminal());
1027 assert!(node.is_function());
1028 assert_eq!(node.depth(), 2);
1029 assert_eq!(node.size(), 3);
1030 }
1031
1032 #[test]
1033 fn test_tree_node_positions() {
1034 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1036 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1037 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1038 let mul = TreeNode::function(ArithmeticFunction::Mul, vec![c1, x1]);
1039 let add = TreeNode::function(ArithmeticFunction::Add, vec![x0, mul]);
1040
1041 let positions = add.positions();
1042 assert_eq!(positions.len(), 5); assert!(positions.contains(&vec![])); assert!(positions.contains(&vec![0])); assert!(positions.contains(&vec![1])); assert!(positions.contains(&vec![1, 0])); assert!(positions.contains(&vec![1, 1])); }
1049
1050 #[test]
1051 fn test_tree_node_get_subtree() {
1052 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1053 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1054 let add: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1055 TreeNode::function(ArithmeticFunction::Add, vec![x0.clone(), c1]);
1056
1057 assert_eq!(add.get_subtree(&[0]), Some(&x0));
1058 assert!(add.get_subtree(&[2]).is_none());
1059 }
1060
1061 #[test]
1062 fn test_tree_genome_evaluate() {
1063 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1065 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1066 let add = TreeNode::function(ArithmeticFunction::Add, vec![x0, x1]);
1067 let tree = TreeGenome::new(add, 5);
1068
1069 assert_eq!(tree.evaluate(&[3.0, 4.0]), 7.0);
1070 }
1071
1072 #[test]
1073 fn test_tree_genome_evaluate_complex() {
1074 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1076 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1077 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1078 let add = TreeNode::function(ArithmeticFunction::Add, vec![x0, c1]);
1079 let mul = TreeNode::function(ArithmeticFunction::Mul, vec![add, x1]);
1080 let tree = TreeGenome::new(mul, 5);
1081
1082 assert_eq!(tree.evaluate(&[2.0, 3.0]), 9.0); }
1084
1085 #[test]
1086 fn test_tree_genome_generate_full() {
1087 let mut rng = rand::thread_rng();
1088 let tree: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1089 TreeGenome::generate_full(&mut rng, 3, 5);
1090
1091 assert!(tree.depth() >= 3);
1094 assert!(tree.size() >= 1);
1095 }
1096
1097 #[test]
1098 fn test_tree_genome_generate_grow() {
1099 let mut rng = rand::thread_rng();
1100 let tree: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1101 TreeGenome::generate_grow(&mut rng, 5, 0.3);
1102
1103 assert!(tree.depth() <= 6);
1105 assert!(tree.size() >= 1);
1106 }
1107
1108 #[test]
1109 fn test_tree_genome_to_sexpr() {
1110 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1111 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1112 let add: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1113 TreeNode::function(ArithmeticFunction::Add, vec![x0, c1]);
1114 let tree = TreeGenome::new(add, 5);
1115
1116 let sexpr = tree.to_sexpr();
1117 assert!(sexpr.contains('+'));
1118 assert!(sexpr.contains("x0"));
1119 assert!(sexpr.contains("1.0"));
1120 }
1121
1122 #[test]
1123 #[cfg(feature = "ppl")]
1124 fn test_tree_genome_trace_roundtrip() {
1125 use crate::genome::trace_genome::TraceGenome;
1129 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1130 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(2.5));
1131 let sub = TreeNode::function(ArithmeticFunction::Sub, vec![x0, c1]);
1133 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1134 let erc = TreeNode::terminal(ArithmeticTerminal::Erc(-0.75));
1135 let mul = TreeNode::function(ArithmeticFunction::Mul, vec![x1, erc]);
1136 let root = TreeNode::function(ArithmeticFunction::Div, vec![sub, mul]);
1137 let original: TreeGenome<ArithmeticTerminal, ArithmeticFunction> = TreeGenome::new(root, 5);
1138
1139 let trace = original.to_trace();
1140 let recovered: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1141 TreeGenome::from_trace(&trace).unwrap();
1142
1143 assert_eq!(original, recovered);
1145 assert_eq!(original.max_depth, recovered.max_depth);
1146 assert_eq!(original.size(), recovered.size());
1147 assert_eq!(recovered.to_sexpr(), original.to_sexpr());
1148
1149 for vars in [[3.0, 4.0], [-1.0, 2.0], [0.5, -0.5]] {
1151 assert_eq!(recovered.evaluate(&vars), original.evaluate(&vars));
1152 }
1153 }
1154
1155 #[test]
1156 #[cfg(feature = "ppl")]
1157 fn test_tree_genome_trace_roundtrip_pow_node() {
1158 use crate::genome::trace_genome::TraceGenome;
1163 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1164 let two = TreeNode::terminal(ArithmeticTerminal::Constant(2.0));
1165 let root = TreeNode::function(ArithmeticFunction::Pow, vec![x0, two]);
1166 let original: TreeGenome<ArithmeticTerminal, ArithmeticFunction> = TreeGenome::new(root, 3);
1167
1168 let recovered: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1169 TreeGenome::from_trace(&original.to_trace()).unwrap();
1170
1171 assert_eq!(original, recovered, "Pow node must survive the round-trip");
1172 assert_eq!(recovered.to_sexpr(), original.to_sexpr());
1173 assert!(
1174 recovered.to_sexpr().contains("pow"),
1175 "recovered tree must still be a pow, got {}",
1176 recovered.to_sexpr()
1177 );
1178 assert_eq!(original.evaluate(&[3.0]), 9.0);
1180 assert_eq!(recovered.evaluate(&[3.0]), original.evaluate(&[3.0]));
1181 for vars in [[3.0], [-2.0], [0.5], [1.5]] {
1182 assert_eq!(recovered.evaluate(&vars), original.evaluate(&vars));
1183 }
1184 }
1185
1186 #[test]
1187 #[cfg(feature = "ppl")]
1188 fn test_every_arithmetic_function_variant_roundtrips_losslessly() {
1189 use crate::genome::trace_genome::TraceGenome;
1194 use ArithmeticFunction::*;
1195 let all = [Add, Sub, Mul, Div, Sin, Cos, Exp, Log, Sqrt, Pow, Neg, Abs];
1196 assert_eq!(
1198 ArithmeticFunction::functions().len(),
1199 all.len(),
1200 "functions() must list every ArithmeticFunction variant"
1201 );
1202 for f in &all {
1203 assert!(
1204 ArithmeticFunction::functions().contains(f),
1205 "functions() is missing variant {f:?}"
1206 );
1207 }
1208
1209 for f in all {
1210 let arity = f.arity();
1211 let children: Vec<_> = (0..arity)
1212 .map(|i| TreeNode::terminal(ArithmeticTerminal::Variable(i)))
1213 .collect();
1214 let root = TreeNode::function(f.clone(), children);
1215 let original: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1216 TreeGenome::new(root, 3);
1217 let recovered: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1218 TreeGenome::from_trace(&original.to_trace()).unwrap();
1219 assert_eq!(
1220 original, recovered,
1221 "variant {f:?} lost identity on round-trip"
1222 );
1223 for vars in [[2.0, 3.0], [-1.5, 0.75]] {
1224 let (a, b) = (recovered.evaluate(&vars), original.evaluate(&vars));
1225 assert!(
1229 a.to_bits() == b.to_bits(),
1230 "variant {f:?} evaluated differently after round-trip: {a} vs {b}"
1231 );
1232 }
1233 }
1234 }
1235
1236 #[test]
1237 fn test_tree_generate_with_depth_explicit() {
1238 let mut rng = rand::thread_rng();
1240 let tree: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1241 TreeGenome::generate_with_depth(&mut rng, 4);
1242 assert!(tree.depth() >= 1);
1243 assert!(tree.depth() <= 5); let _ =
1246 TreeGenome::<ArithmeticTerminal, ArithmeticFunction>::generate_with_depth(&mut rng, 0);
1247 }
1248
1249 #[test]
1250 #[cfg(feature = "ppl")]
1251 fn test_arithmetic_function_ordering_is_stable() {
1252 let funcs = ArithmeticFunction::functions();
1255 assert_eq!(funcs[0], ArithmeticFunction::Add);
1256 assert_eq!(funcs[1], ArithmeticFunction::Sub);
1257 assert_eq!(funcs[2], ArithmeticFunction::Mul);
1258 assert_eq!(funcs[3], ArithmeticFunction::Div);
1259 for (idx, f) in funcs.iter().enumerate() {
1261 let encoded = TreeGenome::<ArithmeticTerminal, ArithmeticFunction>::encode_function(f);
1262 assert_eq!(encoded, idx);
1263 let decoded =
1264 TreeGenome::<ArithmeticTerminal, ArithmeticFunction>::decode_function(encoded)
1265 .unwrap();
1266 assert_eq!(&decoded, f);
1267 }
1268 }
1269
1270 #[test]
1271 fn test_arithmetic_terminal_encode_decode_roundtrip() {
1272 for t in [
1275 ArithmeticTerminal::Variable(0),
1276 ArithmeticTerminal::Variable(7),
1277 ArithmeticTerminal::Constant(3.25),
1278 ArithmeticTerminal::Constant(-100.5),
1279 ArithmeticTerminal::Erc(0.0),
1280 ArithmeticTerminal::Erc(-0.75),
1281 ] {
1282 let (ty, val) = t.encode();
1283 assert_eq!(ArithmeticTerminal::decode(ty, val), t);
1284 }
1285 }
1286
1287 #[test]
1288 fn test_tree_deep_no_stack_overflow() {
1289 let depth = 100_000usize;
1294 let mut root: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1296 TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1297 for _ in 0..depth {
1298 root = TreeNode::function(ArithmeticFunction::Neg, vec![root]);
1299 }
1300 let tree = TreeGenome::new(root, depth + 1);
1301
1302 assert_eq!(tree.size(), depth + 1);
1304 assert_eq!(tree.depth(), depth + 1);
1305 let value = tree.evaluate(&[]);
1307 assert!(value.is_finite());
1308 assert_eq!(value, 1.0);
1309
1310 tree.dismantle();
1312 }
1313
1314 #[test]
1315 fn test_drop_node_iteratively_frees_deep_tree() {
1316 let mut node: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1319 TreeNode::terminal(ArithmeticTerminal::Constant(0.0));
1320 for _ in 0..100_000 {
1321 node = TreeNode::function(ArithmeticFunction::Abs, vec![node]);
1322 }
1323 drop_node_iteratively(node);
1324 }
1325
1326 fn deep_tree(depth: usize) -> TreeGenome<ArithmeticTerminal, ArithmeticFunction> {
1327 let mut root: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1329 TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1330 for _ in 0..depth {
1331 root = TreeNode::function(ArithmeticFunction::Neg, vec![root]);
1332 }
1333 TreeGenome::new(root, depth + 1)
1334 }
1335
1336 #[test]
1337 fn test_deep_tree_implicit_drop_no_overflow() {
1338 let depth = 100_000usize;
1343 {
1344 let tree = deep_tree(depth);
1345 assert_eq!(tree.size(), depth + 1);
1346 }
1348 {
1350 let mut node: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1351 TreeNode::terminal(ArithmeticTerminal::Constant(0.0));
1352 for _ in 0..depth {
1353 node = TreeNode::function(ArithmeticFunction::Abs, vec![node]);
1354 }
1355 let _ = node; }
1357 }
1358
1359 #[test]
1360 fn test_deep_tree_position_collectors_no_overflow() {
1361 let depth = 2_000usize;
1385
1386 let tree = deep_tree(depth);
1393 let tree = std::thread::Builder::new()
1394 .stack_size(64 * 1024)
1395 .spawn(move || {
1396 assert_eq!(tree.root.positions().len(), depth + 1);
1397 assert_eq!(tree.root.terminal_positions().len(), 1);
1399 assert_eq!(tree.root.function_positions().len(), depth);
1400 tree
1401 })
1402 .expect("spawn 64 KiB-stack thread")
1403 .join()
1404 .expect("position collectors overflowed a 64 KiB stack (recursive regression?)");
1405
1406 tree.dismantle();
1407 }
1408
1409 #[test]
1410 fn test_tree_node_replace_subtree() {
1411 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1412 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1413 let mut add: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1414 TreeNode::function(ArithmeticFunction::Add, vec![x0, x1]);
1415
1416 let c5 = TreeNode::terminal(ArithmeticTerminal::Constant(5.0));
1417 add.replace_subtree(&[0], c5);
1418
1419 let tree = TreeGenome::new(add, 5);
1421 assert_eq!(tree.evaluate(&[0.0, 3.0]), 8.0); }
1423
1424 #[test]
1425 fn test_arithmetic_function_protected_div() {
1426 assert_eq!(ArithmeticFunction::Div.apply(&[1.0, 0.0]), 1.0);
1427 assert_eq!(ArithmeticFunction::Div.apply(&[6.0, 2.0]), 3.0);
1428 }
1429
1430 #[test]
1431 fn test_arithmetic_function_protected_log() {
1432 assert_eq!(ArithmeticFunction::Log.apply(&[-1.0]), 0.0);
1433 assert!((ArithmeticFunction::Log.apply(&[std::f64::consts::E]) - 1.0).abs() < 0.001);
1434 }
1435
1436 #[test]
1437 fn test_arithmetic_function_protected_sqrt() {
1438 assert_eq!(ArithmeticFunction::Sqrt.apply(&[4.0]), 2.0);
1439 assert_eq!(ArithmeticFunction::Sqrt.apply(&[-4.0]), 2.0); }
1441
1442 #[test]
1443 fn test_tree_genome_evolutionary_genome_trait() {
1444 let mut rng = rand::thread_rng();
1445 let bounds = MultiBounds::symmetric(5.0, 5);
1446 let tree: TreeGenome<ArithmeticTerminal, ArithmeticFunction> =
1447 TreeGenome::generate(&mut rng, &bounds);
1448
1449 assert!(tree.dimension() >= 1);
1450 let decoded = tree.decode();
1451 assert_eq!(decoded.size(), tree.size());
1452 }
1453
1454 #[test]
1455 fn test_tree_terminal_and_function_positions() {
1456 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1458 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1459 let x1 = TreeNode::terminal(ArithmeticTerminal::Variable(1));
1460 let mul = TreeNode::function(ArithmeticFunction::Mul, vec![c1, x1]);
1461 let add: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1462 TreeNode::function(ArithmeticFunction::Add, vec![x0, mul]);
1463
1464 let terminal_positions = add.terminal_positions();
1465 assert_eq!(terminal_positions.len(), 3); let function_positions = add.function_positions();
1468 assert_eq!(function_positions.len(), 2); }
1470
1471 #[test]
1472 fn test_tree_genome_display() {
1473 let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1474 let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(1.0));
1475 let add: TreeNode<ArithmeticTerminal, ArithmeticFunction> =
1476 TreeNode::function(ArithmeticFunction::Add, vec![x0, c1]);
1477 let tree = TreeGenome::new(add, 5);
1478
1479 let display = format!("{}", tree);
1480 assert!(!display.is_empty());
1481 }
1482}