not-python-rust/src/obj/function.rs

use std::fmt::{self, Debug, Display};
use std::rc::Rc;

use gc::{Finalize, Trace};

use crate::obj::macros::*;
use crate::obj::prelude::*;
use crate::vm::{Argc, Chunk, Frame, Function, Vm};

////////////////////////////////////////////////////////////////////////////////
// FunctionResult, FunctionState
////////////////////////////////////////////////////////////////////////////////

/// A result that instructs the VM what to do after a function finishes its execution.
#[derive(Debug)]
pub enum FunctionResult {
    /// Take this value and push it on the stack as its return value.
    ReturnPush(ObjP),

    /// Return value has already been pushed to the stack.
    Return,

    /// Yield control to the VM with a state marker, and resume execution.
    ///
    /// This means that a new stack frame should have been pushed to the VM and resume execution
    /// starting from there.
    Yield(usize),
}

impl From<ObjP> for FunctionResult {
    fn from(other: ObjP) -> Self {
        FunctionResult::ReturnPush(other)
    }
}

/// A function's resume state.
///
/// When a builtin function is called, it may need to set up a stack frame for a new function,
/// yielding its control back to the VM. If it does this, then presumably, that function would like
/// to resume its execution where it left off. There's not really a way to capture a native
/// function's execution state in Rust (without tons of unsafe and not-really-worth-it black
/// magic), so instead a builtin function will accept a `FunctionState` value to give it a hint of
/// where it left off.
#[derive(Debug, Clone, Copy)]
pub enum FunctionState {
    Begin,
    Resume(usize),
}

////////////////////////////////////////////////////////////////////////////////
// BuiltinFunction
////////////////////////////////////////////////////////////////////////////////

pub type BuiltinFunctionPtr = fn(vm: &mut Vm, function_state: FunctionState) -> FunctionResult;

#[derive(Trace, Finalize)]
pub struct BuiltinFunction {
    base: Obj,
    #[unsafe_ignore_trace]
    name: Rc<String>,
    #[unsafe_ignore_trace]
    function: BuiltinFunctionPtr,
    arity: Argc,
}

impl BuiltinFunction {
    pub fn new(name: impl ToString, function: BuiltinFunctionPtr, arity: Argc) -> Self {
        Self {
            base: Default::default(),
            name: Rc::new(name.to_string()),
            function,
            arity,
        }
    }

    impl_create!(
        name: impl ToString,
        function: BuiltinFunctionPtr,
        arity: Argc,
    );

    pub fn name(&self) -> &Rc<String> {
        &self.name
    }
}

impl Display for BuiltinFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Debug::fmt(self, fmt)
    }
}

impl Debug for BuiltinFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "<BuiltinFunction {}/{} at 0x{:x}>",
            self.name(),
            self.arity().unwrap(),
            self.function as *const BuiltinFunctionPtr as usize
        )
    }
}

impl Object for BuiltinFunction {
    fn arity(&self) -> Option<Argc> {
        Some(self.arity)
    }

    fn call(&self, vm: &mut Vm, argc: Argc) {
        let new_frame = Frame::new(
            Rc::clone(&self.name),
            Function::Builtin(self.function, FunctionState::Begin),
            vm.stack().len() - (argc as usize),
        );
        vm.push_frame(new_frame);
    }

    impl_base_obj!(BuiltinFunction);
}

////////////////////////////////////////////////////////////////////////////////
// UserFunction
////////////////////////////////////////////////////////////////////////////////

#[derive(Clone, Trace, Finalize)]
pub struct UserFunction {
    base: Obj,
    #[unsafe_ignore_trace]
    path: Rc<String>,
    #[unsafe_ignore_trace]
    name: Rc<String>,
    #[unsafe_ignore_trace]
    chunk: Rc<Chunk>,
    arity: Argc,
    captures: Vec<ObjP>,
}

impl UserFunction {
    pub fn new(path: impl ToString, chunk: Chunk, arity: Argc) -> Self {
        Self {
            base: Default::default(),
            path: Rc::new(path.to_string()),
            name: Rc::new("(anonymous)".to_string()),
            chunk: Rc::new(chunk),
            arity,
            captures: Default::default(),
        }
    }

    impl_create!(path: impl ToString, chunk: Chunk, arity: Argc);

    pub fn path(&self) -> &Rc<String> {
        &self.path
    }

    pub fn name(&self) -> &Rc<String> {
        &self.name
    }

    pub fn set_name(&mut self, name: Rc<String>) {
        self.name = name;
    }

    pub fn chunk(&self) -> &Chunk {
        &self.chunk
    }

    pub fn push_capture(&mut self, value: ObjP) {
        self.captures.push(value);
    }
}

impl Display for UserFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Debug::fmt(self, fmt)
    }
}

impl Debug for UserFunction {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "<UserFunction {}/{} at 0x{:x}>",
            self.name(),
            self.arity().unwrap(),
            self as *const _ as usize
        )
    }
}

impl Object for UserFunction {
    fn arity(&self) -> Option<Argc> {
        Some(self.arity)
    }

    fn call(&self, vm: &mut Vm, argc: Argc) {
        assert_eq!(argc, self.arity, "argc must match arity");
        let new_frame = Frame::new(
            Rc::clone(&self.name),
            Function::Chunk(Rc::clone(&self.chunk)),
            vm.stack().len() - (argc as usize),
        );
        vm.push_frame(new_frame);
        for capture in &self.captures {
            vm.push(capture.clone());
        }
    }

    impl_base_obj!(UserFunction);
}

////////////////////////////////////////////////////////////////////////////////
// Method
////////////////////////////////////////////////////////////////////////////////

#[derive(Trace, Finalize)]
pub struct Method {
    base: Obj,
    self_binding: ObjP,
    function: ObjP,
}

impl Method {
    pub fn new(self_binding: ObjP, function: ObjP) -> Self {
        Self {
            base: Default::default(),
            self_binding,
            function,
        }
    }

    pub fn create(self_binding: ObjP, function: ObjP) -> ObjP {
        let ptr = make_ptr(Self::new(self_binding, function));
        ptr.borrow_mut().instantiate();
        ptr
    }

    pub fn self_binding(&self) -> &ObjP {
        &self.self_binding
    }
}

impl Display for Method {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        Debug::fmt(self, fmt)
    }
}

impl Debug for Method {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let function_name: Rc<_> = if let Some(function) = self
            .function
            .borrow()
            .as_any()
            .downcast_ref::<BuiltinFunction>()
        {
            Rc::clone(&function.name())
        } else if let Some(function) = self
            .function
            .borrow()
            .as_any()
            .downcast_ref::<UserFunction>()
        {
            function.name().clone()
        } else {
            unreachable!()
        };
        write!(
            fmt,
            "<Method {}.{}/{} at {}>",
            self.self_binding().borrow().ty_name(),
            function_name,
            self.function.borrow().arity().unwrap(),
            self as *const _ as usize
        )
    }
}

impl Object for Method {
    fn arity(&self) -> Option<Argc> {
        // Subtract one from the arity - this is because the VM uses arity() to check against the
        // number of arguments passed.
        self.function.borrow().arity().map(|arity| arity - 1)
    }

    fn call(&self, vm: &mut Vm, mut argc: Argc) {
        let self_pos = vm.stack().len() - (argc as usize);
        vm.stack_mut().insert(self_pos, self.self_binding().clone());
        argc += 1;

        self.function.borrow().call(vm, argc)
    }

    impl_base_obj!(Method);
}