RustAst/src/ast/compiler/bound_nodes.rs

use crate::ast::nodes::{Node, Symbol};
use crate::ast::types::{Identity, StaticType, Value};
use std::rc::Rc;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Address {
    Local(u32),   // Stack-Slot index (relative to frame base)
    Upvalue(u32), // Index in the closure's upvalue array
    Global(u32),  // Index in the global environment vector
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DeclarationKind {
    Variable,
    Parameter,
}

/// Trait for DSL-specific AST extensions (like Pipe, Series-Access, etc.)
pub trait BoundExtension<T>: std::fmt::Debug {
    fn clone_box(&self) -> Box<dyn BoundExtension<T>>;
    fn display_name(&self) -> String;
}

impl<T> Clone for Box<dyn BoundExtension<T>> {
    fn clone(&self) -> Self {
        self.clone_box()
    }
}

/// Type alias for a node that has been bound. Defaults to untyped (T = ()).
pub type BoundNode<T = ()> = Node<BoundKind<T>, T>;

/// Type alias for a node that has been fully type-checked.
pub type TypedNode = BoundNode<StaticType>;

/// Metrics collected during the analysis phase.
#[derive(Debug, Clone, PartialEq)]
pub struct NodeMetrics {
    pub original: Rc<TypedNode>,
    pub purity: crate::ast::types::Purity,
    pub is_recursive: bool,
}

/// Type alias for a node that has been analyzed.
pub type AnalyzedNode = BoundNode<NodeMetrics>;

#[derive(Debug, Clone)]
pub enum BoundKind<T = ()> {
    Nop,
    Constant(Value),

    // Variable Access (Resolved)
    Get {
        addr: Address,
        name: Symbol,
    },

    // Variable Update (Assignment)
    Set {
        addr: Address,
        value: Box<BoundNode<T>>,
    },

    // Variable Declaration (Unified for Local/Global/Parameter)
    Define {
        name: Symbol,
        addr: Address,
        kind: DeclarationKind,
        value: Box<BoundNode<T>>,
        captured_by: Vec<Identity>,
    },

    If {
        cond: Box<BoundNode<T>>,
        then_br: Box<BoundNode<T>>,
        else_br: Option<Box<BoundNode<T>>>,
    },

    /// A destructuring operation (can be a definition or an assignment depending on the pattern)
    Destructure {
        pattern: Box<BoundNode<T>>,
        value: Box<BoundNode<T>>,
    },

    Lambda {
        params: Rc<BoundNode<T>>,
        // The list of variables captured from enclosing scopes
        upvalues: Vec<Address>,
        body: Rc<BoundNode<T>>,
        /// Static optimization: number of positional parameters if the pattern is flat.
        positional_count: Option<u32>,
    },

    Call {
        callee: Box<BoundNode<T>>,
        args: Box<BoundNode<T>>,
    },

    Again {
        args: Box<BoundNode<T>>,
    },

    Block {
        exprs: Vec<BoundNode<T>>,
    },

    Tuple {
        elements: Vec<BoundNode<T>>,
    },

    Record {
        fields: Vec<RecordField<T>>,
    },

    /// An expanded macro call, preserving the original call for debugging and UI.
    Expansion {
        /// The original call from the untyped AST.
        original_call: Rc<Node<crate::ast::nodes::UntypedKind>>,
        /// The result of binding the expanded AST.
        bound_expanded: Box<BoundNode<T>>,
    },

    /// DSL-specific extension slot
    Extension(Box<dyn BoundExtension<T>>),
}

impl<T> PartialEq for BoundKind<T>
where
    T: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (BoundKind::Nop, BoundKind::Nop) => true,
            (BoundKind::Constant(a), BoundKind::Constant(b)) => a == b,
            (BoundKind::Get { addr: aa, name: na }, BoundKind::Get { addr: ab, name: nb }) => {
                aa == ab && na == nb
            }
            (
                BoundKind::Set {
                    addr: aa,
                    value: va,
                },
                BoundKind::Set {
                    addr: ab,
                    value: vb,
                },
            ) => aa == ab && va == vb,
            (
                BoundKind::Define {
                    name: na,
                    addr: aa,
                    kind: ka,
                    value: va,
                    captured_by: ca,
                },
                BoundKind::Define {
                    name: nb,
                    addr: ab,
                    kind: kb,
                    value: vb,
                    captured_by: cb,
                },
            ) => na == nb && aa == ab && ka == kb && va == vb && ca == cb,
            (
                BoundKind::If {
                    cond: ca,
                    then_br: ta,
                    else_br: ea,
                },
                BoundKind::If {
                    cond: cb,
                    then_br: tb,
                    else_br: eb,
                },
            ) => ca == cb && ta == tb && ea == eb,
            (
                BoundKind::Destructure {
                    pattern: pa,
                    value: va,
                },
                BoundKind::Destructure {
                    pattern: pb,
                    value: vb,
                },
            ) => pa == pb && va == vb,
            (
                BoundKind::Lambda {
                    params: pa,
                    upvalues: ua,
                    body: ba,
                    positional_count: pca,
                },
                BoundKind::Lambda {
                    params: pb,
                    upvalues: ub,
                    body: bb,
                    positional_count: pcb,
                },
            ) => pa == pb && ua == ub && ba == bb && pca == pcb,
            (
                BoundKind::Call {
                    callee: ca,
                    args: aa,
                },
                BoundKind::Call {
                    callee: cb,
                    args: ab,
                },
            ) => ca == cb && aa == ab,
            (BoundKind::Again { args: aa }, BoundKind::Again { args: ab }) => aa == ab,
            (BoundKind::Block { exprs: ea }, BoundKind::Block { exprs: eb }) => ea == eb,
            (BoundKind::Tuple { elements: ea }, BoundKind::Tuple { elements: eb }) => ea == eb,
            (BoundKind::Record { fields: fa }, BoundKind::Record { fields: fb }) => fa == fb,
            (
                BoundKind::Expansion {
                    original_call: oa,
                    bound_expanded: ba,
                },
                BoundKind::Expansion {
                    original_call: ob,
                    bound_expanded: bb,
                },
            ) => Rc::ptr_eq(oa, ob) && ba == bb,
            (BoundKind::Extension(_), BoundKind::Extension(_)) => false, // Currently no equality for extensions
            _ => false,
        }
    }
}

/// A single field in a Record literal (Key-Value pair)
pub type RecordField<T> = (BoundNode<T>, BoundNode<T>);

impl<T> BoundKind<T> {
    pub fn display_name(&self) -> String {
        match self {
            BoundKind::Nop => "NOP".to_string(),
            BoundKind::Constant(v) => format!("CONST({})", v),
            BoundKind::Get { addr, name } => format!("GET({}, {:?})", name.name, addr),
            BoundKind::Set { addr, .. } => format!("SET({:?})", addr),
            BoundKind::Define {
                name, addr, kind, ..
            } => {
                let k_str = match kind {
                    DeclarationKind::Variable => "VAR",
                    DeclarationKind::Parameter => "PARAM",
                };
                format!("DEF_{}({}, {:?})", k_str, name.name, addr)
            }
            BoundKind::If { .. } => "IF".to_string(),
            BoundKind::Destructure { .. } => "DESTRUCTURE".to_string(),
            BoundKind::Lambda {
                params, upvalues, ..
            } => {
                let p_str = match &params.kind {
                    BoundKind::Tuple { elements } => format!("p:{}", elements.len()),
                    _ => "p:1".to_string(),
                };
                format!("LAMBDA({}, Captures:{})", p_str, upvalues.len())
            }
            BoundKind::Call { .. } => "CALL".to_string(),
            BoundKind::Again { .. } => "AGAIN".to_string(),
            BoundKind::Block { .. } => "BLOCK".to_string(),
            BoundKind::Tuple { elements } => format!("TUPLE({})", elements.len()),
            BoundKind::Record { fields } => format!("RECORD({})", fields.len()),
            BoundKind::Expansion { .. } => "EXPANSION".to_string(),
            BoundKind::Extension(ext) => ext.display_name(),
        }
    }
}