RustAst/src/ast/vm.rs

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use std::any::Any;
use crate::ast::compiler::bound_nodes::{Address, BoundKind};
use crate::ast::nodes::Node;
use crate::ast::types::{Value, Object, StaticType};

#[derive(Debug, Clone)]
pub struct Closure {
    pub function_node: Rc<Node<BoundKind, StaticType>>,
    pub upvalues: Vec<Rc<RefCell<Value>>>,
}

impl Object for Closure {
    fn type_name(&self) -> &'static str {
        "closure"
    }
    fn as_any(&self) -> &dyn Any {
        self
    }
}

#[derive(Debug)]
struct CallFrame {
    stack_base: usize,
    closure: Option<Rc<Closure>>,
}

/// Hook for observing VM execution (zero-cost when O::ACTIVE is false)
pub trait VMObserver {
    const ACTIVE: bool = false;
    fn before_eval(&mut self, _vm: &VM, _node: &Node<BoundKind, StaticType>) {}
    fn after_eval(&mut self, _vm: &VM, _node: &Node<BoundKind, StaticType>, _res: &Result<Value, String>) {}
}

/// Default observer that does nothing and should be optimized away
pub struct NoOpObserver;
impl VMObserver for NoOpObserver {}

/// Observer that logs the execution tree and scope state
pub struct TracingObserver {
    pub logs: Vec<String>,
    indent: usize,
}

impl TracingObserver {
    pub fn new() -> Self {
        Self { logs: Vec::new(), indent: 0 }
    }
    
    fn pad(&self) -> String {
        "|  ".repeat(self.indent)
    }
}

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

impl VMObserver for TracingObserver {
    const ACTIVE: bool = true;
    
    fn before_eval(&mut self, _vm: &VM, node: &Node<BoundKind, StaticType>) {
        let pad = self.pad();
        self.logs.push(format!("{}{} [{}]: {{", pad, node.kind.display_name(), node.ty));
        self.indent += 1;
    }

    fn after_eval(&mut self, vm: &VM, node: &Node<BoundKind, StaticType>, res: &Result<Value, String>) {
        self.indent = self.indent.saturating_sub(1);
        let pad = self.pad();
        
        // Show scope state on certain nodes (like Delphi ShowScope)
        match &node.kind {
            BoundKind::DefLocal { .. } | BoundKind::Set { .. } | BoundKind::DefGlobal { .. } => {
                let s_pad = format!("{}|  ", pad);
                self.logs.push(format!("{}--- Scope Status ---", s_pad));
                self.logs.push(format!("{}Stack (top 5): {:?}", s_pad, vm.stack.iter().rev().take(5).collect::<Vec<_>>()));
            },
            _ => {}
        }

        let res_str = match res {
            Ok(v) => format!("{}", v),
            Err(e) => format!("ERROR: {}", e),
        };
        self.logs.push(format!("{}}} -> {}", pad, res_str));
    }
}

pub struct VM {
    stack: Vec<Value>,
    globals: Rc<RefCell<Vec<Value>>>,
    frames: Vec<CallFrame>,
}

macro_rules! dispatch_eval {
    ($self:ident, $node:ident, $eval_method:ident $(, $observer:ident)?) => {
        match &$node.kind {
            BoundKind::Nop => Ok(Value::Void),
            BoundKind::Constant(v) => Ok(v.clone()),

            BoundKind::DefGlobal { global_index, value } => {
                let val = $self.$eval_method($($observer,)? value)?;
                let idx = *global_index as usize;
                let mut globals = $self.globals.borrow_mut();
                if idx >= globals.len() {
                    globals.resize(idx + 1, Value::Void);
                }
                globals[idx] = val.clone();
                Ok(val)
            },

            BoundKind::Get(addr) => $self.get_value(*addr),
            
            BoundKind::Set { addr, value } => {
                let val = $self.$eval_method($($observer,)? value)?;
                $self.set_value(*addr, val.clone())?;
                Ok(val)
            },

            BoundKind::DefLocal { slot, value, captured_by } => {
                let val = $self.$eval_method($($observer,)? value)?;
                let final_val = if !captured_by.is_empty() {
                    Value::Cell(Rc::new(RefCell::new(val)))
                } else {
                    val
                };
                
                let frame = $self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (*slot as usize);
                
                if abs_index == $self.stack.len() {
                    $self.stack.push(final_val.clone());
                } else if abs_index < $self.stack.len() {
                    $self.stack[abs_index] = final_val.clone();
                } else {
                    return Err(format!("Stack gap at local {}", slot));
                }
                Ok(final_val)
            },

            BoundKind::If { cond, then_br, else_br } => {
                let c = $self.$eval_method($($observer,)? cond)?;
                if c.is_truthy() {
                    $self.$eval_method($($observer,)? then_br)
                } else if let Some(e) = else_br {
                    $self.$eval_method($($observer,)? e)
                } else {
                    Ok(Value::Void)
                }
            },

            BoundKind::Block { exprs } => {
                let mut last = Value::Void;
                for e in exprs {
                    last = $self.$eval_method($($observer,)? e)?;
                }
                Ok(last)
            },

            BoundKind::Lambda { param_count: _, upvalues, body } => {
                let mut captured = Vec::with_capacity(upvalues.len());
                for addr in upvalues {
                    captured.push($self.capture_upvalue(*addr)?);
                }

                let closure = Closure {
                    function_node: body.clone(),
                    upvalues: captured,
                };

                Ok(Value::Object(Rc::new(closure)))
            },
            
            BoundKind::TailCall { callee, args } => {
                let func_val = $self.$eval_method($($observer,)? callee)?;
                let mut arg_vals = Vec::with_capacity(args.len());
                for arg in args {
                    arg_vals.push($self.$eval_method($($observer,)? arg)?);
                }
                
                match func_val {
                    Value::Object(obj) => Ok(Value::TailCallRequest(Box::new((obj, arg_vals)))),
                    Value::Function(f) => Ok(f(arg_vals)),
                    _ => Err(format!("Tail call target is not a function: {}", func_val)),
                }
            },

            BoundKind::Call { callee, args } => {
                let mut func_val = $self.$eval_method($($observer,)? callee)?;
                
                let mut arg_vals = Vec::with_capacity(args.len());
                for arg in args {
                    arg_vals.push($self.$eval_method($($observer,)? arg)?);
                }

                loop {
                    match func_val {
                        Value::Function(f) => return Ok(f(arg_vals)),
                        Value::Object(obj) => {
                            if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
                                let old_stack_top = $self.stack.len();
                                let closure_rc = Rc::new(closure.clone()); 
                                
                                $self.stack.extend(arg_vals); 
                                
                                $self.frames.push(CallFrame {
                                    stack_base: old_stack_top,
                                    closure: Some(closure_rc.clone()),
                                });

                                let result = $self.$eval_method($($observer,)? &closure.function_node);

                                $self.frames.pop();
                                $self.stack.truncate(old_stack_top);
                                
                                match result {
                                    Ok(Value::TailCallRequest(payload)) => {
                                        let (next_obj, next_args) = *payload;
                                        func_val = Value::Object(next_obj);
                                        arg_vals = next_args;
                                        continue;
                                    },
                                    Ok(val) => return Ok(val),
                                    Err(e) => return Err(e),
                                }
                            } else {
                                return Err(format!("Object is not a closure: {}", obj.type_name()));
                            }
                        },
                        _ => return Err(format!("Attempt to call non-function: {}", func_val)),
                    }
                }
            },

            BoundKind::Tuple { elements } => {
                let mut vals = Vec::with_capacity(elements.len());
                for e in elements {
                    vals.push($self.$eval_method($($observer,)? e)?);
                }
                Ok(Value::List(Rc::new(vals)))
            },

            BoundKind::Map { entries } => {
                let mut map = HashMap::new();
                for (k, v) in entries {
                    let key = $self.$eval_method($($observer,)? k)?;
                    let val = $self.$eval_method($($observer,)? v)?;
                    
                    if let Value::Keyword(kw) = key {
                        map.insert(kw, val);
                    } else {
                        return Err(format!("Map key must be keyword, got {}", key));
                    }
                }
                Ok(Value::Record(Rc::new(map)))
            }
        }
    };
}

impl VM {
    pub fn new(globals: Rc<RefCell<Vec<Value>>>) -> Self {
        Self {
            stack: Vec::new(),
            globals,
            frames: Vec::new(),
        }
    }

    pub fn run(&mut self, root: &Node<BoundKind, StaticType>) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();

        self.frames.push(CallFrame {
            stack_base: 0,
            closure: None,
        });
        
        let result = self.eval(root);
        
        self.frames.pop();
        result
    }

    pub fn run_with_observer<O: VMObserver>(&mut self, observer: &mut O, root: &Node<BoundKind, StaticType>) -> Result<Value, String> {
        self.stack.clear();
        self.frames.clear();

        self.frames.push(CallFrame {
            stack_base: 0,
            closure: None,
        });
        
        let result = self.eval_observed(observer, root);
        
        self.frames.pop();
        result
    }

    /// The pure, original, zero-cost hot path.
    #[inline(always)]
    fn eval(&mut self, node: &Node<BoundKind, StaticType>) -> Result<Value, String> {
        dispatch_eval!(self, node, eval)
    }

    /// The observed path for debugging.
    fn eval_observed<O: VMObserver>(&mut self, observer: &mut O, node: &Node<BoundKind, StaticType>) -> Result<Value, String> {
        observer.before_eval(self, node);
        let result = dispatch_eval!(self, node, eval_observed, observer);
        observer.after_eval(self, node, &result);
        result
    }


    fn capture_upvalue(&mut self, addr: Address) -> Result<Rc<RefCell<Value>>, String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                    // Check if already boxed
                    if let Value::Cell(cell) = &self.stack[abs_index] {
                        Ok(cell.clone())
                    } else {
                        // Box it
                        let val = self.stack[abs_index].clone();
                        let cell = Rc::new(RefCell::new(val));
                        self.stack[abs_index] = Value::Cell(cell.clone());
                        Ok(cell)
                    }
                } else {
                    Err(format!("Stack underflow capture local {}", idx))
                }
            },
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        Ok(closure.upvalues[idx].clone())
                    } else {
                        Err(format!("Upvalue access out of bounds capture {}", idx))
                    }
                } else {
                    Err("Current frame has no closure".to_string())
                }
            },
             Address::Global(_) => Err("Cannot capture global directly".to_string()),
        }
    }

    fn get_value(&self, addr: Address) -> Result<Value, String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                    match &self.stack[abs_index] {
                        Value::Cell(cell) => Ok(cell.borrow().clone()),
                        val => Ok(val.clone()),
                    }
                } else {
                    Err(format!("Stack underflow access local {}", idx))
                }
            },
            Address::Global(idx) => {
                let idx = idx as usize;
                let globals = self.globals.borrow();
                if idx < globals.len() {
                    Ok(globals[idx].clone())
                } else {
                    Err(format!("Global access out of bounds {}", idx))
                }
            },
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        Ok(closure.upvalues[idx].borrow().clone())
                    } else {
                        Err(format!("Upvalue access out of bounds {}", idx))
                    }
                } else {
                    Err("Current frame has no closure (cannot access upvalues)".to_string())
                }
            }
        }
    }

    fn set_value(&mut self, addr: Address, value: Value) -> Result<(), String> {
        match addr {
            Address::Local(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                let abs_index = frame.stack_base + (idx as usize);
                if abs_index < self.stack.len() {
                     if let Value::Cell(cell) = &self.stack[abs_index] {
                        *cell.borrow_mut() = value;
                     } else {
                        self.stack[abs_index] = value;
                     }
                } else if abs_index == self.stack.len() {
                    self.stack.push(value);
                } else {
                    return Err(format!("Stack gap write local {}", idx));
                }
                Ok(())
            },
            Address::Global(idx) => {
                let idx = idx as usize;
                let mut globals = self.globals.borrow_mut();
                if idx >= globals.len() {
                    globals.resize(idx + 1, Value::Void);
                }
                globals[idx] = value;
                Ok(())
            },
            Address::Upvalue(idx) => {
                let frame = self.frames.last().ok_or("No call frame")?;
                if let Some(closure) = &frame.closure {
                    let idx = idx as usize;
                    if idx < closure.upvalues.len() {
                        *closure.upvalues[idx].borrow_mut() = value;
                        Ok(())
                    } else {
                        Err(format!("Upvalue assignment out of bounds {}", idx))
                    }
                } else {
                    Err("Current frame has no closure".to_string())
                }
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::types::{SourceLocation, NodeIdentity};

    fn make_dummy_identity() -> Rc<NodeIdentity> {
        Rc::new(NodeIdentity {
            location: SourceLocation { line: 0, col: 0 },
        })
    }

    #[test]
    fn test_capture_boxing_modification() {
        // Goal:
        // var x = 10; (Local 0)
        // var f = lambda { x = 20; }; (Local 1)
        // f();
        // return x; -> Should be 20

        let id = make_dummy_identity();

        // 1. Define Lambda Body: x = 20 (Upvalue 0)
        let lambda_body = Node {
            identity: id.clone(),
            ty: StaticType::Void,
            kind: BoundKind::Set {
                addr: Address::Upvalue(0),
                value: Box::new(Node {
                    identity: id.clone(),
                    ty: StaticType::Int,
                    kind: BoundKind::Constant(Value::Int(20)),
                }),
            },
        };

        // 2. Define Main Script
        let root = Node {
            identity: id.clone(),
            ty: StaticType::Int,
            kind: BoundKind::Block {
                exprs: vec![
                    // x = 10 (Local 0) - simulate definition by assignment to stack gap
                    Node {
                        identity: id.clone(),
                        ty: StaticType::Int,
                        kind: BoundKind::Set {
                            addr: Address::Local(0),
                            value: Box::new(Node {
                                identity: id.clone(),
                                ty: StaticType::Int,
                                kind: BoundKind::Constant(Value::Int(10)),
                            }),
                        },
                    },
                    // f = lambda capturing x (Local 1)
                    Node {
                        identity: id.clone(),
                        ty: StaticType::Any,
                        kind: BoundKind::Set {
                            addr: Address::Local(1),
                            value: Box::new(Node {
                                identity: id.clone(),
                                ty: StaticType::Any,
                                kind: BoundKind::Lambda {
                                    param_count: 0,
                                    upvalues: vec![Address::Local(0)], // Capture x
                                    body: Rc::new(lambda_body),
                                },
                            }),
                        },
                    },
                    // f()
                    Node {
                        identity: id.clone(),
                        ty: StaticType::Void,
                        kind: BoundKind::Call {
                            callee: Box::new(Node {
                                identity: id.clone(),
                                ty: StaticType::Any,
                                kind: BoundKind::Get(Address::Local(1)),
                            }),
                            args: vec![],
                        },
                    },
                    // return x (Local 0)
                    Node {
                        identity: id.clone(),
                        ty: StaticType::Int,
                        kind: BoundKind::Get(Address::Local(0)),
                    },
                ],
            },
        };

        let globals = Rc::new(RefCell::new(Vec::new()));
        let mut vm = VM::new(globals);
        
        let result = vm.run(&root);
        
        match result {
            Ok(Value::Int(val)) => assert_eq!(val, 20, "Captured variable should be modified to 20"),
            Ok(val) => panic!("Expected Int(20), got {:?}", val),
            Err(e) => panic!("VM Error: {}", e),
        }
    }
}