jvm/src/jvm.rs

use core::fmt::{Display, Formatter};

use std::error::Error as ErrorTrait;

use crate::accessmasks::{ ClassAccessFlagMask, ClassAccessFlag, MethodAccessFlagMask, MethodAccessFlag};
use crate::bytecode::{ Bytecode, Instruction };
use crate::classfile;
use crate::classfile::{ JavaClassFile, MethodInfo, MethodDescriptor, AbstractTypeDescription, AbstractTypeKind, AttributeInfo, AttributeData, CodeAttributeData };
use crate::classstore;
use crate::classstore::ClassStore;
use crate::constantpool::{ ConstantPoolInfo, ConstantClassInfo, ConstantUtf8Info, ConstantMethodRefInfo, ConstantNameAndTypeInfo};
use crate::stackframe::{ StackFrame, Value };

#[derive(Debug)]
pub enum Error {
    ClassStoreError(classstore::Error),
    ClassFileError(classfile::Error),
    BadNameError(String),
    RunTimeError(String),
    OpcodeError(String),
}

impl From<classfile::Error> for Error {
    fn from(value: classfile::Error) -> Self {
        return Error::ClassFileError(value);
    }
}

impl From<classstore::Error> for Error {
    fn from(value: classstore::Error) -> Self {
        return Error::ClassStoreError(value);
    }
}

impl ErrorTrait for Error {}
impl Display for Error {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
            writeln!(formatter, "{self}")?;
            if let Some(e) = self.source() {
                writeln!(formatter, "\tCaused by: {e:?}")?;
            }
            Ok(())
    }
}

#[derive(Debug)]
pub struct JVM {
    class_store: ClassStore,
    stack_frames: Vec<StackFrame>,
}

impl JVM {
    pub fn new() -> Self {
        return JVM {
            class_store: ClassStore::new(),
            stack_frames: Vec::new(),
        }
    }

    pub fn entrypoint(&mut self, class_name: &String, method_name: &String, arguments: &[Value]) -> Result<(), Error> {
        let entry_class = JavaClassFile {
            minor_version: 0,
            major_version: 0,
            constant_pool: Box::new([
                    ConstantPoolInfo::Class(ConstantClassInfo { name_index: 2 }),
                    ConstantPoolInfo::Utf8(ConstantUtf8Info { utf8: "::EntryPoint".to_string() }),
                    ConstantPoolInfo::Utf8(ConstantUtf8Info { utf8: "Code".to_string() }),
                    ConstantPoolInfo::MethodRef(ConstantMethodRefInfo { class_index: 5, name_and_type_index: 6}),
                    ConstantPoolInfo::Class(ConstantClassInfo { name_index: 7 }),
                    ConstantPoolInfo::NameAndType(ConstantNameAndTypeInfo { name_index: 8, descriptor_index: 9 }),
                    ConstantPoolInfo::Utf8(ConstantUtf8Info { utf8: class_name.to_string() }),
                    ConstantPoolInfo::Utf8(ConstantUtf8Info { utf8: method_name.to_string() }),
                    ConstantPoolInfo::Utf8(ConstantUtf8Info { utf8: "()V".to_string() }),
                ]
            ),
            access_flags: ClassAccessFlagMask { mask: ClassAccessFlag::Super.discriminant() },
            this_class: 1,
            super_class: 0,
            interfaces: Box::new([]),
            fields: Box::new([]),
            methods: Box::new([
                MethodInfo {
                    access_flags: MethodAccessFlagMask { 
                        mask: MethodAccessFlag::Public.discriminant() | MethodAccessFlag::Static.discriminant()
                    },
                    name: "call_main".to_string(),
                    descriptor: MethodDescriptor {
                        argument_types: Box::new([]),
                        return_type: AbstractTypeDescription {
                            array_level: 0,
                            kind: AbstractTypeKind::Void(),
                        }
                    },
                    code_attribute_index: 0,
                    attributes: Box::new([
                        AttributeInfo {
                            attribute_name_index: 3,
                            data: AttributeData::Code(
                                CodeAttributeData {
                                    max_stack: 0,
                                    max_locals: 0,
                                    code: Bytecode {
                                        bytes: Box::new([
                                            0xB8_u8.to_be(), // invokestatic
                                            0x04_u16.to_be_bytes()[0], // index 4 into the constant
                                            0x04_u16.to_be_bytes()[1], // pool                      
                                        ]),
                                    },
                                    exception_table: Box::new([]),
                                    attributes: Box::new([]),
                                }
                            )
                        }
                    ])
                }
            ]),
            attributes: Box::new([]),
        };

        self.stack_frames.push(
            StackFrame::new(&entry_class, 0, 0, arguments),
        );

        self.class_store.add_class(entry_class, true);

        Ok(())
    }

    pub fn run(&mut self) -> Result<(), Error> {
        while self.stack_frames.len() != 0 {
            let jvm_op = self.bytecode_loop()?;
            match jvm_op {
                JVMCallbackOperation::PopFrame() => self.stack_frames.truncate(self.stack_frames.len() - 1),
                JVMCallbackOperation::PushFrame(frame) => self.stack_frames.push(frame),
                JVMCallbackOperation::LoadClass(name) => {
                    self.class_store.load_class(&name)?;
                    ()
                },
                JVMCallbackOperation::InitClass(name) => {
                    self.init_class(*self.class_store.class_idx_from_name(&name).unwrap());
                }
            }
        }

        Ok(())
    }

    pub fn init_class(&mut self, class_idx: usize) {
        let class_file = self.class_store.class_file_from_idx(class_idx).unwrap();
        let clinit_idx = class_file.find_method_index(&"<clinit>".to_string());

        // TODO: ConstantValue Attributes (final)
        // TODO: Static Stuff

        self.class_store.set_init(class_idx, true);
    }

    fn prepare_invoke_static(&mut self, class_index: usize, method_name: &String, arguments: &[Value]) -> Result<(), Error> {

        let class_file = self.class_store.class_file_from_idx(class_index).unwrap();

        let method_index = class_file.find_method_index(method_name)
            .ok_or(Error::BadNameError(format!("Could not find method '{}' in class '{}'", method_name, class_file.get_classname()?)))?;

        let new_frame = StackFrame::new(
            class_file, 
            class_index, 
            method_index.try_into().expect(&format!("Bad method index: {}", method_index)),
            arguments
            );

        self.stack_frames.push(new_frame);

        return Ok(());
    }

    fn bytecode_loop(&mut self) -> Result<JVMCallbackOperation, Error> {

        let frame     = {
            let frame_index = self.stack_frames.len() - 1;
            &mut self.stack_frames[frame_index]
        };
        let class     =      self.class_store.class_file_from_idx(frame.class_index).unwrap();
        let method    = &    class.methods[frame.method_index as usize];
        let code_attr =      method.get_code_attribute().unwrap();
        let bytecode  = &    code_attr.code;

        while frame.instruction_pointer as usize != bytecode.bytes.len() {
            let (instruction, offset) = bytecode.next_instruction(frame.instruction_pointer as usize);
            frame.instruction_pointer += offset as u32;

            match instruction {
                Instruction::InvokeStatic(methodref_index) => {
                    let (supplied_class_name, supplied_method_name, supplied_descriptor_string) = class.gather_methodref(methodref_index)?;

                    if ! self.class_store.have_class(supplied_class_name) {
                        // rewind the bytecode offset, I'll need to execute this instruction again
                        frame.instruction_pointer -= offset as u32;

                        return Ok(JVMCallbackOperation::LoadClass(supplied_class_name.to_string()));
                    }
                    if ! self.class_store.was_init(supplied_class_name).unwrap() {
                        // rewind the bytecode offset, I'll need to execute this instruction again
                        frame.instruction_pointer -= offset as u32;

                        return Ok(JVMCallbackOperation::InitClass(supplied_class_name.to_string()));
                    }

                    let (callee_class_file, callee_class_index) = self.class_store.get_class(supplied_class_name)?;
                    // TODO: Throw exception on fail
                    let callee_method_index = callee_class_file.find_method_index(supplied_method_name).unwrap();
                    // TODO: Throw exception on fail
                    let callee_method_info  = &callee_class_file.methods[callee_method_index];

                    if ! (callee_method_info.access_flags & MethodAccessFlag::Static) {
                        // TODO: Throw IncompatibleClassChangeError
                        return Err(Error::RunTimeError(format!(
                                    "Invoked method '{}' in class '{}' does not have Access::Static (from invokestatic from '{}' in class '{}')",
                                    method.name,
                                    class.get_classname().unwrap(),
                                    supplied_method_name,
                                    supplied_class_name,
                                    )));
                    }

                    let supplied_descriptor: MethodDescriptor = supplied_descriptor_string.try_into()?;
                    // TODO: Throw exception on fail

                    if supplied_descriptor != callee_method_info.descriptor {
                        // TODO: Throw exception on fail
                        return Err(Error::RunTimeError(format!(
                                    "Mismatched method descriptors between caller and callee: Caller ({}) wanted '{}' but found '{}' on Callee ({})",
                                    class.get_classname().unwrap(),
                                    supplied_descriptor_string,
                                    callee_method_info.descriptor.source_string(),
                                    supplied_class_name,
                                    )));
                    }

                    let arguments = Vec::new();
                    // TODO: Pass arguments

                    let new_frame = StackFrame::new(
                        callee_class_file,
                        callee_class_index,
                        callee_method_index as u16,
                        &arguments.into_boxed_slice(),
                    );

                    return Ok(JVMCallbackOperation::PushFrame(new_frame));
                },

                Instruction::ReturnVoid() => {
                    let expected_type = AbstractTypeDescription {
                        array_level: 0,
                        kind: AbstractTypeKind::Void(),
                    };

                    if method.descriptor.return_type != expected_type {
                        return Err(Error::OpcodeError(format!("Found opcode '{:?}' on method returning '{:?}'", instruction, method.descriptor.return_type)))
                    }

                    return Ok(JVMCallbackOperation::PopFrame());
                },

                _ => {
                    return Err(Error::RunTimeError(format!("Opcode not implemented yet: {:?}", instruction)))
                },

            }
        }

        Ok(JVMCallbackOperation::PopFrame())
    }

}

enum JVMCallbackOperation {
    PopFrame(),
    PushFrame(StackFrame),
    LoadClass(String),
    InitClass(String),
}