endorphin-aster-vm/crates/core/src/eval.rs

use crate::frame::StackEntry;
use crate::{
    CompareOp, EndorCellObject, EndorCodeObject, EndorFrame, EndorFunctionObject,
    EndorStringObject, ExecutionLimits, ExecutionLocation, Instruction, ObjectRef, Opcode, Value,
    Vm, VmError,
};

fn invalid(offset: usize, reason: &'static str) -> VmError {
    VmError::InvalidBytecode { offset, reason }
}

impl Vm {
    /// Execute a code object in a globals dictionary. Like CPython's EvalCode
    /// entry path, this constructs a temporary function and enters its frame.
    /// Object results own one external reference: release_value() when done.
    pub fn execute(&mut self, code: ObjectRef, globals: ObjectRef) -> Result<Value, VmError> {
        self.execute_with_limits(code, globals, ExecutionLimits::default())
    }

    pub fn execute_with_limits(
        &mut self,
        code: ObjectRef,
        globals: ObjectRef,
        limits: ExecutionLimits,
    ) -> Result<Value, VmError> {
        self.last_traceback.clear();
        let name = self
            .payload::<EndorCodeObject>(code, "code")?
            .code()
            .name()
            .to_owned();
        let function = self.new_function(name, code, globals, Vec::new(), Vec::new())?;
        let result = self.call_with_limits(function, &[], limits);
        self.release(function)?;
        result
    }

    /// Dispatch both language and host functions. Arguments are borrowed;
    /// every Object result owns one external reference, including native
    /// callbacks that return one of their borrowed arguments.
    pub fn call(&mut self, function: ObjectRef, arguments: &[Value]) -> Result<Value, VmError> {
        self.call_with_limits(function, arguments, ExecutionLimits::default())
    }

    pub fn call_with_limits(
        &mut self,
        function: ObjectRef,
        arguments: &[Value],
        limits: ExecutionLimits,
    ) -> Result<Value, VmError> {
        self.last_traceback.clear();
        if !self
            .object(function)?
            .payload()
            .as_any()
            .is::<EndorFunctionObject>()
        {
            return self.invoke_native(function, arguments);
        }
        if limits.max_frames == 0 {
            return Err(VmError::ExecutionLimit("frames"));
        }
        let frame = self.prepare_frame(function, arguments)?;
        let mut frames = vec![frame];
        let result = self.eval_frames(&mut frames, limits);
        if result.is_err() {
            self.last_traceback = frames
                .iter()
                .map(|frame| ExecutionLocation {
                    function: frame.name.clone(),
                    filename: frame.executable.filename().into(),
                    first_line: frame.executable.first_line(),
                    instruction_offset: frame.current_offset,
                })
                .collect();
        }
        while let Some(frame) = frames.pop() {
            self.clear_frame(frame);
        }
        result
    }

    /// Outer-to-inner execution locations from the last failed call/execute.
    /// These are copied metadata, not heap references or language exceptions.
    pub fn last_traceback(&self) -> &[ExecutionLocation] {
        &self.last_traceback
    }

    pub fn release_value(&mut self, value: Value) -> Result<(), VmError> {
        if let Value::Object(reference) = value {
            self.release(reference)?;
        }
        Ok(())
    }

    fn retain_value(&mut self, value: &Value) -> Result<(), VmError> {
        self.validate_value(value)?;
        if let Value::Object(reference) = value {
            self.retain(*reference)?;
        }
        Ok(())
    }

    fn drop_root(&mut self, value: Value) {
        self.release_value(value)
            .expect("execution owns every object root it releases");
    }

    fn clear_frame(&mut self, frame: EndorFrame) {
        for entry in frame.stack {
            if let StackEntry::Value(value) = entry {
                self.drop_root(value);
            }
        }
        for value in frame.localsplus.into_iter().flatten() {
            self.drop_root(value);
        }
        self.release(frame.function)
            .expect("frame owns its function root");
    }

    fn prepare_frame(
        &mut self,
        function: ObjectRef,
        arguments: &[Value],
    ) -> Result<EndorFrame, VmError> {
        let data = self
            .payload::<EndorFunctionObject>(function, "function")?
            .clone();
        let executable = self
            .payload::<EndorCodeObject>(data.code(), "code")?
            .shared_code();
        let maximum = executable.parameters().len();
        let required = maximum - data.defaults().len();
        if arguments.len() < required || arguments.len() > maximum {
            return Err(VmError::ArgumentCount {
                name: data.name().into(),
                required,
                maximum,
                provided: arguments.len(),
            });
        }
        for value in arguments {
            self.validate_value(value)?;
        }
        self.retain(function)?;
        let mut frame = EndorFrame {
            function,
            previous: None,
            code: data.code(),
            globals: data.globals(),
            builtins: data.builtins(),
            name: data.name().into(),
            instruction_offset: 0,
            current_offset: 0,
            localsplus: vec![None; executable.slot_count()],
            stack: Vec::with_capacity(executable.stack_size()),
            executable,
        };
        let initialized = (|| -> Result<(), VmError> {
            let defaults = data.defaults().iter().skip(arguments.len() - required);
            for (slot, value) in frame
                .localsplus
                .iter_mut()
                .zip(arguments.iter().chain(defaults))
            {
                self.retain_value(value)?;
                *slot = Some(value.clone());
            }
            // Captured local cells occupy the original local slots, as in
            // CPython's localsplus layout. Free cells follow the local slots.
            let local_count = frame.executable.layout().locals.len();
            for (index, slot) in frame.localsplus[..local_count].iter_mut().enumerate() {
                if frame.executable.is_cell_slot(index) {
                    let cell = self.new_cell(slot.clone())?;
                    if let Some(old) = slot.replace(Value::Object(cell)) {
                        self.drop_root(old);
                    }
                }
            }
            let first_free = frame.executable.layout().locals.len();
            for (index, cell) in data.closure().iter().enumerate() {
                self.retain(*cell)?;
                frame.localsplus[first_free + index] = Some(Value::Object(*cell));
            }
            Ok(())
        })();
        if let Err(error) = initialized {
            self.clear_frame(frame);
            return Err(error);
        }
        Ok(frame)
    }

    fn invoke_native(
        &mut self,
        function: ObjectRef,
        arguments: &[Value],
    ) -> Result<Value, VmError> {
        self.object(function)?;
        for argument in arguments {
            self.validate_value(argument)?;
        }
        let result = self.object(function)?.payload().call(arguments)?;
        self.retain_value(&result)?;
        Ok(result)
    }

    fn stack_value(frame: &EndorFrame, index: usize) -> Result<Value, VmError> {
        match frame.stack.get(index) {
            Some(StackEntry::Value(value)) => Ok(value.clone()),
            _ => Err(invalid(
                frame.current_offset,
                "NULL call marker used as a value",
            )),
        }
    }

    fn top_value(frame: &EndorFrame) -> Result<Value, VmError> {
        Self::stack_value(frame, frame.stack.len() - 1)
    }

    fn push_borrowed(&mut self, frame: &mut EndorFrame, value: Value) -> Result<(), VmError> {
        self.retain_value(&value)?;
        frame.stack.push(StackEntry::Value(value));
        Ok(())
    }

    fn discard(&mut self, frame: &mut EndorFrame, count: usize) {
        for _ in 0..count {
            if let StackEntry::Value(value) = frame.stack.pop().expect("verified stack height") {
                self.drop_root(value);
            }
        }
    }

    fn cell_slot(&self, frame: &EndorFrame, index: usize) -> Result<ObjectRef, VmError> {
        match &frame.localsplus[index] {
            Some(Value::Object(reference)) => {
                self.payload::<EndorCellObject>(*reference, "cell")?;
                Ok(*reference)
            }
            _ => Err(invalid(
                frame.current_offset,
                "cell slot is not initialized",
            )),
        }
    }

    fn value_is_true(&self, value: &Value) -> Result<bool, VmError> {
        match value {
            Value::Nil => Ok(false),
            Value::Bool(value) => Ok(*value),
            Value::Text(value) => Ok(!value.is_empty()),
            Value::Number(value) => Ok(!value.is_zero()),
            Value::Object(reference) => self.is_true(*reference),
        }
    }

    fn text_value(&self, value: &Value) -> Result<Option<String>, VmError> {
        Ok(match value {
            Value::Text(text) => Some(text.clone()),
            Value::Object(reference) => self
                .object(*reference)?
                .payload()
                .as_any()
                .downcast_ref::<EndorStringObject>()
                .map(|text| text.as_str().to_owned()),
            _ => None,
        })
    }

    fn compare_values(
        &self,
        left: &Value,
        right: &Value,
        operation: CompareOp,
    ) -> Result<bool, VmError> {
        if matches!(left, Value::Number(_)) || matches!(right, Value::Number(_)) {
            return Err(VmError::UnsupportedOperation("numeric comparison"));
        }
        if let (Some(left), Some(right)) = (self.text_value(left)?, self.text_value(right)?) {
            return Ok(match operation {
                CompareOp::Eq => left == right,
                CompareOp::Ne => left != right,
                CompareOp::Lt => left < right,
                CompareOp::Le => left <= right,
                CompareOp::Gt => left > right,
                CompareOp::Ge => left >= right,
            });
        }
        if let (Value::Object(left), Value::Object(right)) = (left, right) {
            return self.compare(*left, *right, operation);
        }
        match operation {
            CompareOp::Eq => Ok(left == right),
            CompareOp::Ne => Ok(left != right),
            _ => Err(VmError::UnsupportedOperation("ordering")),
        }
    }

    fn build_map(&mut self, frame: &EndorFrame, count: usize) -> Result<ObjectRef, VmError> {
        let dictionary = self.new_dict()?;
        let result = (|| -> Result<(), VmError> {
            let start = frame.stack.len() - count * 2;
            for index in 0..count {
                let key = Self::stack_value(frame, start + index * 2)?;
                let value = Self::stack_value(frame, start + index * 2 + 1)?;
                self.dict_set(dictionary, key, value)?;
            }
            Ok(())
        })();
        if let Err(error) = result {
            self.release(dictionary)?;
            return Err(error);
        }
        Ok(dictionary)
    }

    fn make_function(&mut self, frame: &EndorFrame, argument: u32) -> Result<ObjectRef, VmError> {
        let (default_count, closure_count) = Instruction::function_counts(argument);
        let Value::Object(code) = Self::top_value(frame)? else {
            return Err(VmError::ObjectKindMismatch("code"));
        };
        let name = self
            .payload::<EndorCodeObject>(code, "code")?
            .code()
            .name()
            .to_owned();
        let start = frame.stack.len() - default_count - closure_count - 1;
        let defaults = (start..start + default_count)
            .map(|index| Self::stack_value(frame, index))
            .collect::<Result<Vec<_>, _>>()?;
        let closure = (start + default_count..frame.stack.len() - 1)
            .map(|index| match Self::stack_value(frame, index)? {
                Value::Object(cell) => Ok(cell),
                _ => Err(VmError::ObjectKindMismatch("cell")),
            })
            .collect::<Result<Vec<_>, _>>()?;
        // Keep every input root on the stack until allocation has validated
        // and retained all function edges. Failure leaves cleanup to the frame.
        self.new_function(name, code, frame.globals, defaults, closure)
    }

    fn eval_frames(
        &mut self,
        frames: &mut Vec<EndorFrame>,
        limits: ExecutionLimits,
    ) -> Result<Value, VmError> {
        let mut executed = 0;
        loop {
            if executed >= limits.max_instructions {
                return Err(VmError::ExecutionLimit("instructions"));
            }
            executed += 1;
            let frame = frames.last_mut().expect("an evaluation frame is active");
            let decoded = frame.executable.instruction_at(frame.instruction_offset)?;
            frame.current_offset = decoded.offset;
            frame.instruction_offset = decoded.next_offset;
            let opcode = decoded.instruction.opcode;
            let argument = decoded.instruction.argument as usize;
            match opcode {
                Opcode::Nop => {}
                Opcode::LoadConst => {
                    let value = frame.executable.constants()[argument].clone();
                    self.push_borrowed(frame, value)?;
                }
                Opcode::LoadFast => {
                    let value = frame.localsplus[argument].clone().ok_or_else(|| {
                        VmError::UnboundLocal(frame.executable.slot_name(argument).into())
                    })?;
                    self.push_borrowed(frame, value)?;
                }
                Opcode::StoreFast => {
                    let value = Self::top_value(frame)?;
                    // The stack root is moved into the local slot.
                    frame.stack.pop();
                    if let Some(old) = frame.localsplus[argument].replace(value) {
                        self.drop_root(old);
                    }
                }
                Opcode::DeleteFast => {
                    let value = frame.localsplus[argument].take().ok_or_else(|| {
                        VmError::UnboundLocal(frame.executable.slot_name(argument).into())
                    })?;
                    self.drop_root(value);
                }
                Opcode::LoadGlobal => {
                    let name = &frame.executable.layout().names[argument];
                    let key = Value::Text(name.clone());
                    let mut value = self.dict_get(frame.globals, &key)?;
                    if value.is_none()
                        && let Some(builtins) = frame.builtins
                    {
                        value = self.dict_get(builtins, &key)?;
                    }
                    let value = value.ok_or_else(|| VmError::NameNotFound(name.clone()))?;
                    self.push_borrowed(frame, value)?;
                }
                Opcode::StoreGlobal => {
                    let value = Self::top_value(frame)?;
                    let key = Value::Text(frame.executable.layout().names[argument].clone());
                    self.dict_set(frame.globals, key, value)?;
                    self.discard(frame, 1);
                }
                Opcode::DeleteGlobal => {
                    let name = &frame.executable.layout().names[argument];
                    if !self.dict_delete(frame.globals, &Value::Text(name.clone()))? {
                        return Err(VmError::NameNotFound(name.clone()));
                    }
                }
                Opcode::LoadDeref | Opcode::LoadClosure => {
                    let cell = self.cell_slot(frame, argument)?;
                    let value = if opcode == Opcode::LoadClosure {
                        Value::Object(cell)
                    } else {
                        self.cell_get(cell).map_err(|error| {
                            if error == VmError::EmptyCell {
                                VmError::UnboundLocal(frame.executable.slot_name(argument).into())
                            } else {
                                error
                            }
                        })?
                    };
                    self.push_borrowed(frame, value)?;
                }
                Opcode::StoreDeref => {
                    let cell = self.cell_slot(frame, argument)?;
                    self.cell_set(cell, Self::top_value(frame)?)?;
                    self.discard(frame, 1);
                }
                Opcode::DeleteDeref => {
                    let cell = self.cell_slot(frame, argument)?;
                    if self.cell_get(cell).is_err() {
                        return Err(VmError::UnboundLocal(
                            frame.executable.slot_name(argument).into(),
                        ));
                    }
                    self.cell_clear(cell)?;
                }
                Opcode::PopTop => self.discard(frame, 1),
                Opcode::Copy => {
                    let entry = frame.stack[frame.stack.len() - argument].clone();
                    if let StackEntry::Value(value) = &entry {
                        self.retain_value(value)?;
                    }
                    frame.stack.push(entry);
                }
                Opcode::Swap => {
                    let length = frame.stack.len();
                    frame.stack.swap(length - 1, length - argument);
                }
                Opcode::PushNull => frame.stack.push(StackEntry::Null),
                Opcode::Call => {
                    let start = frame.stack.len() - argument - 2;
                    let Value::Object(function) = Self::stack_value(frame, start)? else {
                        return Err(VmError::UnsupportedOperation("call"));
                    };
                    let mut arguments = Vec::with_capacity(argument + 1);
                    if let StackEntry::Value(value) = &frame.stack[start + 1] {
                        arguments.push(value.clone());
                    }
                    for index in start + 2..frame.stack.len() {
                        arguments.push(Self::stack_value(frame, index)?);
                    }
                    let language_function = self
                        .object(function)?
                        .payload()
                        .as_any()
                        .is::<EndorFunctionObject>();
                    if language_function {
                        if frames.len() >= limits.max_frames {
                            return Err(VmError::ExecutionLimit("frames"));
                        }
                        let mut callee = self.prepare_frame(function, &arguments)?;
                        callee.previous = Some(frames.len() - 1);
                        self.discard(frames.last_mut().unwrap(), argument + 2);
                        frames.push(callee);
                    } else {
                        let value = self.invoke_native(function, &arguments)?;
                        self.discard(frames.last_mut().unwrap(), argument + 2);
                        frames
                            .last_mut()
                            .unwrap()
                            .stack
                            .push(StackEntry::Value(value));
                    }
                }
                Opcode::ReturnValue => {
                    let result = Self::top_value(frame)?;
                    frame.stack.pop();
                    self.clear_frame(frames.pop().unwrap());
                    if let Some(caller) = frames.last_mut() {
                        caller.stack.push(StackEntry::Value(result));
                    } else {
                        return Ok(result);
                    }
                }
                Opcode::JumpForward | Opcode::JumpBackward => {
                    frame.instruction_offset = decoded.jump_target()?
                }
                Opcode::PopJumpIfFalse | Opcode::PopJumpIfTrue => {
                    let truth = self.value_is_true(&Self::top_value(frame)?)?;
                    self.discard(frame, 1);
                    if truth == (opcode == Opcode::PopJumpIfTrue) {
                        frame.instruction_offset = decoded.jump_target()?;
                    }
                }
                Opcode::ToBool | Opcode::UnaryNot => {
                    let truth = self.value_is_true(&Self::top_value(frame)?)?;
                    self.discard(frame, 1);
                    frame.stack.push(StackEntry::Value(Value::Bool(
                        if opcode == Opcode::UnaryNot {
                            !truth
                        } else {
                            truth
                        },
                    )));
                }
                Opcode::CompareOp => {
                    let left = Self::stack_value(frame, frame.stack.len() - 2)?;
                    let right = Self::top_value(frame)?;
                    let result = self.compare_values(
                        &left,
                        &right,
                        Instruction::compare_operation(decoded.instruction.argument).unwrap(),
                    )?;
                    self.discard(frame, 2);
                    frame.stack.push(StackEntry::Value(Value::Bool(result)));
                }
                Opcode::LoadAttr => {
                    let Value::Object(owner) = Self::top_value(frame)? else {
                        return Err(VmError::ObjectKindMismatch("object"));
                    };
                    let value = self.get_attr(owner, &frame.executable.layout().names[argument])?;
                    self.retain_value(&value)?;
                    self.discard(frame, 1);
                    frame.stack.push(StackEntry::Value(value));
                }
                Opcode::StoreAttr => {
                    let Value::Object(owner) = Self::top_value(frame)? else {
                        return Err(VmError::ObjectKindMismatch("object"));
                    };
                    let value = Self::stack_value(frame, frame.stack.len() - 2)?;
                    self.set_attr(
                        owner,
                        frame.executable.layout().names[argument].clone(),
                        value,
                    )?;
                    self.discard(frame, 2);
                }
                Opcode::BuildMap => {
                    let dictionary = self.build_map(frame, argument)?;
                    self.discard(frame, argument * 2);
                    frame
                        .stack
                        .push(StackEntry::Value(Value::Object(dictionary)));
                }
                Opcode::MakeFunction => {
                    let function = self.make_function(frame, decoded.instruction.argument)?;
                    let (defaults, closure) =
                        Instruction::function_counts(decoded.instruction.argument);
                    self.discard(frame, defaults + closure + 1);
                    frame.stack.push(StackEntry::Value(Value::Object(function)));
                }
                Opcode::ExtendedArg => unreachable!("EXTENDED_ARG was decoded before evaluation"),
            }
        }
    }
}