RustAst/src/ast/compiler/tco.rs

use std::rc::Rc;
use crate::ast::nodes::Node;
use crate::ast::compiler::bound_nodes::BoundKind;
use crate::ast::types::StaticType;

pub struct TCO;

impl TCO {
    pub fn optimize(node: Node<BoundKind, StaticType>) -> Node<BoundKind, StaticType> {
        // Top-level starts NOT in tail position usually, unless the script result is returned immediately
        // which implies tail position for the script body if viewed as a function.
        // Let's assume standard execution context (not tail call) for the script root.
        Self::transform(node, false)
    }

    fn transform(node: Node<BoundKind, StaticType>, is_tail_position: bool) -> Node<BoundKind, StaticType> {
        match node.kind {
            BoundKind::Call { callee, args } => {
                let new_callee = Box::new(Self::transform(*callee, false));
                let new_args: Vec<_> = args.into_iter().map(|arg| Self::transform(arg, false)).collect();
                
                if is_tail_position {
                    Node {
                        kind: BoundKind::TailCall {
                            callee: new_callee,
                            args: new_args,
                        },
                        ..node
                    }
                } else {
                    Node {
                        kind: BoundKind::Call {
                            callee: new_callee,
                            args: new_args,
                        },
                        ..node
                    }
                }
            },
            
            BoundKind::If { cond, then_br, else_br } => {
                let new_cond = Box::new(Self::transform(*cond, false));
                let new_then = Box::new(Self::transform(*then_br, is_tail_position));
                let new_else = else_br.map(|e| Box::new(Self::transform(*e, is_tail_position)));
                
                Node {
                    kind: BoundKind::If {
                        cond: new_cond,
                        then_br: new_then,
                        else_br: new_else,
                    },
                    ..node
                }
            },

            BoundKind::Block { exprs } => {
                if exprs.is_empty() { 
                     return Node {
                        kind: BoundKind::Block { exprs },
                        ..node
                    };
                }
                
                let last_idx = exprs.len() - 1;
                let mut new_exprs = Vec::with_capacity(exprs.len());
                
                for (i, expr) in exprs.into_iter().enumerate() {
                    let is_last = i == last_idx;
                    // Only the last expression inherits the tail context
                    new_exprs.push(Self::transform(expr, is_tail_position && is_last));
                }

                Node {
                    kind: BoundKind::Block { exprs: new_exprs },
                    ..node
                }
            },

            BoundKind::Lambda { param_count, upvalues, body } => {
                // The body of a lambda is implicitly in tail position when the lambda is called.
                // However, we process the lambda definition here.
                // The body node inside needs to be transformed assuming it IS in tail position relative to the function.
                let new_body = Rc::new(Self::transform((*body).clone(), true));
                
                Node {
                    kind: BoundKind::Lambda {
                        param_count,
                        upvalues,
                        body: new_body,
                    },
                    ..node
                }
            },

            BoundKind::Set { addr, value } => {
                Node {
                    kind: BoundKind::Set { addr, value: Box::new(Self::transform(*value, false)) },
                    ..node
                }
            },
            BoundKind::DefLocal { slot, value, is_boxed } => {
                 Node {
                    kind: BoundKind::DefLocal { slot, value: Box::new(Self::transform(*value, false)), is_boxed },
                    ..node
                }
            },
            BoundKind::DefGlobal { global_index, value } => {
                 Node {
                    kind: BoundKind::DefGlobal { global_index, value: Box::new(Self::transform(*value, false)) },
                    ..node
                }
            },
            BoundKind::Map { entries } => {
                let new_entries = entries.into_iter().map(|(k, v)| {
                    (Self::transform(k, false), Self::transform(v, false))
                }).collect();
                 Node {
                    kind: BoundKind::Map { entries: new_entries },
                    ..node
                }
            },
            BoundKind::Tuple { elements } => {
                let new_elements = elements.into_iter().map(|e| Self::transform(e, false)).collect();
                 Node {
                    kind: BoundKind::Tuple { elements: new_elements },
                    ..node
                }
            },
            
            // Atoms and variables don't change
            _ => node,
        }
    }
}