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, FieldAccessFlag };
use crate::bytecode::{ Bytecode, Instruction };
use crate::classfile;
use crate::classfile::{ JavaClassFile, MethodInfo, MethodDescriptor, AbstractTypeDescription, AbstractTypeKind, AttributeInfo, AttributeData, CodeAttributeData, ConstantValueAttributeData };
use crate::classstore;
use crate::classstore::ClassStore;
use crate::constantpool::{ ConstantPoolInfo, ConstantClassInfo, ConstantUtf8Info, ConstantMethodRefInfo, ConstantNameAndTypeInfo};
use crate::heap_area::{ HeapArea, FieldValue, ObjectReference };
use crate::stackframe;
use crate::stackframe::{ StackFrame, StackValue, 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: &[StackValue]) -> 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);
        self.class_store.load_class(&"java/lang/Class".to_string())?;
        self.class_store.load_class(&"java/lang/Object".to_string())?;

        self.new_object(&"java/lang/Class".to_string())?;

        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) => {
                    // TODO: throw exception
                    self.load_class_hierarchy(&name)?;
                },

                JVMCallbackOperation::InitClass(name) => {
                    // TODO: throw exception
                    self.init_class_hierarchy(&name)?;
                }
            }
        }

        Ok(())
    }

    fn new_object(&mut self, class_name: &String) -> Result<ObjectReference, Error> {
        let (_, class_idx) = self.class_store.get_class(class_name)?;
        Ok(self.heap_area.make_object(&mut self.class_store, class_idx))
    }

    fn load_class_hierarchy(&mut self, name: &String) -> Result<(), Error> {
        self.class_store.load_class(&name)?;
        let super_class_name = {
            let (file, _) = self.class_store.get_class(&name)?;
            file.get_super_class_name()?.clone()
        };
        let mut super_classes = vec![super_class_name];

        while super_classes.len() != 0 {
            let current_super    = super_classes.pop().unwrap();
            let have_super_super = {
                let (super_file, _) = self.class_store.get_class(&current_super)?;
                super_file.has_super_class()
            };

            if have_super_super {
                let super_super_name = {
                    let (super_file, _) = self.class_store.get_class(&current_super)?;
                    super_file.get_super_class_name()?
                };
                if self.class_store.have_class(super_super_name) {
                    self.class_store.load_class(&current_super)?;
                } else {
                    super_classes.push(current_super);
                    super_classes.push(super_super_name.to_string());
                }
            } else {
                self.class_store.load_class(&current_super)?;
            }
        }

        Ok(())
    }

    fn init_class_hierarchy(&mut self, name: &String) -> Result<(), Error> {
        let mut class_stack = vec![name.to_string()];

        while class_stack.len() != 0 {
            let current_name = class_stack.pop().unwrap();
            let was_super_init = {
                let (file, _)  = self.class_store.get_class(&current_name)?;
                if ! file.has_super_class() {
                    true
                } else {
                    let super_name = file.get_super_class_name()?;
                    self.class_store.was_init(super_name).unwrap()
                }
            };

            if was_super_init {
                let class_idx = self.class_store.class_idx_from_name(&current_name).unwrap();
                self.init_class(class_idx)?;
            } else {
                let super_name = {
                    let (file, _)  = self.class_store.get_class(&current_name)?;
                    file.get_super_class_name()?
                };
                class_stack.push(current_name);
                class_stack.push(super_name.to_string());
            }
        }

        Ok(())
    }

    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());

        self.heap_area.make_static(class_file, class_idx);

        for field in &class_file.fields {
            if field.access_flags & FieldAccessFlag::Static {
                let cvalue_attrs: Vec<&ConstantValueAttributeData> = (&field.attributes).iter()
                    .filter(|a| match a.data { AttributeData::ConstantValue(_) => true, _ => false })
                    .map(|a| match &a.data { AttributeData::ConstantValue(c) => c, _ => unreachable!() })
                    .collect();

                assert!(cvalue_attrs.len() < 2);
                // TODO: Throw error

                if cvalue_attrs.len() == 1 {
                    let constant_value_info = cvalue_attrs[0];

                    assert!(field.descriptor.array_level == 0);
                    // TODO: Throw Error

                    let field_value = match field.descriptor.kind {
                        AbstractTypeKind::Boolean() => {
                            let int_entry = class_file.pool_int_entry(constant_value_info.constant_value_index)?;

                            FieldValue::Boolean(int_entry.value != 0)
                        },
                        AbstractTypeKind::Int() => {
                            let int_entry = class_file.pool_int_entry(constant_value_info.constant_value_index)?;

                            FieldValue::Int(int_entry.value)
                        },
                        AbstractTypeKind::Short() => {
                            let int_entry = class_file.pool_int_entry(constant_value_info.constant_value_index)?;

                            FieldValue::Short(int_entry.value as i16)
                        },
                        _ => todo!()
                    };

                    self.heap_area.static_area.set(class_file.get_classname()?, &field.name, field_value)?;
                }
            }
        }

        // TODO: Push clinit function

        self.class_store.set_init(class_idx, true);

        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(StackValue::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(StackValue::Int(factor_1 * factor_2)))?;
                }

                Instruction::PushConstInt5() => {
                    let frame_result = frame.operand_stack.push(StackValue::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(StackValue::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, StackValue::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, StackValue::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, StackValue::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, StackValue::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(StackValue),
    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(StackValue::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<StackValue>, 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(
                StackValue::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(
                        StackValue::Byte(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_byte(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Char()              => {
                    arguments.push_front(
                        StackValue::Char(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_char(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Double()            => {
                    arguments.push_front(
                        StackValue::Double1(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_double1(0)
                            )?
                        )
                    );
                    arguments.push_front(
                        StackValue::Double0(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_double0(0)
                            )?
                        )
                    );
                },
                AbstractTypeKind::Float()             => {
                    arguments.push_front(
                        StackValue::Float(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_float(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Int()               => {
                    arguments.push_front(
                        StackValue::Int(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_int(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Long()              => {
                    arguments.push_front(
                        StackValue::Long1(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_long1(0)
                            )?
                        )
                    );
                    arguments.push_front(
                        StackValue::Long0(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_long0(0)
                            )?
                        )
                    );
                },
                AbstractTypeKind::Classname(ref name)   => {
                    // TODO: Type checking
                    arguments.push_front(
                        StackValue::Reference(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_reference(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Short()             => {
                    arguments.push_front(
                        StackValue::Short(
                            wrap_stackframe_error(
                                class, 
                                method, 
                                stack.pop_short(0)
                            )?
                        )
                    )
                },
                AbstractTypeKind::Boolean()           => {
                    arguments.push_front(
                        StackValue::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()))),
    }
}