RustAst/src/ast/compiler/binder.rs

use crate::ast::nodes::{
    Address, AssignBinding, BoundPhase, DeclarationKind, DefBinding, GlobalIdx,
    IdentifierBinding, LambdaBinding, Node, NodeKind, SyntaxKind, SyntaxNode, Symbol,
    UpvalueIdx, VirtualId,
};
use crate::ast::diagnostics::Diagnostics;
use crate::ast::types::{Identity, NodeIdentity, Purity, RecordLayout, SourceLocation, StaticType, Value};
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Debug, Clone)]
pub struct LocalInfo {
    pub addr: Address<VirtualId>,
    pub identity: Identity,
    pub _ty: StaticType,
    pub purity: Purity,
}

#[derive(Debug, Clone)]
pub struct CompilerScope {
    pub locals: HashMap<Symbol, LocalInfo>,
}

impl Default for CompilerScope {
    fn default() -> Self {
        Self::new()
    }
}

impl CompilerScope {
    pub fn new() -> Self {
        Self {
            locals: HashMap::new(),
        }
    }
}

struct FunctionCompiler {
    identity: Identity,
    scopes: Vec<CompilerScope>,
    slot_count: u32,
    upvalues: Vec<Address<VirtualId>>,
}

impl FunctionCompiler {
    fn new(identity: Identity, initial_scopes: Vec<CompilerScope>, initial_slot_count: u32) -> Self {
        Self {
            identity,
            scopes: if initial_scopes.is_empty() {
                vec![CompilerScope::new()]
            } else {
                initial_scopes
            },
            slot_count: initial_slot_count,
            upvalues: Vec::new(),
        }
    }

    fn push_scope(&mut self) {
        self.scopes.push(CompilerScope::new());
    }

    fn pop_scope(&mut self) {
        if self.scopes.len() > 1 {
            self.scopes.pop();
        }
    }

    fn define_variable(&mut self, name: &Symbol, identity: Identity) -> Result<Address<VirtualId>, String> {
        let current_scope = self.scopes.last_mut().unwrap();
        if current_scope.locals.contains_key(name) {
            return Err(format!(
                "Variable '{}' is already defined in this scope level.",
                name.name
            ));
        }

        let slot = VirtualId(self.slot_count);
        current_scope.locals.insert(
            name.clone(),
            LocalInfo {
                addr: Address::Local(slot),
                identity,
                _ty: StaticType::Any,
                purity: Purity::Impure,
            },
        );
        self.slot_count += 1;
        Ok(Address::Local(slot))
    }

    fn resolve_local(&self, sym: &Symbol) -> Option<(LocalInfo, usize)> {
        for (idx, scope) in self.scopes.iter().enumerate().rev() {
            if let Some(info) = scope.locals.get(sym) {
                return Some((info.clone(), idx));
            }
        }
        None
    }

    fn add_upvalue(&mut self, addr: Address<VirtualId>) -> UpvalueIdx {
        if let Some(idx) = self.upvalues.iter().position(|&a| a == addr) {
            return UpvalueIdx(idx as u32);
        }
        let idx = UpvalueIdx(self.upvalues.len() as u32);
        self.upvalues.push(addr);
        idx
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExprContext {
    Expression,
    Statement,
}

pub type BindingResult = (
    Node<BoundPhase>,
    HashMap<Identity, Vec<Identity>>,
    Vec<CompilerScope>,
    u32,
);

pub struct Binder {
    functions: Vec<FunctionCompiler>,
    capture_map: HashMap<Identity, std::collections::HashSet<Identity>>,
    fixed_scope_idx: i32,
}

impl Binder {
    pub fn new(
        initial_scopes: Vec<CompilerScope>,
        initial_slot_count: u32,
        fixed_scope_idx: i32,
    ) -> Self {
        let mut binder = Self {
            functions: Vec::new(),
            capture_map: HashMap::new(),
            fixed_scope_idx,
        };
        binder.functions.push(FunctionCompiler::new(
            NodeIdentity::new(SourceLocation {
                line: 0,
                col: 0,
            }),
            initial_scopes,
            initial_slot_count,
        ));
        binder
    }

    pub fn bind_root(
        initial_scopes: Vec<CompilerScope>,
        initial_slot_count: u32,
        fixed_scope_idx: i32,
        node: &SyntaxNode,
        diagnostics: &mut Diagnostics,
    ) -> Result<BindingResult, String> {
        let mut binder = Self::new(initial_scopes, initial_slot_count, fixed_scope_idx);
        let bound = binder.bind(node, ExprContext::Expression, diagnostics);

        let final_captures = binder
            .capture_map
            .into_iter()
            .map(|(k, v)| (k, v.into_iter().collect()))
            .collect();

        let root_compiler = binder.functions.pop().unwrap();
        Ok((
            bound,
            final_captures,
            root_compiler.scopes,
            root_compiler.slot_count,
        ))
    }

    fn declare_variable(
        &mut self,
        name: &Symbol,
        identity: Identity,
        _kind: DeclarationKind,
        diag: &mut Diagnostics,
    ) -> Option<Address<VirtualId>> {
        let current_fn_idx = self.functions.len() - 1;

        if current_fn_idx == 0 {
            let current_scope_idx = self.functions[0].scopes.len() as i32 - 1;
            if current_scope_idx <= self.fixed_scope_idx {
                diag.push_error(
                    format!(
                        "Cannot define '{}': Scope {} is frozen/immutable.",
                        name.name, current_scope_idx
                    ),
                    None,
                );
                return None;
            }
        }

        let current_fn = self.functions.last_mut().unwrap();
        match current_fn.define_variable(name, identity) {
            Ok(addr) => Some(addr),
            Err(e) => {
                diag.push_error(e, None);
                None
            }
        }
    }

    pub fn bind(
        &mut self,
        node: &SyntaxNode,
        ctx: ExprContext,
        diag: &mut Diagnostics,
    ) -> Node<BoundPhase> {
        match &node.kind {
            SyntaxKind::Nop => self.make_node(node.identity.clone(), NodeKind::Nop),
            SyntaxKind::Constant(v) => {
                self.make_node(node.identity.clone(), NodeKind::Constant(v.clone()))
            }

            SyntaxKind::Identifier { symbol: sym, .. } => {
                if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) {
                    self.make_node(
                        node.identity.clone(),
                        NodeKind::Identifier {
                            symbol: sym.clone(),
                            binding: IdentifierBinding::Reference(addr),
                        },
                    )
                } else {
                    self.make_node(node.identity.clone(), NodeKind::Error)
                }
            }

            SyntaxKind::FieldAccessor(k) => {
                self.make_node(node.identity.clone(), NodeKind::FieldAccessor(*k))
            }

            SyntaxKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = self.bind(cond.as_ref(), ExprContext::Expression, diag);

                self.functions.last_mut().unwrap().push_scope();
                let then_br = self.bind(then_br.as_ref(), ctx, diag);
                self.functions.last_mut().unwrap().pop_scope();

                let mut else_br_bound = None;
                if let Some(e) = else_br {
                    self.functions.last_mut().unwrap().push_scope();
                    else_br_bound = Some(Rc::new(self.bind(e, ctx, diag)));
                    self.functions.last_mut().unwrap().pop_scope();
                }

                self.make_node(
                    node.identity.clone(),
                    NodeKind::If {
                        cond: Rc::new(cond),
                        then_br: Rc::new(then_br),
                        else_br: else_br_bound,
                    },
                )
            }

            SyntaxKind::Def { pattern: target, value, .. } => {
                if ctx == ExprContext::Expression {
                    diag.push_error(
                        "Statement 'def' cannot be used as an expression.",
                        Some(node.identity.clone()),
                    );
                }
                if let SyntaxKind::Identifier { symbol: ref name, .. } = target.kind {
                    let addr_opt = self.declare_variable(
                        name,
                        node.identity.clone(),
                        DeclarationKind::Variable,
                        diag,
                    );
                    let val_node = self.bind(value, ExprContext::Expression, diag);

                    if let Some(addr) = addr_opt {
                        let pattern_node = self.make_node(
                            target.identity.clone(),
                            NodeKind::Identifier {
                                symbol: name.clone(),
                                binding: IdentifierBinding::Declaration {
                                    addr,
                                    kind: DeclarationKind::Variable,
                                },
                            },
                        );
                        self.make_node(
                            node.identity.clone(),
                            NodeKind::Def {
                                pattern: Rc::new(pattern_node),
                                value: Rc::new(val_node),
                                info: DefBinding {
                                    captured_by: Vec::new(),
                                },
                            },
                        )
                    } else {
                        self.make_node(node.identity.clone(), NodeKind::Error)
                    }
                } else {
                    let val_node = self.bind(value, ExprContext::Expression, diag);
                    let target_node = self.bind_pattern(target.as_ref(), DeclarationKind::Variable, diag);

                    self.make_node(
                        node.identity.clone(),
                        NodeKind::Def {
                            pattern: Rc::new(target_node),
                            value: Rc::new(val_node),
                            info: DefBinding {
                                captured_by: Vec::new(),
                            },
                        },
                    )
                }
            }

            SyntaxKind::Assign { target, value, .. } => {
                let val_node = self.bind(value, ExprContext::Expression, diag);

                if let SyntaxKind::Identifier { symbol: sym, .. } = &target.kind {
                    if let Some(addr) = self.resolve_variable(sym, diag, &target.identity) {
                        let target_node = self.make_node(
                            target.identity.clone(),
                            NodeKind::Identifier {
                                symbol: sym.clone(),
                                binding: IdentifierBinding::Reference(addr),
                            },
                        );
                        self.make_node(
                            node.identity.clone(),
                            NodeKind::Assign {
                                target: Rc::new(target_node),
                                value: Rc::new(val_node),
                                info: AssignBinding { addr: Some(addr) },
                            },
                        )
                    } else {
                        self.make_node(node.identity.clone(), NodeKind::Error)
                    }
                } else {
                    let target_node = self.bind_assign_pattern(target.as_ref(), diag);
                    self.make_node(
                        node.identity.clone(),
                        NodeKind::Assign {
                            target: Rc::new(target_node),
                            value: Rc::new(val_node),
                            info: AssignBinding { addr: None },
                        },
                    )
                }
            }

            SyntaxKind::Lambda { params, body, .. } => {
                let identity = node.identity.clone();
                self.functions
                    .push(FunctionCompiler::new(identity.clone(), vec![], 0));

                let params_bound = self.bind_pattern(params.as_ref(), DeclarationKind::Parameter, diag);
                let body_bound = self.bind(body.as_ref(), ExprContext::Expression, diag);
                let compiled_fn = self.functions.pop().unwrap();

                fn count_params(node: &Node<BoundPhase>) -> Option<u32> {
                    match &node.kind {
                        NodeKind::Identifier {
                            binding: IdentifierBinding::Declaration {
                                kind: DeclarationKind::Parameter,
                                ..
                            },
                            ..
                        } => Some(1),
                        NodeKind::Tuple { elements } => {
                            let mut total = 0;
                            for e in elements {
                                total += count_params(e)?;
                            }
                            Some(total)
                        }
                        NodeKind::Nop => Some(0),
                        _ => None,
                    }
                }

                let positional_count = count_params(&params_bound);

                self.make_node(
                    identity,
                    NodeKind::Lambda {
                        params: Rc::new(params_bound),
                        body: Rc::new(body_bound),
                        info: LambdaBinding {
                            upvalues: compiled_fn.upvalues,
                            positional_count,
                        },
                    },
                )
            }

            SyntaxKind::Call { callee, args } => {
                let callee = self.bind(callee, ExprContext::Expression, diag);
                let args = self.bind(args.as_ref(), ExprContext::Expression, diag);

                self.make_node(
                    node.identity.clone(),
                    NodeKind::Call {
                        callee: Rc::new(callee),
                        args: Rc::new(args),
                    },
                )
            }

            SyntaxKind::Again { args } => {
                if self.functions.len() <= 1 {
                    diag.push_error(
                        "'again' is only allowed inside a function or lambda.",
                        Some(node.identity.clone()),
                    );
                    return self.make_node(node.identity.clone(), NodeKind::Error);
                }
                let args = self.bind(args.as_ref(), ExprContext::Expression, diag);
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Again {
                        args: Rc::new(args),
                    },
                )
            }

            SyntaxKind::Block { exprs } => {
                self.functions.last_mut().unwrap().push_scope();
                let mut bound_exprs = Vec::new();
                for (i, expr) in exprs.iter().enumerate() {
                    let expr_ctx = if i == exprs.len() - 1 {
                        ctx
                    } else {
                        ExprContext::Statement
                    };
                    bound_exprs.push(Rc::new(self.bind(expr, expr_ctx, diag)));
                }
                self.functions.last_mut().unwrap().pop_scope();
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Block { exprs: bound_exprs },
                )
            }

            SyntaxKind::Tuple { elements } => {
                let mut bound_elems = Vec::new();
                for e in elements {
                    bound_elems.push(Rc::new(self.bind(e, ExprContext::Expression, diag)));
                }
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Tuple {
                        elements: bound_elems,
                    },
                )
            }

            SyntaxKind::Record { fields, .. } => {
                let mut bound_fields = Vec::new();
                let mut layout_fields = Vec::new();

                for (k, v) in fields {
                    let key_node = self.bind(k, ExprContext::Expression, diag);
                    let val_node = self.bind(v, ExprContext::Expression, diag);

                    if let NodeKind::Constant(Value::Keyword(kw)) =
                        &key_node.kind
                    {
                        layout_fields.push((*kw, StaticType::Any));
                    } else {
                        diag.push_error(
                            format!(
                                "Record keys must be keywords, found at {:?}",
                                key_node.identity.location
                            ),
                            Some(key_node.identity.clone()),
                        );
                    }
                    bound_fields.push((Rc::new(key_node), Rc::new(val_node)));
                }

                let layout = RecordLayout::get_or_create(layout_fields);

                self.make_node(
                    node.identity.clone(),
                    NodeKind::Record {
                        fields: bound_fields,
                        layout,
                    },
                )
            }

            SyntaxKind::Expansion { original_call, expanded } => {
                let bound_expanded = self.bind(expanded.as_ref(), ctx, diag);
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Expansion {
                        original_call: original_call.clone(),
                        expanded: Rc::new(bound_expanded),
                    },
                )
            }

            SyntaxKind::Template(_)
            | SyntaxKind::Placeholder(_)
            | SyntaxKind::Splice(_)
            | SyntaxKind::MacroDecl { .. } => {
                diag.push_error(
                    format!(
                        "Macro construct {:?} found in Binder. Macros must be expanded before binding.",
                        node.kind
                    ),
                    Some(node.identity.clone()),
                );
                self.make_node(node.identity.clone(), NodeKind::Error)
            }

            SyntaxKind::Extension(_) => {
                diag.push_error(
                    "Custom extensions not supported in Binder yet",
                    Some(node.identity.clone()),
                );
                self.make_node(node.identity.clone(), NodeKind::Error)
            }
            SyntaxKind::GetField { .. } => {
                diag.push_error(
                    "GetField not expected in SyntaxPhase",
                    Some(node.identity.clone()),
                );
                self.make_node(node.identity.clone(), NodeKind::Error)
            }
            SyntaxKind::Error => Node {
                comments: node.comments.clone(),
                identity: node.identity.clone(),
                kind: NodeKind::Error,
                ty: (),
            },
        }
    }

    fn resolve_variable(
        &mut self,
        sym: &Symbol,
        diag: &mut Diagnostics,
        identity: &Identity,
    ) -> Option<Address<VirtualId>> {
        let current_fn_idx = self.functions.len() - 1;

        // 1. Try local in current function
        if let Some((info, _)) = self.functions[current_fn_idx].resolve_local(sym) {
            return Some(info.addr);
        }

        // 2. Try enclosing scopes (capture chain)
        for i in (0..current_fn_idx).rev() {
            if let Some((info, scope_idx)) = self.functions[i].resolve_local(sym) {
                // If it's in the root script and within a frozen scope, translate to Global
                let is_frozen_root = i == 0 && (scope_idx as i32) <= self.fixed_scope_idx;

                if let Address::Global(_) = info.addr {
                    return Some(info.addr);
                }

                let mut addr = info.addr;

                if is_frozen_root
                    && let Address::Local(slot) = addr {
                    addr = Address::Global(GlobalIdx(slot.0));
                }

                if let Address::Global(_) = addr {
                    return Some(addr);
                }

                // Record the capture for each lambda level in between
                for k in (i + 1)..=current_fn_idx {
                    let lambda_id = self.functions[k].identity.clone();
                    self.capture_map
                        .entry(info.identity.clone())
                        .or_default()
                        .insert(lambda_id);
                    addr = Address::Upvalue(self.functions[k].add_upvalue(addr));
                }
                return Some(addr);
            }
        }

        // 3. Global Fallback (search in root level with context removed)
        if sym.context.is_some() {
            let fallback_sym = Symbol {
                name: sym.name.clone(),
                context: None,
            };
            // Search again in all functions, primarily we care about Root Scopes
            for i in (0..=current_fn_idx).rev() {
                if let Some((info, _)) = self.functions[i].resolve_local(&fallback_sym)
                    && let Address::Global(_) = info.addr
                {
                    return Some(info.addr);
                }
            }
        }

        diag.push_error(
            format!("Undefined variable '{}'", sym.name),
            Some(identity.clone()),
        );
        None
    }

    fn bind_pattern(
        &mut self,
        node: &SyntaxNode,
        kind: DeclarationKind,
        diag: &mut Diagnostics,
    ) -> Node<BoundPhase> {
        match &node.kind {
            SyntaxKind::Identifier { symbol: sym, .. } => {
                if let Some(addr) = self.declare_variable(sym, node.identity.clone(), kind, diag) {
                    self.make_node(
                        node.identity.clone(),
                        NodeKind::Identifier {
                            symbol: sym.clone(),
                            binding: IdentifierBinding::Declaration { addr, kind },
                        },
                    )
                } else {
                    self.make_node(node.identity.clone(), NodeKind::Error)
                }
            }
            SyntaxKind::Tuple { elements } => {
                let mut bound_elems = Vec::new();
                for e in elements {
                    bound_elems.push(Rc::new(self.bind_pattern(e, kind, diag)));
                }
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Tuple {
                        elements: bound_elems,
                    },
                )
            }
            _ => {
                diag.push_error(
                    format!("Invalid node in pattern: {:?}", node.kind),
                    Some(node.identity.clone()),
                );
                self.make_node(node.identity.clone(), NodeKind::Error)
            }
        }
    }

    fn bind_assign_pattern(
        &mut self,
        node: &SyntaxNode,
        diag: &mut Diagnostics,
    ) -> Node<BoundPhase> {
        match &node.kind {
            SyntaxKind::Identifier { symbol: sym, .. } => {
                if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) {
                    self.make_node(
                        node.identity.clone(),
                        NodeKind::Identifier {
                            symbol: sym.clone(),
                            binding: IdentifierBinding::Reference(addr),
                        },
                    )
                } else {
                    self.make_node(node.identity.clone(), NodeKind::Error)
                }
            }
            SyntaxKind::Tuple { elements } => {
                let mut bound_elems = Vec::new();
                for e in elements {
                    bound_elems.push(Rc::new(self.bind_assign_pattern(e, diag)));
                }
                self.make_node(
                    node.identity.clone(),
                    NodeKind::Tuple {
                        elements: bound_elems,
                    },
                )
            }
            _ => {
                diag.push_error(
                    format!("Invalid node in assignment pattern: {:?}", node.kind),
                    Some(node.identity.clone()),
                );
                self.make_node(node.identity.clone(), NodeKind::Error)
            }
        }
    }

    fn make_node(&self, identity: Identity, kind: NodeKind<BoundPhase>) -> Node<BoundPhase> {
        Node {
            comments: std::rc::Rc::from([]),
            identity,
            kind,
            ty: (),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::diagnostics::Diagnostics;
    use crate::ast::parser::Parser;

    #[test]
    fn test_upvalue_capture_sets_is_boxed() {
        let source = "(fn [] (do (def x 10) (def f (fn [] x)) x))";
        let mut parser = Parser::new(source);
        let syntax = parser.parse_expression();

        let mut diagnostics = Diagnostics::new();
        let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap();

        if let NodeKind::Lambda { body, .. } = &bound.kind {
            if let NodeKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
                if let NodeKind::Def { pattern, .. } = &x_decl.kind {
                    if let NodeKind::Identifier { binding: IdentifierBinding::Declaration { addr, .. }, .. } = &pattern.kind {
                        assert!(matches!(addr, Address::Local(_)));
                        assert!(
                            captures.contains_key(&x_decl.identity),
                            "Variable 'x' should have capturers"
                        );
                    } else {
                        panic!("Def pattern should be an Identifier with Declaration binding");
                    }
                } else {
                    panic!("First expression should be Def");
                }
            } else {
                panic!("Lambda body should be a Block");
            }
        } else {
            panic!("Root should be a Lambda");
        }
    }

    #[test]
    fn test_no_capture_not_boxed() {
        let source = "(fn [] (do (def x 10) x))";
        let mut parser = Parser::new(source);
        let syntax = parser.parse_expression();

        let mut diagnostics = Diagnostics::new();
        let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap();

        if let NodeKind::Lambda { body, .. } = &bound.kind {
            if let NodeKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
                if let NodeKind::Def { pattern, .. } = &x_decl.kind {
                    if let NodeKind::Identifier { binding: IdentifierBinding::Declaration { addr, .. }, .. } = &pattern.kind {
                        assert!(matches!(addr, Address::Local(_)));
                        assert!(
                            !captures.contains_key(&x_decl.identity),
                            "Variable 'x' should NOT have any capturers"
                        );
                    } else {
                        panic!("Def pattern should be an Identifier with Declaration binding");
                    }
                } else {
                    panic!("First expression should be Def");
                }
            } else {
                panic!("Lambda body should be a Block");
            }
        } else {
            panic!("Root should be a Lambda");
        }
    }

    #[test]
    fn test_redefinition_error() {
        let source = "(do (def x 1) (def x 2) x)";
        let mut parser = Parser::new(source);
        let syntax = parser.parse_expression();

        let mut diagnostics = Diagnostics::new();
        let _ = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics);

        assert!(diagnostics.has_errors());
        assert!(diagnostics.items.iter().any(|i| i.message.contains("already defined")));
    }

    #[test]
    fn test_repro_global_redefinition() {
        let source1 = "(def x 1) 1";
        let syntax1 = Parser::new(source1).parse_expression();
        let mut diagnostics = Diagnostics::new();
        let (_, _, scopes1, slots1) = Binder::bind_root(vec![], 0, -1, &syntax1, &mut diagnostics).unwrap();

        let source2 = "(def x 2) 2";
        let syntax2 = Parser::new(source2).parse_expression();
        let mut diagnostics2 = Diagnostics::new();
        // Here we simulate frozen scope by passing fixed_scope_idx = 0
        let _ = Binder::bind_root(scopes1, slots1, 0, &syntax2, &mut diagnostics2);

        assert!(diagnostics2.has_errors());
        assert!(diagnostics2.items.iter().any(|i| i.message.contains("frozen/immutable")));
    }
}