jvm/src/heap_area.rs

use std::collections::HashMap;

use core::fmt::Debug as DebugTrait;

use crate::accessmasks::FieldAccessFlag;
use crate::classfile::{ JavaClassFile, AbstractTypeDescription, MethodInfo, AbstractTypeKind };
use crate::classstore::ClassStore;
use crate::iterators::ClassFieldIterator;
use crate::jvm::Error;

#[derive(Debug)]
pub struct HeapArea {
    pub memory_capacity: usize,
    pub memory_used: usize,
    pub object_area: ObjectArea,
    pub static_area: StaticArea,
}

impl HeapArea {
    pub fn new(memory_capacity: usize) -> Self {
        HeapArea {
            memory_capacity,
            memory_used: 0,
            object_area: ObjectArea::default(),
            static_area: StaticArea::default(),
        }
    }

    pub fn make_object(&mut self, class_store: &ClassStore, class_index: usize) -> ObjectReference {
        let (object_ref, object_size) = self.object_area.make(class_store, class_index);
        self.memory_used += object_size;

        return object_ref;
    }

    pub fn make_empty_array(&mut self, class_store: &ClassStore, element_type_desc: AbstractTypeDescription, capacity: usize) -> ObjectReference {
        let (array_ref, array_size) = self.object_area.make_empty_array(class_store, element_type_desc, capacity);
        self.memory_used += array_size;

        return array_ref;
    }

    pub fn make_array(&mut self, class_store: &ClassStore, elements: Box<[ObjectReference]>) -> ObjectReference {
        let (array_ref, array_size) = self.object_area.make_array(class_store, elements);
        self.memory_used += array_size;

        return array_ref;
    }

    pub fn make_static(&mut self, class: &JavaClassFile, class_index: usize) {
        self.memory_used += self.static_area.make(class, class_index);
    }

    pub fn byte_array_from_rust_string(&mut self, string: &str, class_store: &ClassStore) -> ObjectReference {
        let mut byte_obj_vec = Vec::<ObjectReference>::with_capacity(string.len());
        for byte in string.as_bytes() {
            // TODO: Take bytes from ByteCache

            let byte_class_idx = class_store.class_idx_from_name(&"java/lang/Byte".to_string()).unwrap();
            let byte_obj = self.make_object(&class_store, byte_class_idx);
            self.object_area.set_object_field(byte_obj, "value", FieldValue::Byte(*byte), byte_class_idx, &class_store).unwrap();

            byte_obj_vec.push(byte_obj);
        }

        let byte_array = self.make_array(&class_store, byte_obj_vec.into_boxed_slice());

        byte_array
    }

}

#[derive(Debug, Clone, Copy)]
pub struct ObjectReference(u32);


impl ObjectReference {
    pub const NULL: ObjectReference = ObjectReference(0);
}


const DEFAULT_COMPARTMENT_CAPACITY: u32 = u16::MAX as u32;
const INVALID_FIRST_OBJECT: usize = u16::MAX as usize;
const INVALID_NEXT_COMPARTMENT: usize = 0;

#[derive(Default, Debug)]
pub struct ObjectArea {
    memory_used: usize,
    compartments: Vec<ObjectCompartment>,
    first_free_compartment: usize,
}

impl ObjectArea {
    fn make_empty_array(&mut self, class_store: &ClassStore, element_type_desc: AbstractTypeDescription, capacity: usize) -> (ObjectReference, usize) {
        //
        // make new type desc
        let array_type_desc = AbstractTypeDescription {
            array_level: 1 + element_type_desc.array_level,
            kind: element_type_desc.kind,
        };

        let array_class_ref = class_store.get_array_class_ref(&array_type_desc).unwrap();

        let array_object = HeapArray {
            class_ref: array_class_ref,
            content: vec![ObjectReference::NULL; capacity].into_boxed_slice(),
        };

        let array_size = std::mem::size_of::<HeapArray>() + std::mem::size_of::<ObjectReference>() * array_object.content.len();

        let array_object_ref = self.store_entry(CompartmentEntry::Array(array_object));

        self.memory_used += array_size;

        (array_object_ref, array_size)
    }

    fn make_array(&mut self, class_store: &ClassStore, elements: Box<[ObjectReference]>) -> (ObjectReference, usize) {
        assert!(elements.len() != 0);
        let array_element_class_name = self.get_reference_native_class_name(elements[0], class_store);
        let array_element_type_desc = AbstractTypeDescription::parse_first(array_element_class_name).unwrap().1;

        let (array_object_ref, array_size) = self.make_empty_array(class_store, array_element_type_desc, elements.len());

        for (index, element) in elements.iter().enumerate() {
            self.set_array_element(array_object_ref, index, *element);
        }

        return (array_object_ref, array_size);
    }

    pub fn make(&mut self, class_store: &ClassStore, target_class_index: usize) -> (ObjectReference, usize) {
        let fields: Vec<_> = ClassFieldIterator::new(target_class_index, class_store)
            .filter(|f| ! (f.access_flags & FieldAccessFlag::Static))
            .map(|f| ObjectField { value: FieldValue::default_for(&f.descriptor) })
            .collect();

        let object_size = std::mem::size_of::<HeapObject>() + std::mem::size_of::<ObjectField>() * fields.len();

        let new_object = HeapObject {
            class_index: target_class_index,
            fields: fields.into_boxed_slice(),
        };


        let object_ref = self.store_entry(CompartmentEntry::Object(new_object));

        self.memory_used += object_size;

        return (object_ref, object_size);
    }

    pub fn get_reference_native_class_name<'a>(&self, reference: ObjectReference, class_store: &'a ClassStore) -> &'a String {
        let class_ref = self.get_reference_class_ref(reference, class_store);
        let class_data_ref = match self.get_object_field(class_ref, "classData", self.get_object_class_index(class_ref), class_store).unwrap() {
            FieldValue::Reference(r) => r,
            _ => unreachable!(),
        };
        let native_name_index = match self.get_object_field(class_data_ref, "native_class_descriptor_index", self.get_object_class_index(class_data_ref), class_store).unwrap() {
            FieldValue::Int(i) => i,
            _ => unreachable!(),
        };

        return class_store.get_native_class_name(native_name_index as usize);
    }

    fn get_object_class_index(&self, reference: ObjectReference) -> usize {
        match self.get_object(reference) {
            CompartmentEntry::Object(o) => o.class_index,
            _ => unreachable!(),
        }
    }

    fn get_reference_class_ref(&self, reference: ObjectReference, class_store: &ClassStore) -> ObjectReference {
        match self.get_object(reference) {
            CompartmentEntry::Object(o) => class_store.get_class_objectref_from_index(o.class_index),
            CompartmentEntry::Array(a) => a.class_ref,
            _ => unreachable!(),
        }
    }

    fn store_entry(&mut self, object: CompartmentEntry) -> ObjectReference {
        if self.first_free_compartment == 0 {
            self.compartments.push(ObjectCompartment::new(INVALID_NEXT_COMPARTMENT));
            self.first_free_compartment = self.compartments.len();
        }

        let compartment_index = self.first_free_compartment - 1;
        let object_index = self.compartments.get_mut(compartment_index).unwrap().store(object);
        if self.compartments[compartment_index].first_free == INVALID_FIRST_OBJECT {
            self.first_free_compartment = self.compartments[compartment_index].next_free_compartment;
        }

        return ObjectReference(object_index + (DEFAULT_COMPARTMENT_CAPACITY * (compartment_index as u32 + 1)));
    }

    fn get_object(&self, reference: ObjectReference) -> &CompartmentEntry {
        let index = reference.0;
        let compartment_index: u32 = (index / DEFAULT_COMPARTMENT_CAPACITY) - 1;
        let object_index: u32      = index % DEFAULT_COMPARTMENT_CAPACITY;

        let compartment = self.compartments.get(compartment_index as usize).unwrap();
        let object = compartment.objects.get(object_index as usize).unwrap();

        return object;
    }

    fn get_object_mut(&mut self, reference: ObjectReference) -> &mut CompartmentEntry {
        let index = reference.0;
        let compartment_index: u32 = (index / DEFAULT_COMPARTMENT_CAPACITY) - 1;
        let object_index: u32      = index % DEFAULT_COMPARTMENT_CAPACITY;

        let compartment = self.compartments.get_mut(compartment_index as usize).unwrap();
        let object = compartment.objects.get_mut(object_index as usize).unwrap();

        return object;
    }

    pub fn get_object_field(&self, reference: ObjectReference, field_name: &str, accessing_class_idx: usize, class_store: &ClassStore) -> Result<FieldValue, Error> {
        // TODO: Check access rights
        let object = match self.get_object(reference) {
            CompartmentEntry::Object(o) => o,
            _ => unreachable!(),
        };

        let field_option = ClassFieldIterator::new(object.class_index, class_store)
            .filter(|f| ! (f.access_flags & FieldAccessFlag::Static))
            .filter(|f| f.name == field_name)
            .enumerate()
            .next();

        match field_option {
            Some((field_index, _)) => {
                Ok(object.fields[field_index].value)
            },
            None => Err(
                Error::RunTimeError(
                    format!(
                        "Trying to get field '{}' on instance of '{}' but there is no such field",
                        field_name,
                        class_store.class_file_from_idx(object.class_index).unwrap().get_classname().unwrap()
                    )
                )
            )
        }
    }

    pub fn set_array_element(&mut self, array_reference: ObjectReference, index: usize, element: ObjectReference) {
        let array = match self.get_object_mut(array_reference) {
            CompartmentEntry::Array(a) => a,
            _ => unreachable!(),
        };

        let array_element = array.content.get_mut(index).unwrap();

        *array_element = element;
    }

    pub fn set_object_field(&mut self, reference: ObjectReference, field_name: &str, value: FieldValue, accessing_class_idx: usize, class_store: &ClassStore) -> Result<(), Error> {
        // TODO: Check access rights
        let object = match self.get_object_mut(reference) {
            CompartmentEntry::Object(o) => o,
            _ => unreachable!(),
        };

        let field_option = ClassFieldIterator::new(object.class_index, class_store)
            .filter(|f| ! (f.access_flags & FieldAccessFlag::Static))
            .filter(|f| f.name == field_name)
            .enumerate()
            .next();

        match field_option {
            Some((field_index, field_info)) => {
                let object_field = object.fields.get_mut(field_index).unwrap();
                if std::mem::discriminant(&value) == std::mem::discriminant(&object_field.value) {
                    // TODO: Check reference types
                    object_field.value = value;

                    Ok(())
                } else {
                    Err(
                        Error::RunTimeError(
                            format!(
                                "Trying set value '{:?}' on a field with descriptor '{:?}' in instance of '{}'",
                                value,
                                field_info.descriptor,
                                class_store.class_file_from_idx(object.class_index).unwrap().get_classname().unwrap()
                            )
                        )
                    )
                }
            },
            None => Err(Error::RunTimeError(format!(
                        "Trying to set field '{}' on instance of class '{}' but there is no such field",
                        field_name,
                        class_store.class_file_from_idx(object.class_index).unwrap().get_classname().unwrap()
                        )))
        }
    }
}

pub struct ObjectCompartment {
    objects: Box<[CompartmentEntry]>,
    first_free: usize,
    reserved_count: usize,
    next_free_compartment: usize,
}

impl ObjectCompartment {
    fn new(next_free: usize) -> Self {
        let mut os = Vec::with_capacity(DEFAULT_COMPARTMENT_CAPACITY as usize);
        for i in 0..DEFAULT_COMPARTMENT_CAPACITY-1 {
            os.push(CompartmentEntry::EmptyNext(i as usize + 1));
        }
        os.push(CompartmentEntry::EmptyTail());

        ObjectCompartment {
            objects: os.into_boxed_slice(),
            first_free: 0,
            reserved_count: 0,
            next_free_compartment: next_free,
        }
    }

    fn store(&mut self, object: CompartmentEntry) -> u32 {
        let store_slot = self.objects.get(self.first_free).unwrap();
        let compartment_index = self.first_free;

        match store_slot {
            CompartmentEntry::EmptyNext(next) => {
                self.first_free = *next;
            },
            CompartmentEntry::EmptyTail() => {
                // TODO: Maybe change to something else
                self.first_free = INVALID_FIRST_OBJECT;
            },

            CompartmentEntry::Object(_) => unreachable!(),
            CompartmentEntry::Array(_)  => unreachable!(),
        }
        
        *self.objects.get_mut(compartment_index).unwrap() = object;
        self.reserved_count += 1;

        return compartment_index as u32;
    }
}

impl DebugTrait for ObjectCompartment {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
        f.debug_struct("ObjectCompartment")
            .field("first_free", &self.first_free)
            .field("reserved_count", &self.reserved_count)
            .field("next_free_compartment", &self.next_free_compartment)
            .field("objects", &self.objects
                .iter()
                .enumerate()
                .filter(|e| match e { (_, CompartmentEntry::EmptyNext(_)) | (_, CompartmentEntry::EmptyTail()) => false, _ => true})
                .collect::<Vec<_>>()
            ).finish()
    }
}

#[derive(Debug)]
pub enum CompartmentEntry {
    Object(HeapObject),
    Array(HeapArray),
    EmptyNext(usize),
    EmptyTail(), // last empty value
}

#[derive(Debug)]
pub struct HeapArray {
    class_ref: ObjectReference,
    content: Box<[ObjectReference]>,
}

#[derive(Debug)]
pub struct HeapObject {
    class_index: usize,
    fields: Box<[ObjectField]>,
}

#[derive(Default, Debug)]
pub struct StaticArea {
    memory_used: usize,
    static_objects: HashMap<String, StaticObject>,
}

impl StaticArea {

    pub fn set(&mut self, class_name: &String, field_name: &String, field_value: FieldValue) -> Result<(), Error> {
        // TODO: Access permission checking
        let static_object = match self.static_objects.get_mut(class_name) {
            Some(o) => o,
            None => return Err(Error::RunTimeError(format!("Trying to set '{}.{}={:?}' but there is no such static object.", class_name, field_name, field_value))),
        };

        let field_index = match static_object.field_index_from_name(&field_name) {
            Some(i) => i,
            None => return Err(Error::RunTimeError(format!("Trying to set '{}.{}={:?}' but there is no such static field on this object.", class_name, field_name, field_value))),
        };

        field_value.check_type(&static_object.fields[field_index].descriptor, class_name)?;

        static_object.fields[field_index].value = field_value;

        Ok(())
    }

    fn make(&mut self, class: &JavaClassFile, class_index: usize) -> usize {
        let mut fields = Vec::new();

        let mut fields_cumulative_size = 0;
        for field in &class.fields {
            if field.access_flags & FieldAccessFlag::Static {
                let new_field = StaticField {
                    name: field.name.clone(),
                    descriptor: field.descriptor.clone(),
                    value: FieldValue::default_for(&field.descriptor),
                };
                fields_cumulative_size += std::mem::size_of_val(&new_field);
                fields.push(new_field);
            }
        }

        let new_object = StaticObject {
            class_index,
            fields: fields.into_boxed_slice(),
            methods: Box::new([]),
        };

        let object_memory_size = std::mem::size_of_val(&new_object);
        let field_array_size   = std::mem::size_of_val(&new_object.fields);
        let method_array_size  = std::mem::size_of_val(&new_object.methods);

        let _ = self.static_objects.insert(class.get_classname().unwrap().to_string(), new_object);

        let total_object_size = object_memory_size + field_array_size + fields_cumulative_size + method_array_size;

        self.memory_used += total_object_size;

        return total_object_size;
    }
}

#[derive(Debug)]
pub struct StaticObject {
    pub class_index: usize,
    pub fields: Box<[StaticField]>,
    pub methods: Box<[MethodInfo]>,
}

impl StaticObject {
    fn field_index_from_name(&self, name: &String) -> Option<usize> {
        for (index, field) in self.fields.iter().enumerate() {
            if field.name == *name {
                return Some(index)
            }
        }
        return None;
    }
}

#[derive(Debug)]
pub struct StaticField {
    pub name: String,
    pub descriptor: AbstractTypeDescription,
    pub value: FieldValue,
}

#[derive(Debug)]
pub struct ObjectField {
    pub value: FieldValue,
}

#[derive(Debug, Clone, Copy)]
pub enum FieldValue {
    Boolean(bool),
    Byte(u8),
    Char(u16),
    Short(i16),
    Int(i32),
    Float(f32),
    Reference(ObjectReference),
    Double(f64),
    Long(i64),
}

impl FieldValue {
    fn check_type(&self, t: &AbstractTypeDescription, class_name: &String) -> Result<(), Error> {
        let default_val = FieldValue::default_for(t);

        return if std::mem::discriminant(&default_val) != std::mem::discriminant(self) {
            Err(Error::RunTimeError(format!("Type Mismatch: trying to set {:?} on a field of type {:?} in class {}", self, default_val, class_name)))
        } else {
            Ok(())
        }
    }

    fn default_for(t: &AbstractTypeDescription) -> Self {
        if t.array_level != 0 {
            return Self::Reference(ObjectReference::NULL);
        }
        match &t.kind {
            AbstractTypeKind::Void()       => unreachable!(),
            AbstractTypeKind::Byte()       => Self::Byte(0),
            AbstractTypeKind::Char()       => Self::Char(0),
            AbstractTypeKind::Double()     => Self::Double(0.0),
            AbstractTypeKind::Float()      => Self::Float(0.0),
            AbstractTypeKind::Int()        => Self::Int(0),
            AbstractTypeKind::Long()       => Self::Long(0),
            AbstractTypeKind::Classname(_) => Self::Reference(ObjectReference::NULL),
            AbstractTypeKind::Short()      => Self::Short(0),
            AbstractTypeKind::Boolean()    => Self::Boolean(true),
        }
    }
}