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::rc::Rc;
use std::fmt::Debug;

#[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>,
}

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)
            .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: Box::new(Self::transform(Rc::new((**callee).clone()), false)),
                args: Box::new(Self::transform(Rc::new((**args).clone()), false)),
            },

            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => BoundKind::If {
                cond: Box::new(Self::transform(Rc::new((**cond).clone()), false)),
                then_br: Box::new(Self::transform(
                    Rc::new((**then_br).clone()),
                    is_tail_position,
                )),
                else_br: else_br
                    .as_ref()
                    .map(|e| Box::new(Self::transform(Rc::new((**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(Self::transform(
                            Rc::new(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: Box::new(Self::transform(Rc::new((**value).clone()), false)),
            },
            BoundKind::DefLocal {
                name,
                slot,
                value,
                captured_by,
            } => BoundKind::DefLocal {
                name: name.clone(),
                slot: *slot,
                value: Box::new(Self::transform(Rc::new((**value).clone()), false)),
                captured_by: captured_by.clone(),
            },
            BoundKind::DefGlobal {
                name,
                global_index,
                value,
            } => BoundKind::DefGlobal {
                name: name.clone(),
                global_index: *global_index,
                value: Box::new(Self::transform(Rc::new((**value).clone()), false)),
            },
            BoundKind::Record { fields } => {
                let new_fields = fields
                    .iter()
                    .map(|(k, v)| {
                        (
                            Self::transform(Rc::new(k.clone()), false),
                            Self::transform(Rc::new(v.clone()), false),
                        )
                    })
                    .collect();
                BoundKind::Record { fields: new_fields }
            }
            BoundKind::Tuple { elements } => {
                let new_elements = elements
                    .iter()
                    .map(|e| Self::transform(Rc::new(e.clone()), false))
                    .collect();
                BoundKind::Tuple {
                    elements: new_elements,
                }
            }
            BoundKind::Parameter { name, slot } => BoundKind::Parameter {
                name: name.clone(),
                slot: *slot,
            },
            BoundKind::Constant(v) => BoundKind::Constant(v.clone()),
            BoundKind::Get { addr, name } => BoundKind::Get {
                addr: *addr,
                name: name.clone(),
            },
            BoundKind::Nop => BoundKind::Nop,
            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => BoundKind::Expansion {
                original_call: original_call.clone(),
                bound_expanded: Box::new(Self::transform(
                    Rc::new((**bound_expanded).clone()),
                    is_tail_position,
                )),
            },
            BoundKind::Extension(_) => BoundKind::Nop,
        };

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