jvm/src/stackframe.rs

use crate::classfile::{ JavaClassFile, AttributeData };
use crate::heap_area::{ ObjectReference, FieldValue };

#[derive(Copy, Clone, Debug)]
pub enum StackValue {
    Boolean(bool),
    Byte(u8),
    Char(u16),
    Short(u16),
    Int(i32),
    Float(f32),
    Reference(ObjectReference),
    ReturnAddress(u32),
    Double0(u32),
    Double1(u32),
    Long0(u32),
    Long1(u32),
    Empty(),
}

#[derive(Debug)]
pub struct OperandStack {
    stack: Box<[StackValue]>,
    depth: u16,
}

#[derive(Debug)]
pub enum Error {
    PushError(String),
    LocalError(String),
}

impl OperandStack {
    fn new(size: u16) -> Self {
        return OperandStack {
            stack: vec![StackValue::Empty(); size.into()].into_boxed_slice(),
            depth: 0,
        }
    }

    pub fn push(&mut self, value: StackValue) -> Result<(), Error> {
        if self.depth as usize == self.stack.len() {
            return Err(Error::PushError(format!("Trying to push onto full operand stack, capacity: {}, value: {:?}", self.depth, value)))
        }

        self.stack[self.depth as usize] = value;
        self.depth += 1;

        Ok(())
    }

    pub fn push_field_value(&mut self, value: FieldValue) -> Result<(), Error> {
        match value {
            FieldValue::Reference(r) => {
                self.push(StackValue::Reference(r))
            }

            FieldValue::Boolean(b) => {
                self.push(StackValue::Int(if b { 1 } else { 0 }))
            }

            FieldValue::Byte(b) => {
                self.push(StackValue::Int(b as i32))
            }

            _ => {
                println!("{value:?}");
                todo!();
            }
        }
    }
    
    pub fn pop_computational_1(&mut self, index: usize) -> Result<StackValue, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        
        match value {
            StackValue::Long0(_) | StackValue::Long1(_) | StackValue::Double0(_) | StackValue::Double1(_) => {
                Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected type with computational type 1 but found '{:?}'", index, value)))
            },
            _ => Ok(value),
        }
    }

    pub fn pop_boolean(&mut self, index: usize) -> Result<bool, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Boolean(b) => Ok(b),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Boolean but found '{:?}'", index, value)))
        }
    }

    pub fn pop_byte(&mut self, index: usize) -> Result<u8, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        match value {
            StackValue::Byte(b) => Ok(b),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Byte but found '{:?}'", index, value)))
        }
    }

    pub fn pop_char(&mut self, index: usize) -> Result<u16, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Char(c) => Ok(c),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Char but found '{:?}'", index, value)))
        }
    }

    pub fn pop_short(&mut self, index: usize) -> Result<u16, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Short(s) => Ok(s),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Short but found '{:?}'", index, value)))
        }
    }

    pub fn pop_int(&mut self, index: usize) -> Result<i32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Int(i) => Ok(i),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Int but found '{:?}'", index, value)))
        }
    }

    pub fn pop_float(&mut self, index: usize) -> Result<f32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Float(f) => Ok(f),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Float but found '{:?}'", index, value)))
        }
    }

    pub fn pop_reference(&mut self, index: usize) -> Result<ObjectReference, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Reference(o) => Ok(o),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Reference but found '{:?}'", index, value)))
        }
    }

    pub fn pop_returnaddress(&mut self, index: usize) -> Result<u32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::ReturnAddress(a) => Ok(a),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected ReturnAddress but found '{:?}'", index, value)))
        }
    }

    pub fn pop_double0(&mut self, index: usize) -> Result<u32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Double0(a) => Ok(a),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Double0 but found '{:?}'", index, value)))
        }
    }

    pub fn pop_double1(&mut self, index: usize) -> Result<u32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Double1(a) => Ok(a),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Double1 but found '{:?}'", index, value)))
        }
    }

    pub fn pop_long0(&mut self, index: usize) -> Result<u32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Long0(a) => Ok(a),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Long0 but found '{:?}'", index, value)))
        }
    }

    pub fn pop_long1(&mut self, index: usize) -> Result<u32, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let value = self.stack[absolute_index];
        self.depth -= 1;
        match value {
            StackValue::Long1(a) => Ok(a),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function operand stack, expected Long1 but found '{:?}'", index, value)))
        }
    }

    pub fn pop_double(&mut self, index: usize) -> Result<f64, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let higher_bytes = self.stack[absolute_index];
        let lower_bytes  = self.stack[absolute_index + 1];
        self.depth -= 2;
        match (higher_bytes, lower_bytes) {
            (StackValue::Double0(hi), StackValue::Double1(lo)) => {
                let v: u64 = ((hi as u64) << 32) | lo as u64;
                Ok(
                    f64::from_bits(v)
                )
            },
            _ => Err(Error::LocalError(format!("Mismatched types at index {} of the function operand stack, expected (Double0, Double1) but found '{:?}'", index, (higher_bytes, lower_bytes))))
        }
    }

    pub fn pop_long(&mut self, index: usize) -> Result<u64, Error> {
        let absolute_index = self.depth as usize - 1 - index;
        let higher_bytes = self.stack[absolute_index];
        let lower_bytes  = self.stack[absolute_index + 1];
        self.depth -= 2;
        match (higher_bytes, lower_bytes) {
            (StackValue::Long0(hi), StackValue::Long1(lo)) => {
                Ok(((hi as u64) << 32) | lo as u64)
            },
            _ => Err(Error::LocalError(format!("Mismatched types at index {} of the function operand stack, expected (Long0, Long1) but found '{:?}'", index, (higher_bytes, lower_bytes))))
        }
    }
}

#[derive(Debug)]
pub struct StackFrame {
    pub locals: Box<[StackValue]>,
    pub operand_stack: OperandStack,
    pub class_index: usize,
    pub method_index: u16,
    pub instruction_pointer: u32,
}

impl StackFrame {
    pub fn new(classfile: &JavaClassFile, class_index: usize, method_index: u16, arguments: &[StackValue]) -> Self {
        let method_info = &classfile.methods[method_index as usize];
        let code_data = match &method_info.attributes[method_info.code_attribute_index].data {
            AttributeData::Code(data) => data,
            _ => unreachable!(),
        };
        let mut locals = vec![StackValue::Empty(); code_data.max_locals.into()].into_boxed_slice();

        assert!(locals.len() >= arguments.len());

        for (index, v) in arguments.iter().enumerate() {
            locals[index] = *v;
        }

        StackFrame {
            locals,
            operand_stack: OperandStack::new(code_data.max_stack),
            class_index,
            method_index,
            instruction_pointer: 0,
        }
    }

    pub fn load_local_int(&self, index: u16) -> Result<i32, Error> {
        let local = self.locals[index as usize];
        match local {
            StackValue::Int(i) => Ok(i),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function locals, expected Int but found '{:?}'", index, local)))
        }
    }

    pub fn load_local_reference(&self, index: u16) -> Result<ObjectReference, Error> {
        let local = self.locals[index as usize];
        match local {
            StackValue::Reference(r) => Ok(r),
            _ => Err(Error::LocalError(format!("Mismatched type at index {} of the function locals, expected Reference but found '{:?}'", index, local)))
        }
    }

    pub fn store_local(&mut self, index: u16, value: StackValue) -> Result<(), Error> {
        let field = self.locals.get_mut(index as usize);
        match field {
            Some(f) => {
                *f = value;
                Ok(())
            },
            None => {
                Err(Error::LocalError(format!("Tried to get local at index {} when max_locals is {}", index, self.locals.len())))
            }
        }
    }
}