jvm/src/jvm.rs

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

use std::collections::VecDeque;
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::heap_area::HeapArea;
use crate::stackframe;
use crate::stackframe::{ StackFrame, Value, OperandStack };

#[derive(Debug)]
pub enum Error {
    ClassStoreError(classstore::Error),
    ClassFileError(classfile::Error),
    StackFrameError(stackframe::Error, String),
    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 {
    pub class_store: ClassStore,
    pub stack_frames: Vec<StackFrame>,
    pub heap_area: HeapArea,
}

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

    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::ReturnFrame(value) => {
                    // Pop returning frame
                    self.stack_frames.truncate(self.stack_frames.len() - 1);

                    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];
                    wrap_stackframe_error(class, method, self.stack_frames.last_mut().unwrap().operand_stack.push(value))?;
                }

                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) -> Result<(), Error> {
        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: Static Stuff
        self.heap_area.static_area.make(class_file, class_idx);

        // TODO: ConstantValue Attributes (final)
        // TODO: Push clinit function

        self.class_store.set_init(class_idx, true);

        Ok(())
    }

    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 mut arguments = VecDeque::new();
                    fill_arguments(class, method, &mut arguments, &callee_method_info.descriptor.argument_types, &mut frame.operand_stack)?;

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

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

                Instruction::LoadByteImmediate(byte) => {
                    // sign extend into int
                    let i8_int = i8::from_be_bytes([byte]);

                    let frame_result = frame.operand_stack.push(Value::Int(i8_int as i32));
                    match frame_result {
                        Ok(_) => (),
                        Err(err) => return Err(Error::StackFrameError(err, format!("in '{}', in class '{}'", method.name, class.get_classname().unwrap()))),
                    }
                },

                Instruction::LoadLocalInt0() => {
                    load_local_int(class, method, frame, 0)?;
                }
                Instruction::LoadLocalInt1() => {
                    load_local_int(class, method, frame, 1)?;
                }
                Instruction::LoadLocalInt2() => {
                    load_local_int(class, method, frame, 2)?;
                }
                Instruction::LoadLocalInt3() => {
                    load_local_int(class, method, frame, 3)?;
                }

                Instruction::MultiplyInt() => {
                    let factor_1 = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(0))?;
                    let factor_2 = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(0))?;
                    wrap_stackframe_error(class, method, frame.operand_stack.push(Value::Int(factor_1 * factor_2)))?;
                }

                Instruction::PushConstInt5() => {
                    let frame_result = frame.operand_stack.push(Value::Int(5));
                    match frame_result {
                        Ok(_) => (),
                        Err(err) => return Err(Error::StackFrameError(err, format!("in '{}', in class '{}'", method.name, class.get_classname().unwrap()))),
                    }
                }

                Instruction::ReturnInt() => {
                    let expected_type = AbstractTypeDescription {
                        array_level: 0,
                        kind: AbstractTypeKind::Int(),
                    };
                    if method.descriptor.return_type != expected_type {
                        return Err(Error::OpcodeError(format!("Found opcode '{:?}' on method returning '{:?}'", instruction, method.descriptor.return_type)))
                    }

                    let int = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(0))?;

                    return Ok(JVMCallbackOperation::ReturnFrame(Value::Int(int)));
                }
                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());
                },

                Instruction::StoreLocalInt0() => {
                    let int = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(0))?;

                    wrap_stackframe_error(class, method, frame.store_local(0, Value::Int(int)))?;
                },
                Instruction::StoreLocalInt1() => {
                    let int = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(1))?;

                    wrap_stackframe_error(class, method, frame.store_local(0, Value::Int(int)))?;
                },
                Instruction::StoreLocalInt2() => {
                    let int = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(2))?;

                    wrap_stackframe_error(class, method, frame.store_local(0, Value::Int(int)))?;
                },
                Instruction::StoreLocalInt3() => {
                    let int = wrap_stackframe_error(class, method, frame.operand_stack.pop_int(3))?;

                    wrap_stackframe_error(class, method, frame.store_local(0, Value::Int(int)))?;
                },

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

            }
        }

        // TODO: Review this, maybe crash when there is no return?
        Ok(JVMCallbackOperation::PopFrame())
    }

}

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

fn load_local_int(class: &JavaClassFile, method: &MethodInfo, frame: &mut StackFrame, index: usize) -> Result<(), Error> {
    let frame_result = frame.load_local_int(index as u16);
    let local_int = match frame_result {
        Ok(i) => {
            i
        },
        Err(err) => return Err(Error::StackFrameError(err, format!("in '{}', in class '{}'", method.name, class.get_classname().unwrap()))),
    };

    let frame_result = frame.operand_stack.push(Value::Int(local_int));
    match frame_result {
        Ok(_) => Ok(()),
        Err(err) => return Err(Error::StackFrameError(err, format!("in '{}', in class '{}'", method.name, class.get_classname().unwrap()))),
    }
}

fn fill_arguments(class: &JavaClassFile, method: &MethodInfo, arguments: &mut VecDeque<Value>, argument_types: &Box<[AbstractTypeDescription]>, stack: &mut OperandStack) -> Result<(), Error> {
    for argument_type in argument_types {
        if argument_type.array_level != 0 {
            // TODO: Type checking
            arguments.push_front(
                Value::Reference(wrap_stackframe_error(class, method, stack.pop_reference(0))?),
            )
        } else {
            match argument_type.kind {
                AbstractTypeKind::Void()              => return Err(Error::RunTimeError("Functions cannot take arguments of type void".to_string())),
                // TODO: Add better description
                AbstractTypeKind::Byte()              => {
                    arguments.push_front(
                        Value::Byte(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_byte(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Char()              => {
                    arguments.push_front(
                        Value::Char(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_char(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Double()            => {
                    arguments.push_front(
                        Value::Double1(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_double1(0)
                            )?
                        )
                    );
                    arguments.push_front(
                        Value::Double0(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_double0(0)
                            )?
                        )
                    );
                },
                AbstractTypeKind::Float()             => {
                    arguments.push_front(
                        Value::Float(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_float(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Int()               => {
                    arguments.push_front(
                        Value::Int(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_int(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Long()              => {
                    arguments.push_front(
                        Value::Long1(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_long1(0)
                            )?
                        )
                    );
                    arguments.push_front(
                        Value::Long0(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_long0(0)
                            )?
                        )
                    );
                },
                AbstractTypeKind::Classname(ref name)   => {
                    // TODO: Type checking
                    arguments.push_front(
                        Value::Reference(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_reference(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Short()             => {
                    arguments.push_front(
                        Value::Short(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_short(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Boolean()           => {
                    arguments.push_front(
                        Value::Boolean(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_boolean(0)
                            )?
                        )
                    )
                },

            }
        }
    }

    Ok(())
}

fn wrap_stackframe_error<T>(class: &JavaClassFile, method: &MethodInfo, frame_result: Result<T, stackframe::Error>) -> Result<T, Error> {
    match frame_result {
        Ok(t) => Ok(t),
        Err(err) => return Err(Error::StackFrameError(err, format!("in '{}', in class '{}'", method.name, class.get_classname().unwrap()))),
    }
}