RustAst/src/ast/compiler/tco.rs

use crate::ast::compiler::bound_nodes::{AnalyzedNode, BoundKind};
use crate::ast::nodes::Node;
use crate::ast::types::StaticType;
use std::fmt::Debug;
use std::rc::Rc;

#[derive(Clone)]
pub struct RuntimeMetadata {
    pub ty: StaticType,
    pub is_tail: bool,
    /// The analyzed node, containing metrics and a link to the original TypedNode.
    pub original: Rc<AnalyzedNode>,
    pub stack_size: u32,
}

impl Debug for RuntimeMetadata {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Metadata")
            .field("ty", &self.ty)
            .field("is_tail", &self.is_tail)
            .field("stack_size", &self.stack_size)
            .finish()
    }
}

/// The ExecNode is the AST used by the VM. It carries TCO flags and links to metrics.
pub type ExecNode = Node<BoundKind<RuntimeMetadata>, RuntimeMetadata>;

pub struct TCO;

impl TCO {
    /// Lowers an AnalyzedNode to an ExecNode and marks tail positions.
    pub fn optimize(node: AnalyzedNode) -> ExecNode {
        Self::transform(Rc::new(node), true)
    }

    fn transform(node_rc: Rc<AnalyzedNode>, is_tail_position: bool) -> ExecNode {
        let node = &*node_rc;
        let new_kind = match &node.kind {
            BoundKind::Call { callee, args } => BoundKind::Call {
                callee: Rc::new(Self::transform(callee.clone(), false)),
                args: Rc::new(Self::transform(args.clone(), false)),
            },

            BoundKind::Again { args } => {
                if !is_tail_position {
                    panic!("'again' is only allowed in tail position to avoid dead code.");
                }
                BoundKind::Again {
                    args: Rc::new(Self::transform(args.clone(), false)),
                }
            }
            BoundKind::Pipe {
                inputs,
                lambda,
                out_type,
            } => {
                let mut t_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    t_inputs.push(Rc::new(Self::transform(input.clone(), false)));
                }
                BoundKind::Pipe {
                    inputs: t_inputs,
                    lambda: Rc::new(Self::transform(lambda.clone(), false)),
                    out_type: out_type.clone(),
                }
            }

            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => BoundKind::If {
                cond: Rc::new(Self::transform(cond.clone(), false)),
                then_br: Rc::new(Self::transform(
                    then_br.clone(),
                    is_tail_position,
                )),
                else_br: else_br
                    .as_ref()
                    .map(|e| Rc::new(Self::transform(e.clone(), is_tail_position))),
            },

            BoundKind::Block { exprs } => {
                if exprs.is_empty() {
                    BoundKind::Block { exprs: vec![] }
                } else {
                    let last_idx = exprs.len() - 1;
                    let mut new_exprs = Vec::with_capacity(exprs.len());

                    for (i, expr) in exprs.iter().enumerate() {
                        let is_last = i == last_idx;
                        new_exprs.push(Rc::new(Self::transform(
                            expr.clone(),
                            is_tail_position && is_last,
                        )));
                    }
                    BoundKind::Block { exprs: new_exprs }
                }
            }

            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => BoundKind::Lambda {
                params: Rc::new(Self::transform(params.clone(), false)),
                upvalues: upvalues.clone(),
                body: Rc::new(Self::transform(body.clone(), true)),
                positional_count: *positional_count,
            },

            BoundKind::Set { addr, value } => BoundKind::Set {
                addr: *addr,
                value: Rc::new(Self::transform(value.clone(), false)),
            },
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                captured_by,
            } => BoundKind::Define {
                name: name.clone(),
                addr: *addr,
                kind: *kind,
                value: Rc::new(Self::transform(value.clone(), false)),
                captured_by: captured_by.clone(),
            },
            BoundKind::Destructure { pattern, value } => BoundKind::Destructure {
                pattern: Rc::new(Self::transform(pattern.clone(), false)),
                value: Rc::new(Self::transform(value.clone(), false)),
            },
            BoundKind::Record { layout, values } => {
                let new_values = values
                    .iter()
                    .map(|v| Rc::new(Self::transform(v.clone(), false)))
                    .collect();
                BoundKind::Record {
                    layout: layout.clone(),
                    values: new_values,
                }
            }
            BoundKind::Tuple { elements } => {
                let new_elements = elements
                    .iter()
                    .map(|e| Rc::new(Self::transform(e.clone(), false)))
                    .collect();
                BoundKind::Tuple {
                    elements: new_elements,
                }
            }
            BoundKind::Constant(v) => BoundKind::Constant(v.clone()),
            BoundKind::Get { addr, name } => BoundKind::Get {
                addr: *addr,
                name: name.clone(),
            },
            BoundKind::FieldAccessor(k) => BoundKind::FieldAccessor(*k),
            BoundKind::GetField { rec, field } => BoundKind::GetField {
                rec: Rc::new(Self::transform(rec.clone(), false)),
                field: *field,
            },
            BoundKind::Nop => BoundKind::Nop,
            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => BoundKind::Expansion {
                original_call: original_call.clone(),
                bound_expanded: Rc::new(Self::transform(
                    bound_expanded.clone(),
                    is_tail_position,
                )),
            },
            BoundKind::Extension(_) => BoundKind::Nop,
            BoundKind::Error => BoundKind::Error,
        };

        let stack_size = match &new_kind {
            BoundKind::Lambda { .. } => {
                // Pre-calculate stack size for the lambda body.
                // Note: 'node' here is AnalyzedNode (Node<BoundKind<NodeMetrics>, NodeMetrics>)
                // 'calculate_stack_size' works for any T.
                node.calculate_stack_size()
            }
            _ => 0,
        };

        Node {
            identity: node.identity.clone(),
            kind: new_kind,
            ty: RuntimeMetadata {
                ty: node.ty.original.ty.clone(),
                is_tail: is_tail_position,
                original: node_rc,
                stack_size,
            },
        }
    }
}