RustAst/src/ast/compiler/macros.rs

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

/// Trait for evaluating expressions during macro expansion.
pub trait MacroEvaluator {
    fn evaluate(
        &self,
        node: &Node<UntypedKind>,
        bindings: &HashMap<Rc<str>, Node<UntypedKind>>,
    ) -> Result<Value, String>;
}

/// Internal state for template expansion (Hygiene context + Parameter binding)
struct ExpansionState<'a> {
    bindings: &'a HashMap<Rc<str>, Node<UntypedKind>>,
    /// The identity of the current expansion instance, used to "color" internal symbols.
    expansion_id: Identity,
}

type MacroMap = HashMap<Rc<str>, (Vec<Rc<str>>, Node<UntypedKind>)>;

/// A registry for macro declarations.
#[derive(Clone)]
pub struct MacroRegistry {
    scopes: Vec<Rc<MacroMap>>,
}

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

impl MacroRegistry {
    pub fn new() -> Self {
        Self {
            scopes: vec![Rc::new(HashMap::new())],
        }
    }

    pub fn push(&mut self) {
        self.scopes.push(Rc::new(HashMap::new()));
    }

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

    pub fn define(&mut self, name: Rc<str>, params: Vec<Rc<str>>, body: Node<UntypedKind>) {
        if let Some(current_rc) = self.scopes.last_mut() {
            // Copy-on-Write: If the Rc is shared, clone the map before modifying.
            let current = Rc::make_mut(current_rc);
            current.insert(name, (params, body));
        }
    }

    pub fn lookup(&self, name: &str) -> Option<(Vec<Rc<str>>, Node<UntypedKind>)> {
        for scope in self.scopes.iter().rev() {
            if let Some(m) = scope.get(name) {
                return Some(m.clone());
            }
        }
        None
    }
}

pub struct MacroExpander<E: MacroEvaluator> {
    registry: MacroRegistry,
    evaluator: E,
}

impl<E: MacroEvaluator> MacroExpander<E> {
    pub fn new(registry: MacroRegistry, evaluator: E) -> Self {
        Self {
            registry,
            evaluator,
        }
    }

    pub fn into_registry(self) -> MacroRegistry {
        self.registry
    }

    pub fn expand(&mut self, node: Node<UntypedKind>) -> Result<Node<UntypedKind>, String> {
        self.expand_recursive(node)
    }

    fn expand_recursive(&mut self, node: Node<UntypedKind>) -> Result<Node<UntypedKind>, String> {
        match node.kind {
            UntypedKind::MacroDecl { name, params, body } => {
                let p_names = self.extract_param_names(&params)?;
                self.registry.define(name.name, p_names, *body);
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Nop,
                    ty: (),
                })
            }

            UntypedKind::Call { callee, args } => {
                if let UntypedKind::Identifier(ref sym) = callee.kind
                    && let Some((params, body)) = self.registry.lookup(&sym.name)
                {
                    let expanded_args = self.expand_recursive(*args)?;

                    // Extract argument nodes from the expanded tuple
                    let arg_elements = if let UntypedKind::Tuple { elements } = expanded_args.kind {
                        elements
                    } else {
                        vec![expanded_args]
                    };

                    let expanded = self.expand_call(
                        node.identity.clone(),
                        &sym.name,
                        params,
                        arg_elements,
                        body,
                    )?;
                    return self.expand_recursive(expanded);
                }

                let expanded_callee = self.expand_recursive(*callee)?;
                let expanded_args = self.expand_recursive(*args)?;

                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Call {
                        callee: Box::new(expanded_callee),
                        args: Box::new(expanded_args),
                    },
                    ty: (),
                })
            }

            UntypedKind::Block { exprs } => {
                self.registry.push();
                let mut expanded_exprs = Vec::new();
                for expr in exprs {
                    expanded_exprs.push(self.expand_recursive(expr)?);
                }
                self.registry.pop();

                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Block {
                        exprs: expanded_exprs,
                    },
                    ty: (),
                })
            }

            UntypedKind::Pipe { inputs, lambda } => {
                let mut expanded_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    expanded_inputs.push(self.expand_recursive(input)?);
                }
                let expanded_lambda = Box::new(self.expand_recursive(*lambda)?);
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Pipe {
                        inputs: expanded_inputs,
                        lambda: expanded_lambda,
                    },
                    ty: (),
                })
            }

            UntypedKind::Lambda { params, body } => {
                self.registry.push();
                let expanded_params = self.expand_recursive(*params)?;
                let expanded_body = self.expand_recursive(Node::clone(&body))?;
                self.registry.pop();

                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Lambda {
                        params: Box::new(expanded_params),
                        body: Rc::new(expanded_body),
                    },
                    ty: (),
                })
            }

            UntypedKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = self.expand_recursive(*cond)?;
                let then_br = self.expand_recursive(*then_br)?;
                let else_br = if let Some(e) = else_br {
                    Some(Box::new(self.expand_recursive(*e)?))
                } else {
                    None
                };
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::If {
                        cond: Box::new(cond),
                        then_br: Box::new(then_br),
                        else_br,
                    },
                    ty: (),
                })
            }

            UntypedKind::Def { target, value } => {
                let target = self.expand_recursive(*target)?;
                let value = self.expand_recursive(*value)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Def {
                        target: Box::new(target),
                        value: Box::new(value),
                    },
                    ty: (),
                })
            }

            UntypedKind::Assign { target, value } => {
                let target = self.expand_recursive(*target)?;
                let value = self.expand_recursive(*value)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Assign {
                        target: Box::new(target),
                        value: Box::new(value),
                    },
                    ty: (),
                })
            }

            UntypedKind::Tuple { elements } => {
                let mut expanded_elements = Vec::new();
                for e in elements {
                    expanded_elements.push(self.expand_recursive(e)?);
                }
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Tuple {
                        elements: expanded_elements,
                    },
                    ty: (),
                })
            }

            UntypedKind::Record { fields } => {
                let mut expanded_fields = Vec::new();
                for (k, v) in fields {
                    expanded_fields.push((self.expand_recursive(k)?, self.expand_recursive(v)?));
                }
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Record {
                        fields: expanded_fields,
                    },
                    ty: (),
                })
            }

            UntypedKind::Expansion { call, expanded } => {
                let expanded = self.expand_recursive(*expanded)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Expansion {
                        call,
                        expanded: Box::new(expanded),
                    },
                    ty: (),
                })
            }

            UntypedKind::Again { args } => {
                let expanded_args = self.expand_recursive(*args)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Again {
                        args: Box::new(expanded_args),
                    },
                    ty: (),
                })
            }

            _ => Ok(node),
        }
    }

    fn expand_call(
        &self,
        identity: Identity,
        name: &str,
        params: Vec<Rc<str>>,
        args: Vec<Node<UntypedKind>>,
        body: Node<UntypedKind>,
    ) -> Result<Node<UntypedKind>, String> {
        if params.len() != args.len() {
            return Err(format!(
                "Macro {} expects {} arguments, but got {}",
                name,
                params.len(),
                args.len()
            ));
        }

        let mut bindings = HashMap::new();
        for (p, a) in params.into_iter().zip(args.into_iter()) {
            bindings.insert(p, a);
        }

        // AST-Authority: Substitution and Hygiene ONLY happen if there is a Template node.
        let expanded_body = if let UntypedKind::Template(inner) = body.kind {
            let mut state = ExpansionState {
                bindings: &bindings,
                expansion_id: identity.clone(),
            };
            self.expand_template(*inner, &mut state)?
        } else {
            // Static AST fragment: No substitution, no hygiene.
            body
        };

        let original_call = Node {
            identity: identity.clone(),
            kind: UntypedKind::Call {
                callee: Box::new(Node {
                    identity: identity.clone(),
                    kind: UntypedKind::Identifier(Symbol::from(name)),
                    ty: (),
                }),
                args: Box::new(Node {
                    identity: identity.clone(),
                    kind: UntypedKind::Tuple {
                        elements: bindings.into_values().collect(),
                    },
                    ty: (),
                }),
            },
            ty: (),
        };

        Ok(Node {
            identity,
            kind: UntypedKind::Expansion {
                call: Box::new(original_call),
                expanded: Box::new(expanded_body),
            },
            ty: (),
        })
    }

    fn extract_param_names(&self, node: &Node<UntypedKind>) -> Result<Vec<Rc<str>>, String> {
        match &node.kind {
            UntypedKind::Identifier(sym) => Ok(vec![sym.name.clone()]),
            UntypedKind::Tuple { elements } => {
                let mut names = Vec::new();
                for el in elements {
                    names.extend(self.extract_param_names(el)?);
                }
                Ok(names)
            }
            _ => Err("Invalid parameter pattern in macro declaration".to_string()),
        }
    }

    fn expand_template(
        &self,
        node: Node<UntypedKind>,
        state: &mut ExpansionState,
    ) -> Result<Node<UntypedKind>, String> {
        match node.kind {
            UntypedKind::Identifier(mut sym) => {
                // Inside a template, all internal identifiers are colored.
                sym.context = Some(state.expansion_id.clone());
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Identifier(sym),
                    ty: (),
                })
            }

            UntypedKind::Placeholder(inner) => {
                // Break out of template for substitution/evaluation
                if let UntypedKind::Identifier(ref sym) = inner.kind
                    && let Some(arg) = state.bindings.get(&sym.name)
                {
                    return Ok(arg.clone());
                }
                let val = self.evaluator.evaluate(&inner, state.bindings)?;
                Ok(self.value_to_node(val, node.identity))
            }

            UntypedKind::Splice(inner) => {
                // Splicing is also a form of substitution
                if let UntypedKind::Identifier(ref sym) = inner.kind
                    && let Some(arg) = state.bindings.get(&sym.name)
                {
                    return Ok(arg.clone());
                }
                let val = self.evaluator.evaluate(&inner, state.bindings)?;
                Ok(self.value_to_node(val, node.identity))
            }

            UntypedKind::Def { target, value } => {
                let target = self.expand_template(*target, state)?;
                let val = self.expand_template(*value, state)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Def {
                        target: Box::new(target),
                        value: Box::new(val),
                    },
                    ty: (),
                })
            }

            UntypedKind::Assign { target, value } => {
                let target = self.expand_template(*target, state)?;
                let value = self.expand_template(*value, state)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Assign {
                        target: Box::new(target),
                        value: Box::new(value),
                    },
                    ty: (),
                })
            }

            UntypedKind::Tuple { elements } => {
                let mut new_elements = Vec::new();
                for e in elements {
                    if let UntypedKind::Splice(ref inner) = e.kind {
                        let val = self.evaluator.evaluate(inner, state.bindings)?;
                        self.handle_splice_value(val, node.identity.clone(), &mut new_elements)?;
                    } else {
                        new_elements.push(self.expand_template(e, state)?);
                    }
                }
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Tuple {
                        elements: new_elements,
                    },
                    ty: (),
                })
            }

            UntypedKind::Block { exprs } => {
                let mut new_exprs = Vec::new();
                for e in exprs {
                    if let UntypedKind::Splice(ref inner) = e.kind {
                        let val = self.evaluator.evaluate(inner, state.bindings)?;
                        self.handle_splice_value(val, node.identity.clone(), &mut new_exprs)?;
                    } else {
                        new_exprs.push(self.expand_template(e, state)?);
                    }
                }
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Block { exprs: new_exprs },
                    ty: (),
                })
            }

            UntypedKind::Record { fields } => {
                let mut new_fields = Vec::new();
                for (k, v) in fields {
                    new_fields.push((
                        self.expand_template(k, state)?,
                        self.expand_template(v, state)?,
                    ));
                }
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Record { fields: new_fields },
                    ty: (),
                })
            }

            UntypedKind::Call { callee, args } => {
                let expanded_callee = self.expand_template(*callee, state)?;
                let expanded_args = self.expand_template(*args, state)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Call {
                        callee: Box::new(expanded_callee),
                        args: Box::new(expanded_args),
                    },
                    ty: (),
                })
            }

            UntypedKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = self.expand_template(*cond, state)?;
                let then_br = self.expand_template(*then_br, state)?;
                let else_br = if let Some(e) = else_br {
                    Some(Box::new(self.expand_template(*e, state)?))
                } else {
                    None
                };
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::If {
                        cond: Box::new(cond),
                        then_br: Box::new(then_br),
                        else_br,
                    },
                    ty: (),
                })
            }

            UntypedKind::Pipe { inputs, lambda } => {
                let mut expanded_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    expanded_inputs.push(self.expand_template(input, state)?);
                }
                let expanded_lambda = Box::new(self.expand_template(*lambda, state)?);
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Pipe {
                        inputs: expanded_inputs,
                        lambda: expanded_lambda,
                    },
                    ty: (),
                })
            }

            UntypedKind::Lambda { params, body } => {
                let expanded_params = self.expand_template(*params, state)?;
                let body = self.expand_template(Node::clone(&body), state)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Lambda {
                        params: Box::new(expanded_params),
                        body: Rc::new(body),
                    },
                    ty: (),
                })
            }

            UntypedKind::Again { args } => {
                let expanded_args = self.expand_template(*args, state)?;
                Ok(Node {
                    identity: node.identity,
                    kind: UntypedKind::Again {
                        args: Box::new(expanded_args),
                    },
                    ty: (),
                })
            }

            _ => Ok(node),
        }
    }

    fn handle_splice_value(
        &self,
        val: Value,
        _identity: Identity,
        target: &mut Vec<Node<UntypedKind>>,
    ) -> Result<(), String> {
        match val {
            Value::Object(obj) => {
                if let Some(node) = obj.as_any().downcast_ref::<Node<UntypedKind>>() {
                    match &node.kind {
                        UntypedKind::Tuple { elements } => {
                            for e in elements {
                                target.push(e.clone());
                            }
                            return Ok(());
                        }
                        UntypedKind::Block { exprs } => {
                            for e in exprs {
                                target.push(e.clone());
                            }
                            return Ok(());
                        }
                        _ => {}
                    }
                }
                Err(format!(
                    "Strict Splicing: Expected Tuple or Block node, but found {}",
                    obj.type_name()
                ))
            }
            _ => Err(format!(
                "Strict Splicing: Expected sequence (Tuple/Block Node), but found Value: {}",
                val
            )),
        }
    }

    fn value_to_node(&self, val: Value, identity: Identity) -> Node<UntypedKind> {
        match val {
            Value::Object(obj) => {
                if let Some(node) = obj.as_any().downcast_ref::<Node<UntypedKind>>() {
                    return node.clone();
                }
                Node {
                    identity,
                    kind: UntypedKind::Constant(Value::Object(obj)),
                    ty: (),
                }
            }
            _ => Node {
                identity,
                kind: UntypedKind::Constant(val),
                ty: (),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::compiler::bound_nodes::Address;
    use crate::ast::compiler::Binder;
    use crate::ast::parser::Parser;
    use crate::ast::types::{Object, Value};

    struct SimpleEvaluator;
    impl MacroEvaluator for SimpleEvaluator {
        fn evaluate(
            &self,
            node: &Node<UntypedKind>,
            bindings: &HashMap<Rc<str>, Node<UntypedKind>>,
        ) -> Result<Value, String> {
            if let UntypedKind::Identifier(ref sym) = node.kind
                && let Some(arg) = bindings.get(&sym.name)
            {
                return Ok(Value::Object(Rc::new(arg.clone()) as Rc<dyn Object>));
            }
            Err(format!(
                "SimpleEvaluator cannot evaluate complex expression: {:?}",
                node.kind
            ))
        }
    }

    #[test]
    fn test_macro_expansion_hygiene() {
        let source = "
            (do
                (macro m [] `(def y 10))
                (m))
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        if let UntypedKind::Block { exprs } = &expanded.kind
            && let UntypedKind::Expansion {
                expanded: result,
                call,
            } = &exprs[1].kind
        {
            if let UntypedKind::Def { target, .. } = &result.kind {
                if let UntypedKind::Identifier(sym) = &target.kind {
                    assert_eq!(sym.context, Some(call.identity.clone()));
                    assert_eq!(sym.name.as_ref(), "y");
                } else {
                    panic!("Expected Identifier target, got {:?}", target.kind);
                }
            } else {
                panic!("Expected Def result, got {:?}", result.kind);
            }
        }
    }

    #[test]
    fn test_macro_expansion_unless() {
        let source = "
            (do
                (macro unless [c b] `(if ~c ... ~b))
                (unless false 42))
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        if let UntypedKind::Block { exprs } = &expanded.kind
            && let UntypedKind::Expansion {
                call: _,
                expanded: result,
            } = &exprs[1].kind
            && let UntypedKind::If {
                cond,
                then_br,
                else_br,
            } = &result.kind
        {
            if let UntypedKind::Identifier(sym) = &cond.kind {
                assert_eq!(sym.name.as_ref(), "false");
            } else {
                panic!("Expected identifier 'false', got {:?}", cond.kind);
            }
            assert!(matches!(then_br.kind, UntypedKind::Nop));
            if let Some(eb) = else_br {
                if let UntypedKind::Constant(Value::Int(n)) = &eb.kind {
                    assert_eq!(*n, 42);
                } else {
                    panic!("Expected 42, got {:?}", eb.kind);
                }
            } else {
                panic!("Else branch missing");
            }
        }
    }

    #[test]
    fn test_macro_expansion_splice() {
        let source = "
            (do
                (macro wrap [items] `[0 ~@items 4])
                (def t (wrap [1 2 3])))
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        if let UntypedKind::Block { exprs } = &expanded.kind
            && let UntypedKind::Def { value, .. } = &exprs[1].kind
            && let UntypedKind::Expansion {
                expanded: result, ..
            } = &value.kind
            && let UntypedKind::Tuple { elements } = &result.kind
        {
            assert_eq!(elements.len(), 5);
            if let UntypedKind::Constant(Value::Int(n)) = &elements[0].kind {
                assert_eq!(*n, 0);
            }
            if let UntypedKind::Constant(Value::Int(n)) = &elements[4].kind {
                assert_eq!(*n, 4);
            }
        }
    }

    #[test]
    fn test_repro_macro_global_lookup() {
        let source = "
            (do
                (macro square [x] `(* ~x ~x))
                (square 3))
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        // Convert test globals into a CompilerScope
        let mut locals = HashMap::new();
        locals.insert(
            Symbol::from("*"),
            crate::ast::compiler::binder::LocalInfo {
                addr: Address::Local(crate::ast::compiler::bound_nodes::LocalSlot(0)),
                identity: crate::ast::types::NodeIdentity::new(crate::ast::types::SourceLocation {
                    line: 0,
                    col: 0,
                }),
                _ty: crate::ast::types::StaticType::Any,
                purity: crate::ast::types::Purity::Pure,
            },
        );
        let initial_scopes = vec![crate::ast::compiler::binder::CompilerScope { locals }];

        let mut diag = crate::ast::diagnostics::Diagnostics::new();
        // fixed_scope_idx = 0 means scope 0 is frozen (Global)
        let result = Binder::bind_root(initial_scopes, 1, 0, &expanded, &mut diag);
        assert!(
            result.is_ok(),
            "Should find global '*' Error: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_repro_macro_hygiene_success() {
        let source = "
            (do
                (def y 1)
                (macro m [] `(def y 10))
                (m)
                y)
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        let mut diag = crate::ast::diagnostics::Diagnostics::new();
        let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag);
        assert!(result.is_ok(), "Hygiene failed: {:?}", result.err());
    }

    #[test]
    fn test_repro_macro_hygiene_clash_fixed() {
        let source = "
            (do
                (def y 1)
                (macro m1 [x] `(do (def y 10) (assign y 4)))
                (m1 8)
                y)
        ";
        let mut parser = Parser::new(source);
        let untyped = parser.parse_expression();

        let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator);
        let expanded = expander.expand(untyped).unwrap();

        let mut diag = crate::ast::diagnostics::Diagnostics::new();
        let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag);
        assert!(
            result.is_ok(),
            "Explicit Hygiene with Backticks failed: {:?}",
            result.err()
        );
    }
}