Skip to main content

fugue_evo/genome/
tree.rs

1//! Tree genomes for genetic programming
2//!
3//! This module provides tree-based genomes for symbolic regression and
4//! genetic programming applications.
5//!
6//! # Deep trees and stack safety
7//!
8//! Every traversal in this module uses an explicit work stack instead of
9//! recursion, so no operation overflows the call stack even for pathologically
10//! deep trees. This covers the read-side traversals ([`TreeNode::depth`],
11//! [`TreeNode::size`], [`TreeGenome::evaluate`], and the position collectors
12//! backing [`TreeNode::positions`]/[`TreeNode::terminal_positions`]/
13//! [`TreeNode::function_positions`]), the [`crate::operators`] point-mutation
14//! traversal, AND teardown: [`TreeNode`] implements a stack-safe [`Drop`]
15//! (EV-60) that frees an arbitrarily deep tree iteratively, so even dropping a
16//! deep tree *implicitly* (never calling [`TreeGenome::dismantle`]) cannot
17//! overflow the stack. The `Drop` impl moves each node's children out with
18//! `mem::take` — which is permitted under `Drop`, unlike moving a field out by
19//! value (E0509) — so the operator layer must likewise use `mem::take`/in-place
20//! mutation rather than by-value destructuring of an owned `TreeNode`.
21//! [`TreeGenome::dismantle`] and [`drop_node_iteratively`] remain as explicit,
22//! self-documenting entry points but are no longer required for correctness.
23
24#[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/// A node in a GP tree
35#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
36#[serde(bound = "")]
37pub enum TreeNode<T: Terminal, F: Function> {
38    /// Terminal node (leaf)
39    Terminal(T),
40    /// Function node (internal)
41    Function(F, Vec<TreeNode<T, F>>),
42}
43
44impl<T: Terminal, F: Function> TreeNode<T, F> {
45    /// Create a new terminal node
46    pub fn terminal(value: T) -> Self {
47        Self::Terminal(value)
48    }
49
50    /// Create a new function node
51    pub fn function(func: F, children: Vec<Self>) -> Self {
52        Self::Function(func, children)
53    }
54
55    /// Check if this is a terminal node
56    pub fn is_terminal(&self) -> bool {
57        matches!(self, Self::Terminal(_))
58    }
59
60    /// Check if this is a function node
61    pub fn is_function(&self) -> bool {
62        matches!(self, Self::Function(_, _))
63    }
64
65    /// Get the depth of this subtree.
66    ///
67    /// Uses an explicit work stack rather than recursion so that pathologically
68    /// deep trees cannot overflow the call stack (see the module note on deep
69    /// trees).
70    pub fn depth(&self) -> usize {
71        let mut max_depth = 0;
72        // (node, depth-of-node) pairs; a terminal has depth 1.
73        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    /// Get the number of nodes in this subtree.
88    ///
89    /// Uses an explicit work stack rather than recursion (see the module note on
90    /// deep trees).
91    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    /// Get all node positions (preorder traversal indices)
106    pub fn positions(&self) -> Vec<Vec<usize>> {
107        let mut positions = Vec::new();
108        self.collect_positions(&[], &mut positions);
109        positions
110    }
111
112    /// Iterative preorder collector (EV-60: explicit stack, no recursion) that
113    /// records the path to every node.
114    fn collect_positions(&self, path: &[usize], positions: &mut Vec<Vec<usize>>) {
115        // (node, path-to-node). Children are pushed in reverse so they pop in
116        // left-to-right order, preserving the original preorder traversal.
117        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    /// Get a subtree at the given path
131    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    /// Get a mutable subtree at the given path
149    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    /// Replace a subtree at the given path
167    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    /// Get all terminal positions
191    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    /// Iterative preorder collector (EV-60: explicit stack, no recursion) that
198    /// records the path to every terminal (leaf) node.
199    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    /// Get all function positions
216    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    /// Iterative preorder collector (EV-60: explicit stack, no recursion) that
223    /// records the path to every function (internal) node.
224    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
239/// Trait for terminal nodes in GP trees
240pub trait Terminal:
241    Clone + Send + Sync + PartialEq + fmt::Debug + Serialize + for<'de> Deserialize<'de> + 'static
242{
243    /// Generate a random terminal
244    fn random<R: Rng>(rng: &mut R) -> Self;
245
246    /// Get the set of available terminals
247    fn terminals() -> &'static [Self];
248
249    /// Evaluate this terminal with the given variable bindings
250    fn evaluate(&self, variables: &[f64]) -> f64;
251
252    /// Convert to string representation
253    fn to_string(&self) -> String;
254
255    /// Encode this terminal as a `(type_code, payload)` pair for lossless trace
256    /// round-tripping.
257    ///
258    /// `type_code` identifies the terminal variant (a discriminant) and
259    /// `payload` carries its associated value. The pair must satisfy the
260    /// inverse relationship `Self::decode(self.encode()) == *self` so that
261    /// [`TreeGenome::from_trace`](crate::genome::trace_genome::TraceGenome::from_trace)
262    /// reproduces the exact terminal that
263    /// [`TreeGenome::to_trace`](crate::genome::trace_genome::TraceGenome::to_trace)
264    /// serialized.
265    fn encode(&self) -> (f64, f64);
266
267    /// Decode a terminal previously produced by [`encode`](Self::encode).
268    ///
269    /// This is the inverse of [`encode`](Self::encode) and must reconstruct an
270    /// equal terminal for any `(type_code, payload)` this type emits.
271    fn decode(type_code: f64, payload: f64) -> Self;
272}
273
274/// Trait for function nodes in GP trees
275pub trait Function:
276    Clone + Send + Sync + PartialEq + fmt::Debug + Serialize + for<'de> Deserialize<'de> + 'static
277{
278    /// Get the arity (number of arguments) of this function
279    fn arity(&self) -> usize;
280
281    /// Generate a random function
282    fn random<R: Rng>(rng: &mut R) -> Self;
283
284    /// Get the set of available functions
285    fn functions() -> &'static [Self];
286
287    /// Apply this function to the given arguments
288    fn apply(&self, args: &[f64]) -> f64;
289
290    /// Convert to string representation
291    fn to_string(&self) -> String;
292}
293
294/// Standard arithmetic terminals for symbolic regression
295#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
296pub enum ArithmeticTerminal {
297    /// Variable x_i
298    Variable(usize),
299    /// Constant value
300    Constant(f64),
301    /// Ephemeral random constant (ERC)
302    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        // Return a representative set; actual terminals depend on context
317        &[]
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        // type_code: 0 = Variable, 1 = Constant, 2 = Erc
336        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            // Unknown discriminant (corrupt trace): preserve the payload as a
349            // constant rather than fabricating a random terminal.
350            _ => Self::Constant(payload),
351        }
352    }
353}
354
355/// Standard arithmetic functions for symbolic regression
356#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
357pub enum ArithmeticFunction {
358    /// Addition
359    Add,
360    /// Subtraction
361    Sub,
362    /// Multiplication
363    Mul,
364    /// Protected division (returns 1.0 for division by zero)
365    Div,
366    /// Sine
367    Sin,
368    /// Cosine
369    Cos,
370    /// Exponential
371    Exp,
372    /// Natural logarithm (protected)
373    Log,
374    /// Square root (protected)
375    Sqrt,
376    /// Power
377    Pow,
378    /// Negation (unary)
379    Neg,
380    /// Absolute value (unary)
381    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        // EV-04: this slice is the canonical, stable index table used by
399        // TreeGenome trace encode/decode (encode_function looks up a function's
400        // position here). EVERY variant of the enum MUST appear exactly once, in
401        // enum-declaration order, or a node of the missing variant round-trips to
402        // the wrong function (previously `Pow` was absent, so a Pow node silently
403        // decoded to index 0 = Add). Keep this in sync with the enum above.
404        &[
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 // Protected division
430                } 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 // Overflow protection
440                } 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 // Protected log
448                } 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() // Protected sqrt
456                } 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                // Protected power (EV-04): now that `Pow` is part of the function
464                // set drawn by the generators and point mutation, it must never
465                // produce NaN/Inf — matching the protection every other function
466                // in this set already provides. `powf` returns NaN for a negative
467                // base with a fractional exponent, so we guard both the base≈0
468                // (negative exponent) case and any non-finite result.
469                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                        // e.g. (-1.5)^0.75: fall back to a finite value.
475                        1.0
476                    } else {
477                        // `clamp` maps ±inf to ±1e10 and leaves finite values as-is.
478                        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/// Tree genome for genetic programming
506#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
507#[serde(bound = "")]
508pub struct TreeGenome<T: Terminal = ArithmeticTerminal, F: Function = ArithmeticFunction> {
509    /// Root node of the tree
510    pub root: TreeNode<T, F>,
511    /// Maximum allowed depth
512    pub max_depth: usize,
513}
514
515impl<T: Terminal, F: Function> TreeGenome<T, F> {
516    /// Create a new tree genome
517    pub fn new(root: TreeNode<T, F>, max_depth: usize) -> Self {
518        Self { root, max_depth }
519    }
520
521    /// Get the depth of the tree
522    pub fn depth(&self) -> usize {
523        self.root.depth()
524    }
525
526    /// Get the number of nodes in the tree
527    pub fn size(&self) -> usize {
528        self.root.size()
529    }
530
531    /// Evaluate the tree with given variable bindings.
532    ///
533    /// Uses an explicit work stack (iterative post-order traversal) rather than
534    /// recursion, so a pathologically deep tree cannot overflow the call stack
535    /// (see the module note on deep trees).
536    pub fn evaluate(&self, variables: &[f64]) -> f64 {
537        // Two task kinds: `Eval` expands a node; `Apply` combines the results of
538        // a function node's already-evaluated children.
539        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                        // Schedule the apply, then push children in reverse so
553                        // they evaluate left-to-right and land on `values` in
554                        // argument order.
555                        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    /// Free this tree without deep recursion.
573    ///
574    /// Consumes the genome and dismantles its tree iteratively (see the module
575    /// note on deep trees). Use this for pathologically deep trees that would
576    /// otherwise overflow the stack when dropped implicitly.
577    pub fn dismantle(self) {
578        drop_node_iteratively(self.root);
579    }
580
581    /// Generate a random tree using the "full" method
582    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    /// Generate a random tree using the "grow" method
605    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    /// Generate using ramped half-and-half
631    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    /// Generate a random tree with an explicit maximum depth.
645    ///
646    /// This is the honest constructor for random generation: unlike
647    /// [`EvolutionaryGenome::generate`],
648    /// which overloads `MultiBounds` and remaps its *dimension count* to a depth,
649    /// this takes the maximum depth directly. It uses ramped half-and-half
650    /// between depth 2 and `max_depth` (both clamped to at least 1).
651    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    /// Convert tree to S-expression string
658    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    /// Get a random node position
674    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    /// Get a random terminal position
680    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    /// Get a random function position
690    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    /// Generate a random tree.
713    ///
714    /// Only `bounds.dimension()` is consulted — it is remapped (clamped to
715    /// `[3, 10]`) to a maximum tree depth — and the per-dimension `min`/`max`
716    /// values are ignored. Prefer [`TreeGenome::generate_with_depth`] to make the
717    /// depth explicit instead of overloading `MultiBounds`.
718    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        // Tree edit distance approximation based on size difference
725        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        // Any two trees are comparable (size/depth deltas), so this never errs.
732        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        // Store max_depth and total size
743        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                // Store is_terminal flag (true = terminal)
776                trace.insert_choice(
777                    addr!("tree_is_terminal", current_index),
778                    ChoiceValue::Bool(true),
779                    0.0,
780                );
781                // For ArithmeticTerminal, store the variant type and value
782                // We encode using f64 for simplicity
783                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                // Store is_terminal flag (false = function)
797                trace.insert_choice(
798                    addr!("tree_is_terminal", current_index),
799                    ChoiceValue::Bool(false),
800                    0.0,
801                );
802                // Store function type as index and arity
803                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                // Recurse into children
815                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    // Encode a terminal as a (type_code, payload) pair via the type's own
868    // lossless `Terminal::encode`, so trace round-tripping preserves the exact
869    // terminal (variant + value) rather than fabricating a random one.
870    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    // Encode a function as its index in the stable `F::functions()` ordering.
879    // The ordering returned by `functions()` is a fixed `&'static` slice, so the
880    // index is stable across encode/decode. EV-04: the built-in
881    // `ArithmeticFunction::functions()` now lists every variant (including `Pow`,
882    // which was previously missing and silently collapsed to Add), so this path
883    // is lossless for the built-in type. A function absent from the set can now
884    // only arise from a *custom* `Function` impl that violates the `functions()`
885    // contract (it must enumerate every variant it can produce). We surface that
886    // as a debug-build panic to catch the bug during development, and fall back to
887    // index 0 in release rather than silently indexing out of range on decode.
888    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
924/// Trait for tree genome types (marker trait for operators)
925pub trait TreeGenomeType: EvolutionaryGenome {
926    /// The terminal type
927    type Term: Terminal;
928    /// The function type
929    type Func: Function;
930
931    /// Get the root of the tree
932    fn root(&self) -> &TreeNode<Self::Term, Self::Func>;
933
934    /// Get a mutable reference to the root
935    fn root_mut(&mut self) -> &mut TreeNode<Self::Term, Self::Func>;
936
937    /// Get the maximum depth
938    fn max_depth(&self) -> usize;
939
940    /// Create a new tree from a root node
941    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    /// Stack-safe teardown (EV-60).
967    ///
968    /// The compiler-generated drop glue for a recursive `enum` like [`TreeNode`]
969    /// recurses one stack frame per level, so dropping a pathologically deep tree
970    /// would overflow the stack. This impl instead frees the subtree with an
971    /// explicit work stack. It takes each node's children out with
972    /// [`std::mem::take`] (leaving an empty `Vec` behind) — which is permitted
973    /// under `Drop`, unlike moving a field out of `self` by value (E0509). Since
974    /// every node's `children` `Vec` is emptied *before* that node is dropped,
975    /// the reentrant `Drop::drop` invoked when the node itself is freed always
976    /// finds an empty `Vec` and returns in O(1); no deep recursion occurs.
977    fn drop(&mut self) {
978        // Only function nodes own children that could recurse.
979        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                // Detach grandchildren so `node` drops shallowly (its own
986                // reentrant Drop then finds an empty Vec).
987                stack.append(&mut std::mem::take(grandchildren));
988            }
989            // `node` drops here in O(1): a Terminal, or a Function whose
990            // children Vec is now empty.
991        }
992    }
993}
994
995/// Free a tree node and all of its descendants without deep recursion.
996///
997/// EV-60: [`TreeNode`] now has a stack-safe [`Drop`] impl, so simply dropping a
998/// node already frees an arbitrarily deep tree iteratively. This helper is
999/// retained as an explicit, self-documenting entry point (and for source
1000/// compatibility) but is no longer required to avoid a stack overflow — a plain
1001/// `drop(node)` or letting the node fall out of scope is equally safe.
1002pub 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        // Create: (+ x0 (* 1.0 x1))
1035        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); // root, left, right, right-left, right-right
1043        assert!(positions.contains(&vec![])); // root
1044        assert!(positions.contains(&vec![0])); // left child (x0)
1045        assert!(positions.contains(&vec![1])); // right child (mul)
1046        assert!(positions.contains(&vec![1, 0])); // mul's left child
1047        assert!(positions.contains(&vec![1, 1])); // mul's right child
1048    }
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        // Create: (+ x0 x1)
1064        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        // Create: (* (+ x0 1) x1) = (x0 + 1) * x1
1075        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); // (2 + 1) * 3 = 9
1083    }
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        // Full tree with target depth 3 creates: Function -> Function -> Function -> Terminal
1092        // Which has depth 4 (counting levels from root to leaf)
1093        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        // Grow can create trees up to max_depth + 1 levels (due to counting from 0)
1104        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        // regression: EV-04 — from_trace(to_trace(g)) must reproduce g *exactly*
1126        // (function identity and terminal values), not fabricate Add nodes and
1127        // fresh random terminals as the previous implementation did.
1128        use crate::genome::trace_genome::TraceGenome;
1129        let x0 = TreeNode::terminal(ArithmeticTerminal::Variable(0));
1130        let c1 = TreeNode::terminal(ArithmeticTerminal::Constant(2.5));
1131        // Use a non-Add function and mixed terminals to expose the old data loss.
1132        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        // Exact structural + semantic equality (variant identity, terminal values).
1144        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        // Evaluation results must agree across several inputs.
1150        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        // regression: EV-04 — a TreeGenome containing a `Pow` node must round-trip
1159        // losslessly. `Pow` was previously absent from ArithmeticFunction::functions(),
1160        // so encode_function's `.position(...).unwrap_or(0)` silently mapped it to
1161        // index 0 = Add, corrupting the tree with no error.
1162        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        // x^2 at x=3 must be 9, not (Add) 3+2=5 as the pre-fix collapse produced.
1179        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        // regression: EV-04 — prove the round-trip is lossless for EVERY public
1190        // ArithmeticFunction variant, not just the ones the generators draw from.
1191        // Each variant is exercised as a real node whose arity matches, and both
1192        // exact structural equality and evaluation equality are asserted.
1193        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        // functions() must contain every variant exactly once (index table).
1197        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                // Bit-identical (NaN-aware): the same function on the same inputs
1226                // must produce the same result, including matching NaN (e.g.
1227                // Pow(-1.5, 0.75)) which `==` would otherwise report as unequal.
1228                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        // EV-94: honest constructor takes an explicit maximum depth.
1239        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); // ramped/grow can reach max_depth (+1 level)
1244                                    // Degenerate depth is clamped to at least 1 and must not panic.
1245        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        // regression: EV-04 — encode_function relies on the stable ordering of
1253        // F::functions(); pin that ordering so encode/decode stays consistent.
1254        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        // Every function decodes back to itself from its own index.
1260        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        // regression: EV-04 — terminal encode/decode must be exact for every
1273        // variant, including distinguishing Constant from Erc.
1274        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        // regression: EV-60 — a ~100k-deep degenerate tree must evaluate, report
1290        // depth/size, and be torn down without overflowing the call stack. The
1291        // previous recursive eval/depth/size and the implicit recursive drop
1292        // would all overflow at this depth.
1293        let depth = 100_000usize;
1294        // Build bottom-up in a loop (no recursion during construction).
1295        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        // Iterative traversals must not overflow.
1303        assert_eq!(tree.size(), depth + 1);
1304        assert_eq!(tree.depth(), depth + 1);
1305        // Neg applied an even number of times to 1.0 yields +1.0.
1306        let value = tree.evaluate(&[]);
1307        assert!(value.is_finite());
1308        assert_eq!(value, 1.0);
1309
1310        // Iterative teardown must not overflow (implicit drop would recurse).
1311        tree.dismantle();
1312    }
1313
1314    #[test]
1315    fn test_drop_node_iteratively_frees_deep_tree() {
1316        // regression: EV-60 — the standalone iterative teardown handles a bare
1317        // deep TreeNode (not wrapped in a TreeGenome) without recursion.
1318        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        // Build bottom-up in a loop (no recursion during construction).
1328        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        // regression: EV-60 — dropping a ~100k-deep tree *implicitly* (never
1339        // calling dismantle()) must not overflow. The stack-safe Drop impl frees
1340        // it iteratively; the compiler-generated recursive drop glue would blow
1341        // the stack here.
1342        let depth = 100_000usize;
1343        {
1344            let tree = deep_tree(depth);
1345            assert_eq!(tree.size(), depth + 1);
1346            // Intentionally let `tree` fall out of scope here: implicit Drop only.
1347        }
1348        // A bare deep TreeNode dropped implicitly must also be safe.
1349        {
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; // dropped implicitly at end of scope
1356        }
1357    }
1358
1359    #[test]
1360    fn test_deep_tree_position_collectors_no_overflow() {
1361        // regression: EV-60 — the position collectors must be *iterative*
1362        // (explicit work stack), never recursive: calling
1363        // positions()/terminal_positions()/function_positions() on a
1364        // pathologically deep tree must not overflow the call stack.
1365        //
1366        // We prove that property directly and cheaply by running the collectors
1367        // inside a thread whose stack is capped at 64 KiB. A recursive collector
1368        // uses ~one call frame per tree level; at `depth = 2_000` that is 2000
1369        // frames, and even a minimal debug-build frame here (self ptr, the
1370        // `path` fat pointer, a `&mut Vec`, a per-node `child_path` local, saved
1371        // frame pointer + return address — well over the 64 KiB / 2000 ≈ 32 B
1372        // budget per frame, in practice ~100 B) blows past 64 KiB by several
1373        // multiples, so a recursive regression overflows and aborts the thread.
1374        // The iterative implementation uses a single frame plus a heap-allocated
1375        // work stack, so its stack footprint is independent of depth and fits
1376        // comfortably.
1377        //
1378        // `depth = 2_000` also keeps the collectors' *output* modest: each
1379        // returns the full root-path for every node, ~depth²/2 usizes ≈ 16 MB.
1380        // The former `depth = 100_000` produced Σ path lengths ≈ 5e9 usizes
1381        // ≈ 40 GB, which OOM-killed the 16 GB Linux CI runner (it only survived
1382        // on macOS because every path element is 0 and memory compression
1383        // flattened the pages). The positions() API is deliberately unchanged.
1384        let depth = 2_000usize;
1385
1386        // Build (and below, dismantle) the tree on the main thread. Construction
1387        // is an explicit bottom-up loop and Drop is iterative, so both are
1388        // stack-safe, but keeping them off the tiny stack isolates the property
1389        // under test to the collectors alone. The tree is `move`d into the
1390        // small-stack thread and handed straight back out so it is never dropped
1391        // on the 64 KiB stack.
1392        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                // One terminal (the single leaf) and `depth` function nodes.
1398                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        // Now tree should be (+ 5.0 x1)
1420        let tree = TreeGenome::new(add, 5);
1421        assert_eq!(tree.evaluate(&[0.0, 3.0]), 8.0); // 5 + 3 = 8
1422    }
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); // Protected
1440    }
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        // Create: (+ x0 (* 1.0 x1))
1457        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); // x0, 1.0, x1
1466
1467        let function_positions = add.function_positions();
1468        assert_eq!(function_positions.len(), 2); // add, mul
1469    }
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}