RustAst/src/ast/compiler/binder.rs

use crate::ast::compiler::bound_nodes::{Address, BoundKind, BoundNode};
use crate::ast::nodes::{Node, Symbol, UntypedKind};
use crate::ast::types::{Identity, StaticType};
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Debug, Clone)]
struct LocalInfo {
    slot: u32,
    // Note: Binder doesn't strictly need the type anymore,
    // but it might be useful for built-ins during resolution.
    // For now we keep it as Any or Unknown.
    _ty: StaticType,
}

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

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

    fn define(&mut self, sym: &Symbol) -> Result<u32, String> {
        if self.locals.contains_key(sym) {
            return Err(format!(
                "Variable '{}' is already defined in this scope level.",
                sym.name
            ));
        }
        let slot = self.slot_count;
        self.locals.insert(
            sym.clone(),
            LocalInfo {
                slot,
                _ty: StaticType::Any,
            },
        );
        self.slot_count += 1;
        Ok(slot)
    }

    fn resolve(&self, sym: &Symbol) -> Option<LocalInfo> {
        self.locals.get(sym).cloned()
    }
}

struct FunctionCompiler {
    scope: CompilerScope,
    upvalues: Vec<Address>,
}

impl FunctionCompiler {
    fn new() -> Self {
        Self {
            scope: CompilerScope::new(),
            upvalues: Vec::new(),
        }
    }

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

pub struct Binder {
    functions: Vec<FunctionCompiler>,
    // Globals mapping: Symbol -> Index
    globals: Rc<RefCell<HashMap<Symbol, u32>>>,
    // Map of Declaration Identity -> List of Lambda Identities that capture it
    capture_map: HashMap<Identity, Vec<Identity>>,
}

impl Binder {
    pub fn new(globals: Rc<RefCell<HashMap<Symbol, u32>>>) -> Self {
        Self::with_boxed(globals, HashMap::new())
    }

    fn with_boxed(
        globals: Rc<RefCell<HashMap<Symbol, u32>>>,
        captures: HashMap<Identity, Vec<Identity>>,
    ) -> Self {
        let mut binder = Self {
            functions: Vec::new(),
            globals,
            capture_map: captures,
        };
        binder.functions.push(FunctionCompiler::new());
        binder
    }

    pub fn bind_root(
        globals: Rc<RefCell<HashMap<Symbol, u32>>>,
        node: &Node<UntypedKind>,
    ) -> Result<BoundNode, String> {
        let captures = crate::ast::compiler::upvalues::UpvalueAnalyzer::analyze(node);
        let mut binder = Self::with_boxed(globals, captures);
        binder.bind(node)
    }

    pub fn bind(&mut self, node: &Node<UntypedKind>) -> Result<BoundNode, String> {
        match &node.kind {
            UntypedKind::Nop => Ok(self.make_node(node.identity.clone(), BoundKind::Nop)),
            UntypedKind::Constant(v) => {
                Ok(self.make_node(node.identity.clone(), BoundKind::Constant(v.clone())))
            }

            UntypedKind::Identifier(sym) => {
                let addr = self.resolve_variable(sym)?;
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Get {
                        addr,
                        name: sym.clone(),
                    },
                ))
            }

            UntypedKind::Parameter(sym) => {
                // Parameters are ONLY allowed if we are inside a function (Binder stack > 1)
                if self.functions.len() <= 1 {
                    return Err(format!(
                        "Parameter '{}' is not allowed in root scope.",
                        sym.name
                    ));
                }
                let current_fn = self.functions.last_mut().unwrap();
                let slot = current_fn.scope.define(sym)?;
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Parameter {
                        name: sym.clone(),
                        slot,
                    },
                ))
            }

            UntypedKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = self.bind(cond)?;
                let then_br = self.bind(then_br)?;
                let mut else_br_bound = None;
                if let Some(e) = else_br {
                    else_br_bound = Some(Box::new(self.bind(e)?));
                }

                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::If {
                        cond: Box::new(cond),
                        then_br: Box::new(then_br),
                        else_br: else_br_bound,
                    },
                ))
            }

            UntypedKind::Def { name, value } => {
                // 1. Pre-declare name to support recursion
                let slot_or_idx = if self.functions.len() == 1 {
                    let mut globals = self.globals.borrow_mut();
                    if globals.contains_key(name) {
                        return Err(format!(
                            "Global variable '{}' is already defined.",
                            name.name
                        ));
                    }
                    let idx = globals.len() as u32;
                    globals.insert(name.clone(), idx);
                    idx
                } else {
                    let current_fn = self.functions.last_mut().unwrap();
                    current_fn.scope.define(name)?
                };

                // 2. Bind Value (now 'name' is visible)
                let val_node = self.bind(value)?;

                // 3. Return Node
                if self.functions.len() == 1 {
                    Ok(self.make_node(
                        node.identity.clone(),
                        BoundKind::DefGlobal {
                            name: name.clone(),
                            global_index: slot_or_idx,
                            value: Box::new(val_node),
                        },
                    ))
                } else {
                    let captured_by = self
                        .capture_map
                        .get(&node.identity)
                        .cloned()
                        .unwrap_or_default();
                    Ok(self.make_node(
                        node.identity.clone(),
                        BoundKind::DefLocal {
                            name: name.clone(),
                            slot: slot_or_idx,
                            value: Box::new(val_node),
                            captured_by,
                        },
                    ))
                }
            }

            UntypedKind::Assign { target, value } => {
                let val_node = self.bind(value)?;

                if let UntypedKind::Identifier(sym) = &target.kind {
                    let addr = self.resolve_variable(sym)?;
                    Ok(self.make_node(
                        node.identity.clone(),
                        BoundKind::Set {
                            addr,
                            value: Box::new(val_node),
                        },
                    ))
                } else {
                    Err("Assignment target must be an identifier".to_string())
                }
            }

            UntypedKind::Lambda { params, body } => {
                let identity = node.identity.clone();
                self.functions.push(FunctionCompiler::new());

                // 1. Bind the parameter pattern/tuple
                let params_bound = self.bind(params)?;

                // 2. Bind the body
                let body_bound = self.bind(body)?;

                let compiled_fn = self.functions.pop().unwrap();

                // 3. Static optimization: count total parameters needed in flat argument list
                fn count_params(node: &BoundNode) -> Option<u32> {
                    match &node.kind {
                        BoundKind::Parameter { .. } => Some(1),
                        BoundKind::Tuple { elements } => {
                            let mut total = 0;
                            for e in elements {
                                total += count_params(e)?;
                            }
                            Some(total)
                        }
                        BoundKind::Nop => Some(0),
                        _ => None,
                    }
                }

                let positional_count = count_params(&params_bound);

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

            UntypedKind::Call { callee, args } => {
                let callee = self.bind(callee)?;
                let args = self.bind(args)?;

                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Call {
                        callee: Box::new(callee),
                        args: Box::new(args),
                    },
                ))
            }

            UntypedKind::Block { exprs } => {
                let mut bound_exprs = Vec::new();
                for expr in exprs {
                    bound_exprs.push(self.bind(expr)?);
                }
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Block { exprs: bound_exprs },
                ))
            }

            UntypedKind::Tuple { elements } => {
                let mut bound_elems = Vec::new();
                for e in elements {
                    bound_elems.push(self.bind(e)?);
                }
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Tuple {
                        elements: bound_elems,
                    },
                ))
            }

            UntypedKind::Record { fields } => {
                let mut bound_fields = Vec::new();
                for (k, v) in fields {
                    bound_fields.push((self.bind(k)?, self.bind(v)?));
                }
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Record {
                        fields: bound_fields,
                    },
                ))
            }

            UntypedKind::Expansion { call, expanded } => {
                let bound_expanded = self.bind(expanded)?;
                Ok(self.make_node(
                    node.identity.clone(),
                    BoundKind::Expansion {
                        original_call: Rc::from(call.as_ref().clone()),
                        bound_expanded: Box::new(bound_expanded),
                    },
                ))
            }

            UntypedKind::Template(_)
            | UntypedKind::Placeholder(_)
            | UntypedKind::Splice(_)
            | UntypedKind::MacroDecl { .. } => Err(format!(
                "Macro construct {:?} found in Binder. Macros must be expanded before binding.",
                node.kind
            )),

            UntypedKind::Extension(_) => {
                // Future: Delegate to extension binder
                Err("Custom extensions not supported in Binder yet".to_string())
            }
        }
    }

    fn resolve_variable(&mut self, sym: &Symbol) -> Result<Address, String> {
        let current_fn_idx = self.functions.len() - 1;

        // 1. Try local in current function
        if let Some(info) = self.functions[current_fn_idx].scope.resolve(sym) {
            return Ok(Address::Local(info.slot));
        }

        // 2. Try enclosing scopes (capture chain)
        for i in (0..current_fn_idx).rev() {
            if let Some(info) = self.functions[i].scope.resolve(sym) {
                let mut addr = Address::Local(info.slot);

                for k in (i + 1)..=current_fn_idx {
                    addr = Address::Upvalue(self.functions[k].add_upvalue(addr));
                }
                return Ok(addr);
            }
        }

        // 3. Try Global
        let globals = self.globals.borrow();
        if let Some(idx) = globals.get(sym) {
            return Ok(Address::Global(*idx));
        }

        // 4. Global Fallback
        if sym.context.is_some() {
            let fallback_sym = Symbol {
                name: sym.name.clone(),
                context: None,
            };
            if let Some(idx) = globals.get(&fallback_sym) {
                return Ok(Address::Global(*idx));
            }
        }

        Err(format!("Undefined variable '{}'", sym.name))
    }

    fn make_node(&self, identity: Identity, kind: BoundKind<()>) -> BoundNode {
        Node {
            identity,
            kind,
            ty: (),
        }
    }
}

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

    #[test]
    fn test_upvalue_capture_sets_is_boxed() {
        // Wrap in a lambda to ensure 'x' is a local variable, not a global
        let source = "(fn [] (do (def x 10) (def f (fn [] x)) x))";
        let mut parser = Parser::new(source).unwrap();
        let untyped = parser.parse_expression().unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
        let bound = Binder::bind_root(globals, &untyped).unwrap();

        // Structure: Lambda -> Block -> [ DefLocal(x), DefLocal(f), Get(x) ]
        if let BoundKind::Lambda { body, .. } = &bound.kind {
            if let BoundKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
                if let BoundKind::DefLocal { captured_by, .. } = &x_decl.kind {
                    assert!(
                        !captured_by.is_empty(),
                        "Variable 'x' should have capturers because it is used in lambda 'f'"
                    );
                } else {
                    panic!(
                        "First expression in block should be DefLocal, got {:?}",
                        x_decl.kind
                    );
                }
            } else {
                panic!("Lambda body should be a Block, got {:?}", body.kind);
            }
        } else {
            panic!("Root should be a Lambda, got {:?}", bound.kind);
        }
    }

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

        let globals = Rc::new(RefCell::new(HashMap::new()));
        let bound = Binder::bind_root(globals, &untyped).unwrap();

        if let BoundKind::Lambda { body, .. } = &bound.kind {
            if let BoundKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
                if let BoundKind::DefLocal { captured_by, .. } = &x_decl.kind {
                    assert!(
                        captured_by.is_empty(),
                        "Variable 'x' should NOT have any capturers"
                    );
                } else {
                    panic!("First expression should be DefLocal");
                }
            } 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))";
        let mut parser = Parser::new(source).unwrap();
        let untyped = parser.parse_expression().unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
        let result = Binder::bind_root(globals, &untyped);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("already defined"));
    }

    #[test]
    fn test_repro_global_redefinition() {
        let globals = Rc::new(RefCell::new(HashMap::new()));

        // First run: defines 'x'
        let source1 = "(def x 1)";
        let untyped1 = Parser::new(source1).unwrap().parse_expression().unwrap();
        assert!(Binder::bind_root(globals.clone(), &untyped1).is_ok());

        // Second run: attempts to redefine 'x' in the same global environment
        let source2 = "(def x 2)";
        let untyped2 = Parser::new(source2).unwrap().parse_expression().unwrap();
        let result = Binder::bind_root(globals.clone(), &untyped2);

        assert!(result.is_err());
        assert!(result.unwrap_err().contains("already defined"));
    }
}