use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use crate::debug::JvmDebugger;
use crate::error::MemoryError;
use crate::security::Sanitizer;
use super::heap_object::{HeapArray, HeapObject};
use super::monitor::Monitor;
use super::Value;
#[derive(Debug)]
pub struct Heap {
objects: HashMap<u32, HeapObject>,
arrays: HashMap<u32, HeapArray>,
next_addr: u32,
marked: HashSet<u32>,
debugger: Option<JvmDebugger>,
sanitizer: Option<Arc<Sanitizer>>,
}
impl Heap {
pub fn new() -> Self {
Heap {
objects: HashMap::new(),
arrays: HashMap::new(),
next_addr: 1,
marked: HashSet::new(),
debugger: None,
sanitizer: None,
}
}
pub fn with_debugger(debugger: JvmDebugger) -> Self {
Heap {
objects: HashMap::new(),
arrays: HashMap::new(),
next_addr: 1,
marked: HashSet::new(),
debugger: Some(debugger),
sanitizer: None,
}
}
pub fn set_sanitizer(&mut self, sanitizer: Option<Arc<Sanitizer>>) {
self.sanitizer = sanitizer;
}
pub fn allocate(&mut self, class_name: String) -> u32 {
let addr = self.next_addr;
self.next_addr += 1;
if let Some(debugger) = &self.debugger {
debugger.log_memory_allocation(addr, 0, &class_name);
}
self.objects.insert(
addr,
HeapObject {
fields: HashMap::new(),
class_name,
marked: false,
string_data: None,
monitor: None,
},
);
addr
}
pub fn allocate_string(&mut self, data: String) -> u32 {
let addr = self.next_addr;
self.next_addr += 1;
if let Some(debugger) = &self.debugger {
debugger.log_memory_allocation(addr, data.len(), "java/lang/String");
}
self.objects.insert(
addr,
HeapObject {
fields: HashMap::new(),
class_name: "java/lang/String".to_string(),
marked: false,
string_data: Some(data),
monitor: None,
},
);
addr
}
pub fn allocate_array(&mut self, array_type: HeapArray) -> u32 {
let addr = self.next_addr;
self.next_addr += 1;
self.arrays.insert(addr, array_type);
addr
}
pub fn get_object(&self, addr: u32) -> Option<&HeapObject> {
self.objects.get(&addr)
}
pub fn get_object_mut(&mut self, addr: u32) -> Option<&mut HeapObject> {
self.objects.get_mut(&addr)
}
pub fn get_field(&self, addr: u32, field_name: &str) -> Option<Value> {
self.objects
.get(&addr)
.and_then(|obj| obj.fields.get(field_name).cloned())
}
pub fn set_field(
&mut self,
addr: u32,
field_name: String,
value: Value,
) -> Result<(), crate::error::MemoryError> {
let obj = self
.objects
.get_mut(&addr)
.ok_or_else(|| crate::error::MemoryError::InvalidReference(addr))?;
obj.fields.insert(field_name, value);
Ok(())
}
pub fn get_string_data(&self, addr: u32) -> Option<&String> {
self.objects.get(&addr).and_then(|obj| {
if obj.class_name == "java/lang/String" {
obj.string_data.as_ref()
} else {
None
}
})
}
pub fn is_string(&self, addr: u32) -> bool {
self.objects
.get(&addr)
.map(|obj| obj.class_name == "java/lang/String")
.unwrap_or(false)
}
pub fn get_array(&self, addr: u32) -> Option<&HeapArray> {
self.arrays.get(&addr)
}
pub fn get_array_mut(&mut self, addr: u32) -> Option<&mut HeapArray> {
self.arrays.get_mut(&addr)
}
pub fn mark(&mut self, addr: u32) {
if let Some(obj) = self.objects.get_mut(&addr) {
obj.marked = true;
self.marked.insert(addr);
}
}
pub fn unmark_all(&mut self) {
for obj in self.objects.values_mut() {
obj.marked = false;
}
self.marked.clear();
}
pub fn collect_garbage(&mut self, roots: &[u32]) -> Result<usize, MemoryError> {
self.unmark_all();
let mut to_mark: Vec<u32> = roots.to_vec();
while let Some(addr) = to_mark.pop() {
if self.marked.contains(&addr) {
continue;
}
self.mark(addr);
if let Some(obj) = self.objects.get(&addr) {
for value in obj.fields.values() {
if let Some(ref_addr) = value.as_reference() {
if !self.marked.contains(&ref_addr) {
to_mark.push(ref_addr);
}
}
}
}
if let Some(array) = self.arrays.get(&addr) {
if let HeapArray::ReferenceArray(_, refs) = array {
for ref_addr in refs {
if !self.marked.contains(ref_addr) {
to_mark.push(*ref_addr);
}
}
}
}
}
let mut freed = 0;
let objects_to_remove: Vec<_> = self
.objects
.iter()
.filter(|(_, obj)| !obj.marked)
.map(|(addr, _)| *addr)
.collect();
let arrays_to_remove: Vec<_> = self
.arrays
.keys()
.filter(|addr| !self.marked.contains(addr))
.cloned()
.collect();
for addr in objects_to_remove {
self.objects.remove(&addr);
freed += 1;
}
for addr in arrays_to_remove {
self.arrays.remove(&addr);
freed += 1;
}
Ok(freed)
}
pub fn object_count(&self) -> usize {
self.objects.len()
}
pub fn array_count(&self) -> usize {
self.arrays.len()
}
pub fn memory_used(&self) -> usize {
let mut total = 0;
for obj in self.objects.values() {
total += obj.class_name.len() + obj.fields.len() * 16;
}
for array in self.arrays.values() {
total += match array {
HeapArray::IntArray(v) => v.len() * 4,
HeapArray::FloatArray(v) => v.len() * 4,
HeapArray::LongArray(v) => v.len() * 8,
HeapArray::DoubleArray(v) => v.len() * 8,
HeapArray::ReferenceArray(_, v) => v.len() * 4,
HeapArray::ByteArray(v) => v.len(),
HeapArray::CharArray(v) => v.len() * 2,
HeapArray::ShortArray(v) => v.len() * 2,
HeapArray::BooleanArray(v) => v.len(),
};
}
total
}
pub fn array_length(&self, addr: u32) -> Result<usize, MemoryError> {
if let Some(ref s) = self.sanitizer {
s.check_null(Some(addr))
.map_err(MemoryError::InvalidArrayOperation)?;
}
self.arrays
.get(&addr)
.map(|array| match array {
HeapArray::IntArray(v) => v.len(),
HeapArray::FloatArray(v) => v.len(),
HeapArray::LongArray(v) => v.len(),
HeapArray::DoubleArray(v) => v.len(),
HeapArray::ReferenceArray(_, v) => v.len(),
HeapArray::ByteArray(v) => v.len(),
HeapArray::CharArray(v) => v.len(),
HeapArray::ShortArray(v) => v.len(),
HeapArray::BooleanArray(v) => v.len(),
})
.ok_or_else(|| MemoryError::InvalidHeapAddress(addr))
}
pub fn array_get(&self, addr: u32, index: usize) -> Result<Value, MemoryError> {
if let Some(ref s) = self.sanitizer {
s.check_null(Some(addr))
.map_err(MemoryError::InvalidArrayOperation)?;
}
let array = self
.arrays
.get(&addr)
.ok_or_else(|| MemoryError::InvalidHeapAddress(addr))?;
let len = match array {
HeapArray::IntArray(v) => v.len(),
HeapArray::FloatArray(v) => v.len(),
HeapArray::LongArray(v) => v.len(),
HeapArray::DoubleArray(v) => v.len(),
HeapArray::ReferenceArray(_, v) => v.len(),
HeapArray::ByteArray(v) => v.len(),
HeapArray::CharArray(v) => v.len(),
HeapArray::ShortArray(v) => v.len(),
HeapArray::BooleanArray(v) => v.len(),
};
if let Some(ref s) = self.sanitizer {
s.check_array_bounds(index, len)
.map_err(MemoryError::InvalidArrayOperation)?;
}
match array {
HeapArray::IntArray(v) => v
.get(index)
.map(|&val| Value::Int(val))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::FloatArray(v) => v
.get(index)
.map(|&val| Value::Float(val))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::LongArray(v) => v
.get(index)
.map(|&val| Value::Long(val))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::DoubleArray(v) => v
.get(index)
.map(|&val| Value::Double(val))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::ReferenceArray(_, v) => v
.get(index)
.map(|&val| Value::Reference(val))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::ByteArray(v) => v
.get(index)
.map(|&val| Value::Int(val as i8 as i32))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::CharArray(v) => v
.get(index)
.map(|&val| Value::Int(val as i32))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::ShortArray(v) => v
.get(index)
.map(|&val| Value::Int(val as i32))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
HeapArray::BooleanArray(v) => v
.get(index)
.map(|&val| Value::Int(if val { 1 } else { 0 }))
.ok_or_else(|| MemoryError::InvalidArrayLength(index)),
}
}
pub fn array_set(&mut self, addr: u32, index: usize, value: Value) -> Result<(), MemoryError> {
if let Some(ref s) = self.sanitizer {
s.check_null(Some(addr))
.map_err(MemoryError::InvalidArrayOperation)?;
}
let array = self
.arrays
.get_mut(&addr)
.ok_or_else(|| MemoryError::InvalidHeapAddress(addr))?;
let len = match array {
HeapArray::IntArray(v) => v.len(),
HeapArray::FloatArray(v) => v.len(),
HeapArray::LongArray(v) => v.len(),
HeapArray::DoubleArray(v) => v.len(),
HeapArray::ReferenceArray(_, v) => v.len(),
HeapArray::ByteArray(v) => v.len(),
HeapArray::CharArray(v) => v.len(),
HeapArray::ShortArray(v) => v.len(),
HeapArray::BooleanArray(v) => v.len(),
};
if let Some(ref s) = self.sanitizer {
s.check_array_bounds(index, len)
.map_err(MemoryError::InvalidArrayOperation)?;
}
match array {
HeapArray::IntArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_int();
}
HeapArray::FloatArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_float();
}
HeapArray::LongArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_long();
}
HeapArray::DoubleArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_double();
}
HeapArray::ReferenceArray(_, v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_reference().ok_or_else(|| {
MemoryError::InvalidArrayType("Expected reference".to_string())
})?;
}
HeapArray::ByteArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = (value.as_int() as i8) as u8;
}
HeapArray::CharArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_int() as u16;
}
HeapArray::ShortArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_int() as i16;
}
HeapArray::BooleanArray(v) => {
if index >= v.len() {
return Err(MemoryError::InvalidArrayLength(index));
}
v[index] = value.as_int() != 0;
}
}
Ok(())
}
pub fn monitor_enter(&mut self, addr: u32, thread_id: u32) -> Result<(), MemoryError> {
let obj = self
.get_object_mut(addr)
.ok_or(MemoryError::InvalidReference(addr))?;
if obj.monitor.is_none() {
obj.monitor = Some(Monitor::new());
}
let monitor = obj.monitor.as_mut().unwrap();
if monitor.enter(thread_id) {
Ok(())
} else {
Err(MemoryError::InvalidMonitorState)
}
}
pub fn monitor_exit(&mut self, addr: u32, thread_id: u32) -> Result<(), MemoryError> {
let obj = self
.get_object_mut(addr)
.ok_or(MemoryError::InvalidReference(addr))?;
if let Some(monitor) = &mut obj.monitor {
if Monitor::exit(monitor, thread_id) {
Ok(())
} else {
Err(MemoryError::IllegalMonitorState)
}
} else {
Err(MemoryError::IllegalMonitorState)
}
}
pub fn owns_monitor(&self, addr: u32, thread_id: u32) -> bool {
if let Some(obj) = self.get_object(addr) {
if let Some(monitor) = &obj.monitor {
return monitor.is_owned_by(thread_id);
}
}
false
}
}
impl Default for Heap {
fn default() -> Self {
Self::new()
}
}