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, PartialOrd, Ord)]
pub struct LocalSlot(pub u32);

impl std::fmt::Display for LocalSlot {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "L{}", self.0)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct UpvalueIdx(pub u32);

impl std::fmt::Display for UpvalueIdx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "U{}", self.0)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct GlobalIdx(pub u32);

impl std::fmt::Display for GlobalIdx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "G{}", self.0)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Address {
    Local(LocalSlot),    // Stack-Slot index (relative to frame base)
    Upvalue(UpvalueIdx), // Index in the closure's upvalue array
    Global(GlobalIdx),   // 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()
    }
}

/// A bound AST node, decorated with type or metric information 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>;

/// Type alias for the global function registry.
pub type GlobalFunctionRegistry = std::collections::HashMap<GlobalIdx, Rc<BoundNode>>;

/// 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>;

/// Type alias for the global analyzed function registry.
pub type GlobalAnalyzedRegistry = std::collections::HashMap<GlobalIdx, Rc<AnalyzedNode>>;

#[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: Rc<BoundNode<T>>,
    },

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

    /// A first-class field accessor (e.g. .name)
    FieldAccessor(crate::ast::types::Keyword),

    /// Specialized field access (O(1) via RecordLayout)
    GetField {
        rec: Rc<BoundNode<T>>,
        field: crate::ast::types::Keyword,
    },

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

    /// A destructuring operation (can be a definition or an assignment depending on the pattern)
    Destructure {
        pattern: Rc<BoundNode<T>>,
        value: Rc<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: Rc<BoundNode<T>>,
        args: Rc<BoundNode<T>>,
    },

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

    Pipe {
        inputs: Vec<Rc<BoundNode<T>>>,
        lambda: Rc<BoundNode<T>>,
        out_type: crate::ast::types::StaticType,
    },
    Block {
        exprs: Vec<Rc<BoundNode<T>>>,
    },

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

    Record {
        layout: std::sync::Arc<crate::ast::types::RecordLayout>,
        values: Vec<Rc<BoundNode<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: Rc<BoundNode<T>>,
    },

    /// A diagnostic poison node, allowing compilation to continue after an error.
    Error,

    /// 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 && Rc::ptr_eq(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 && Rc::ptr_eq(va, vb) && ca == cb,
            (BoundKind::FieldAccessor(a), BoundKind::FieldAccessor(b)) => a == b,
            (
                BoundKind::GetField { rec: ra, field: fa },
                BoundKind::GetField { rec: rb, field: fb },
            ) => Rc::ptr_eq(ra, rb) && fa == fb,
            (
                BoundKind::If {
                    cond: ca,
                    then_br: ta,
                    else_br: ea,
                },
                BoundKind::If {
                    cond: cb,
                    then_br: tb,
                    else_br: eb,
                },
            ) => Rc::ptr_eq(ca, cb) && Rc::ptr_eq(ta, tb) && match (ea, eb) {
                (Some(a), Some(b)) => Rc::ptr_eq(a, b),
                (None, None) => true,
                _ => false,
            },
            (
                BoundKind::Destructure {
                    pattern: pa,
                    value: va,
                },
                BoundKind::Destructure {
                    pattern: pb,
                    value: vb,
                },
            ) => Rc::ptr_eq(pa, pb) && Rc::ptr_eq(va, vb),
            (
                BoundKind::Lambda {
                    params: pa,
                    upvalues: ua,
                    body: ba,
                    positional_count: pca,
                },
                BoundKind::Lambda {
                    params: pb,
                    upvalues: ub,
                    body: bb,
                    positional_count: pcb,
                },
            ) => Rc::ptr_eq(pa, pb) && ua == ub && Rc::ptr_eq(ba, bb) && pca == pcb,
            (
                BoundKind::Call {
                    callee: ca,
                    args: aa,
                },
                BoundKind::Call {
                    callee: cb,
                    args: ab,
                },
            ) => Rc::ptr_eq(ca, cb) && Rc::ptr_eq(aa, ab),
            (BoundKind::Again { args: aa }, BoundKind::Again { args: ab }) => Rc::ptr_eq(aa, ab),
            (BoundKind::Block { exprs: ea }, BoundKind::Block { exprs: eb }) => {
                ea.len() == eb.len() && ea.iter().zip(eb.iter()).all(|(a, b)| Rc::ptr_eq(a, b))
            }
            (BoundKind::Tuple { elements: ea }, BoundKind::Tuple { elements: eb }) => {
                ea.len() == eb.len() && ea.iter().zip(eb.iter()).all(|(a, b)| Rc::ptr_eq(a, b))
            }
            (
                BoundKind::Record {
                    layout: la,
                    values: va,
                },
                BoundKind::Record {
                    layout: lb,
                    values: vb,
                },
            ) => std::sync::Arc::ptr_eq(la, lb) && va.len() == vb.len() && va.iter().zip(vb.iter()).all(|(a, b)| Rc::ptr_eq(a, b)),
            (
                BoundKind::Expansion {
                    original_call: ca,
                    bound_expanded: ea,
                },
                BoundKind::Expansion {
                    original_call: cb,
                    bound_expanded: eb,
                },
            ) => Rc::ptr_eq(ca, cb) && Rc::ptr_eq(ea, eb),
            (BoundKind::Error, BoundKind::Error) => true,
            (BoundKind::Extension(_), BoundKind::Extension(_)) => false,
            _ => 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::FieldAccessor(k) => format!("FIELD_ACCESSOR(.{})", k.name()),
            BoundKind::GetField { field, .. } => format!("GET_FIELD(.{})", field.name()),
            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::Pipe { .. } => "PIPE".to_string(),
            BoundKind::Block { .. } => "BLOCK".to_string(),
            BoundKind::Tuple { elements } => format!("TUPLE({})", elements.len()),
            BoundKind::Record { values, .. } => format!("RECORD({})", values.len()),
            BoundKind::Expansion { .. } => "EXPANSION".to_string(),
            BoundKind::Extension(ext) => ext.display_name(),
            BoundKind::Error => "ERROR".to_string(),
        }
    }
}