RustAst/src/ast/compiler/lowering.rs

use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, ExecNode, Node, RuntimeMetadata, StackOffset, VirtualId};
use std::collections::HashMap;
use std::rc::Rc;

struct StackAllocator {
    mapping: HashMap<u32, u32>,
    next_slot: u32,
}

impl StackAllocator {
    fn new() -> Self {
        Self {
            mapping: HashMap::new(),
            next_slot: 0,
        }
    }

    fn map_slot(&mut self, slot: VirtualId) -> StackOffset {
        let entry = self.mapping.entry(slot.0).or_insert_with(|| {
            let s = self.next_slot;
            self.next_slot += 1;
            s
        });
        StackOffset(*entry)
    }

    fn map_address(&mut self, addr: Address<VirtualId>) -> Address<StackOffset> {
        match addr {
            Address::Local(slot) => Address::Local(self.map_slot(slot)),
            Address::Upvalue(idx) => Address::Upvalue(idx),
            Address::Global(idx) => Address::Global(idx),
        }
    }
}

pub struct Lowering;

impl Lowering {
    /// Lowers an AnalyzedNode to an ExecNode, marking tail positions and calculating stack sizes.
    pub fn lower(node: AnalyzedNode) -> ExecNode {
        let mut allocator = StackAllocator::new();
        let mut exec_node = Self::transform(Rc::new(node), true, &mut allocator);

        // If the top-level node is a Lambda, it already has its internal stack_size
        // calculated during transform(). For non-lambdas (like raw expressions),
        // we use the allocator's next_slot to determine the required root stack size.
        if !matches!(exec_node.kind, BoundKind::Lambda { .. }) {
            exec_node.ty.stack_size = allocator.next_slot;
        }
        exec_node
    }

    fn transform(
        node_rc: Rc<AnalyzedNode>,
        is_tail_position: bool,
        allocator: &mut StackAllocator,
    ) -> ExecNode {
        let node = &*node_rc;
        let mut lambda_stack_size = 0;

        let new_kind = match &node.kind {
            BoundKind::Call { callee, args } => BoundKind::Call {
                callee: Rc::new(Self::transform(callee.clone(), false, allocator)),
                args: Rc::new(Self::transform(args.clone(), false, allocator)),
            },

            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, allocator)),
                }
            }
            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, allocator)));
                }
                BoundKind::Pipe {
                    inputs: t_inputs,
                    lambda: Rc::new(Self::transform(lambda.clone(), false, allocator)),
                    out_type: out_type.clone(),
                }
            }

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

            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,
                            allocator,
                        )));
                    }
                    BoundKind::Block { exprs: new_exprs }
                }
            }

            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => {
                // Upvalues refer to the PARENT scope's addresses.
                let mapped_upvalues = upvalues
                    .iter()
                    .map(|a| allocator.map_address(*a))
                    .collect();

                // New allocator for the lambda's own stack frame
                let mut lambda_allocator = StackAllocator::new();
                let t_params = Rc::new(Self::transform(params.clone(), false, &mut lambda_allocator));
                let t_body = Rc::new(Self::transform(body.clone(), true, &mut lambda_allocator));
                lambda_stack_size = lambda_allocator.next_slot;

                BoundKind::Lambda {
                    params: t_params,
                    upvalues: mapped_upvalues,
                    body: t_body,
                    positional_count: *positional_count,
                }
            }

            BoundKind::Set { addr, value } => BoundKind::Set {
                addr: allocator.map_address(*addr),
                value: Rc::new(Self::transform(value.clone(), false, allocator)),
            },
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                captured_by,
            } => BoundKind::Define {
                name: name.clone(),
                addr: allocator.map_address(*addr),
                kind: *kind,
                value: Rc::new(Self::transform(value.clone(), false, allocator)),
                captured_by: captured_by.clone(),
            },
            BoundKind::Destructure { pattern, value } => BoundKind::Destructure {
                pattern: Rc::new(Self::transform(pattern.clone(), false, allocator)),
                value: Rc::new(Self::transform(value.clone(), false, allocator)),
            },
            BoundKind::Record { layout, values } => {
                let new_values = values
                    .iter()
                    .map(|v| Rc::new(Self::transform(v.clone(), false, allocator)))
                    .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, allocator)))
                    .collect();
                BoundKind::Tuple {
                    elements: new_elements,
                }
            }
            BoundKind::Constant(v) => BoundKind::Constant(v.clone()),
            BoundKind::Get { addr, name } => BoundKind::Get {
                addr: allocator.map_address(*addr),
                name: name.clone(),
            },
            BoundKind::FieldAccessor(k) => BoundKind::FieldAccessor(*k),
            BoundKind::GetField { rec, field } => BoundKind::GetField {
                rec: Rc::new(Self::transform(rec.clone(), false, allocator)),
                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,
                    allocator,
                )),
            },
            BoundKind::Extension(_) => BoundKind::Nop,
            BoundKind::Error => BoundKind::Error,
        };

        let stack_size = if let BoundKind::Lambda { .. } = &new_kind {
            lambda_stack_size
        } else {
            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,
            },
        }
    }
}