rlox/interpreter/src/class.rs

use std::fmt::Display;
use std::rc::Rc;

use rlox2_frontend::parser::{Expr, Stmt};
use smol_str::SmolStr;

use crate::value::Attrs;
use crate::{LoxFunction, LoxReference, Value};

#[derive(Debug, Clone)]
pub struct LoxClass {
    superclass: Option<Rc<LoxClass>>,

    name: SmolStr,

    methods: Attrs,
}

/// Representation of a class in Lox. Always behind an Rc to ensure uniqueness. Should never be
/// handled raw.
impl LoxClass {
    pub fn new(
        name: impl Into<SmolStr>,
        methods: Attrs,
        superclass: Option<Rc<LoxClass>>,
    ) -> Rc<Self> {
        let name = name.into();
        let mut methods = methods;

        // if class has an init method: insert an implicit "return this;" at its end
        if let Some(Value::Function(init)) = methods.get("init") {
            let name = init.name().clone();
            let closure = init.closure().clone();
            let param_names = init.param_names().to_vec();

            let mut body = init.body().as_ref().clone();

            if let Stmt::Block { ref mut statements } = body {
                statements.push(Stmt::return_stmt(Expr::local_variable("this", 1)));
            } else {
                panic!("Body of init method of class {name} wasn't a block");
            }

            let body = Rc::new(body);

            let new_init = Value::function(LoxFunction::new(name, closure, param_names, body));

            methods.insert("init".into(), new_init);
        }

        if let Some(ref superclass) = superclass {
            let mut new_methods = Attrs::default();

            // Rc<LoxFunction> is immutable, so we need to drain, change, and replace
            for (name, value) in methods {
                if let Value::Function(method_ref) = value {
                    let superclass = superclass.clone();
                    let mut method = method_ref.as_ref().clone();

                    method.inject_closed_var("super", Value::Class(superclass));

                    new_methods.insert(name, Value::function(method));
                }
            }

            methods = new_methods;
        }

        Rc::new(LoxClass {
            superclass,
            name,
            methods,
        })
    }

    pub fn superclass(&self) -> Option<Rc<LoxClass>> {
        self.superclass.clone()
    }

    pub fn name(&self) -> &SmolStr {
        &self.name
    }

    pub fn arity(&self) -> usize {
        if let Some(Value::Function(method)) = self.methods.get("init") {
            return method.arity();
        }

        0
    }

    pub fn get_method(&self, name: &str, this: LoxReference) -> Option<Value> {
        if let Some(method) = self.methods.get(name) {
            match method {
                Value::Function(method) => {
                    let mut method = method.as_ref().clone();
                    method.inject_closed_var("this", Value::Object(this));
                    return Some(Value::function(method));
                }
                method => panic!("method {name} on class {self} is {method}"),
            }
        }

        if let Some(ref superclass) = self.superclass {
            return superclass.get_method(name, this);
        }

        None
    }
}

impl Display for LoxClass {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "<class {} at {:#X}>",
            self.name,
            (self as *const LoxClass) as usize
        )
    }
}

// slightly hacky: two classes are equal if they are stored at the same location
// since a LoxClass is always (supposed to be) behind an Rc anyways, we might as well compare their
// pointers
impl PartialEq for LoxClass {
    fn eq(&self, other: &Self) -> bool {
        let self_ptr = self as *const LoxClass;
        let other_ptr = other as *const LoxClass;
        self_ptr == other_ptr
    }
}