jvm-rs/crates/core/src/objects/object_manager.rs

use crate::class::RuntimeClass;
use crate::class_file::ClassFlags;
use crate::class_file::attributes::ArrayType;
use crate::error::VmError;
use crate::objects::array::{Array, ArrayReference};
use crate::objects::object::{
	Object, ObjectReference, Reference, ReferenceKind, string_from_bytes,
};
use crate::value::Value;
use crate::{BaseType, ThreadId, VmSymbols};
use dashmap::DashMap;
use jni::sys::{jint, jlong};
use parking_lot::{Condvar, Mutex};
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};

#[derive(Default)]
pub struct ObjectManager {
	pub objects: HashMap<u32, ReferenceKind>,
	strings: HashMap<String, u32>,
	monitors: DashMap<u32, Arc<Monitor>>,
	max_memory: usize,
	next_id: AtomicU32,
}

impl ObjectManager {
	pub fn max_memory(&self) -> usize {
		self.max_memory
	}

	fn next_id(&self) -> u32 {
		self.next_id.fetch_add(1, Ordering::Relaxed)
	}

	/// Returns the total bytes in use by all managed objects and arrays
	pub fn bytes_in_use(&self) -> usize {
		self.objects
			.values()
			.map(|kind| match kind {
				ReferenceKind::ObjectReference(r) => {
					let guard = r.lock();
					guard
						.fields
						.iter()
						.map(|entry| entry.value().size_bytes())
						.sum()
				}
				ReferenceKind::ArrayReference(a) => a.size_bytes(),
			})
			.sum()
	}

	pub fn new_object(&mut self, class: Arc<RuntimeClass>) -> ObjectReference {
		let id = self.next_id();
		if self.objects.contains_key(&id) {
			let next = self.next_id();
			let key_pairs = self
				.objects
				.iter()
				.map(|(id, obj)| format!("id: {id} obj: {}", obj.class().this_class))
				.collect::<Vec<_>>();
			eprintln!("next: {} all: \n{:#?}", next, key_pairs)
		}
		assert!(
			!self.objects.contains_key(&id),
			"Generated ID already exists!"
		);
		let object = Arc::new(Mutex::new(Object {
			id,
			class: class.clone(),
			fields: Default::default(),
		}));
		self.objects.insert(id, ReferenceKind::from(object.clone()));
		object
	}

	pub fn new_primitive_array(&mut self, class: Arc<RuntimeClass>, count: i32) -> ArrayReference {
		let id = self.next_id();
		assert!(
			!self.objects.contains_key(&id),
			"Generated ID already exists!"
		);

		let symbol = class
			.this_class
			.strip_prefix("[")
			.unwrap()
			.chars()
			.next()
			.unwrap();
		let array_type = symbol.into();

		let array_ref = match array_type {
			BaseType::Int => {
				ArrayReference::Int(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Byte => {
				ArrayReference::Byte(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Short => {
				ArrayReference::Short(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Long => {
				ArrayReference::Long(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Float => {
				ArrayReference::Float(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Double => {
				ArrayReference::Double(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Char => {
				ArrayReference::Char(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
			BaseType::Boolean => {
				ArrayReference::Boolean(Arc::new(Mutex::new(Array::new(id, class, count))))
			}
		};

		self.objects
			.insert(id, ReferenceKind::ArrayReference(array_ref.clone()));
		array_ref
	}

	pub fn new_object_array(&mut self, class: Arc<RuntimeClass>, count: i32) -> ArrayReference {
		let id = self.next_id();
		assert!(
			!self.objects.contains_key(&id),
			"Generated ID already exists!"
		);
		let array_ref = ArrayReference::Object(Arc::new(Mutex::new(Array::new(id, class, count))));

		self.objects
			.insert(id, ReferenceKind::ArrayReference(array_ref.clone()));
		array_ref
	}

	pub fn new_object_array_from(
		&mut self,
		class: Arc<RuntimeClass>,
		values: Box<[Reference]>,
	) -> ArrayReference {
		let id = self.next_id();
		assert!(
			!self.objects.contains_key(&id),
			"Generated ID already exists!"
		);
		let array_ref =
			ArrayReference::Object(Arc::new(Mutex::new(Array::from((id, class, values)))));

		self.objects
			.insert(id, ReferenceKind::ArrayReference(array_ref.clone()));
		array_ref
	}

	pub fn new_byte_array(
		&mut self,
		class: Arc<RuntimeClass>,
		vector: Box<[i8]>,
	) -> ArrayReference {
		let id = self.next_id();
		assert!(
			!self.objects.contains_key(&id),
			"Generated ID already exists!"
		);
		let array_ref =
			ArrayReference::Byte(Arc::new(Mutex::new(Array::from((id, class, vector)))));

		self.objects
			.insert(id, ReferenceKind::ArrayReference(array_ref.clone()));
		array_ref
	}

	pub fn get(&self, id: u32) -> ReferenceKind {
		self.objects
			.get(&id)
			.unwrap_or_else(|| {
				let objs = self
					.objects
					.iter()
					.map(|(x, y)| format!("{x} : {y}"))
					.collect::<Vec<_>>();
				panic!("Object must be present id: {id}\n{objs:#?}")
			})
			.clone()
	}

	pub fn get_interned_string(&self, utf: &str) -> Option<ObjectReference> {
		self.strings
			.get(utf)
			.and_then(|e| self.get(*e).try_into_object_reference().ok())
	}

	pub fn new_string(
		&mut self,
		byte_array_class: Arc<RuntimeClass>,
		string_class: Arc<RuntimeClass>,
		utf8: &str,
	) -> ObjectReference {
		let key = utf8.to_owned();
		let jstr = self.new_object(string_class);
		let byte_vec = utf8
			.encode_utf16()
			.flat_map(|e| e.to_le_bytes())
			.map(|e| e as i8)
			.collect::<Vec<_>>();
		let barray = self.new_byte_array(byte_array_class, byte_vec.into_boxed_slice());

		jstr.lock().fields.insert(
			"value".to_string(),
			Value::from(Some(ReferenceKind::ArrayReference(barray))),
		);
		jstr.lock()
			.fields
			.insert("coder".to_string(), Value::from(1i8));
		jstr.lock()
			.fields
			.insert("hash".to_string(), Value::from(0i32));
		jstr.lock()
			.fields
			.insert("hashIsZero".to_string(), Value::from(false));

		let id = jstr.lock().id;
		debug_assert!(!self.strings.contains_key(&key), "String already interned");
		self.strings.insert(key, id);
		jstr
	}
	pub fn intern_existing_string(&mut self, utf: String, obj_id: u32) -> u32 {
		*self.strings.entry(utf).or_insert(obj_id)
	}

	pub fn new_class(
		&mut self,
		class_class: Arc<RuntimeClass>,
		name: Reference,
		module: Reference,
		modifiers: ClassFlags,
		primitive: bool,
		component_type: Reference,
	) -> ObjectReference {
		let string = self
			.transmute_string(name.clone().unwrap().try_into_object_reference().unwrap())
			.unwrap();
		if string.contains("byte") {
			println!("!!!!!");
		}
		let module = None;
		let modifiers = 17u16;

		let class_redefined_count = 0i32;
		let clazz = self.new_object(class_class);
		{
			let clakz = clazz.lock();
			clakz.set_field("name", Value::from(name));
			clakz.set_field("module", Value::from(module));
			clakz.set_field("modifiers", Value::from(modifiers));
			clakz.set_field("primitive", Value::from(primitive));
			clakz.set_field("classRedefinedCount", Value::from(class_redefined_count));
			clakz.set_field("componentType", Value::from(component_type))
		}
		clazz
	}
	/// using pure wizardry, transmute a string object into a rust string
	pub fn transmute_string(&self, string: ObjectReference) -> Result<String, VmError> {
		let guard = string.lock();
		let value_field = guard
			.fields
			.get("value")
			.ok_or_else(|| VmError::InvariantError("String must have value field".to_string()))?;

		let array_ref = match *value_field {
			Value::Reference(Some(ReferenceKind::ArrayReference(ref arr))) => arr.clone(),
			_ => {
				return Err(VmError::InvariantError(
					"String value field must be a byte array".to_string(),
				));
			}
		};

		let bytes = match array_ref {
			ArrayReference::Byte(byte_arr) => {
				let arr_guard = byte_arr.lock();
				arr_guard.backing.clone()
			}
			_ => {
				return Err(VmError::InvariantError(
					"String value must be a byte array".to_string(),
				));
			}
		};
		Ok(string_from_bytes(&bytes))
	}
	pub fn monitor_enter(&self, thread_id: ThreadId, obj_id: u32) {
		let monitor = self
			.monitors
			.entry(obj_id)
			.or_insert_with(|| Arc::new(Monitor::new()))
			.clone();

		// DashMap entry guard dropped here, map unlocked

		let mut inner = monitor.inner.lock();

		while inner.owner.is_some() && inner.owner != Some(thread_id) {
			monitor.condvar.wait(&mut inner);
		}

		inner.owner = Some(thread_id);
		inner.entry_count += 1;
	}

	pub fn monitor_exit(&self, thread_id: ThreadId, obj_id: u32) -> Result<(), VmError> {
		let monitor = self
			.monitors
			.get(&obj_id)
			.ok_or_else(|| VmError::InvariantError("Monitor not found for object".to_string()))?
			.clone();

		let mut inner = monitor.inner.lock();

		if inner.owner != Some(thread_id) {
			return Err(VmError::InvariantError(
				"Thread does not own monitor".to_string(),
			));
		}

		inner.entry_count -= 1;

		if inner.entry_count == 0 {
			inner.owner = None;
			monitor.condvar.notify_one();
		}

		Ok(())
	}
	pub fn wait(
		&self,
		thread_id: ThreadId,
		obj_id: u32,
		timeout_millis: jlong,
	) -> Result<(), VmError> {
		let monitor = self
			.monitors
			.get(&obj_id)
			.ok_or_else(|| VmError::InvariantError("Monitor not found".into()))?
			.clone();

		let mut inner = monitor.inner.lock();

		if inner.owner != Some(thread_id) {
			return Err(VmError::InvariantError(
				"IllegalMonitorStateException".into(),
			));
		}

		// Save entry count, release monitor
		let saved_count = inner.entry_count;
		inner.entry_count = 0;
		inner.owner = None;
		inner.wait_count += 1;

		// Wake a thread waiting to enter
		monitor.condvar.notify_one();

		// Wait for notify
		if timeout_millis > 0 {
			monitor.wait_condvar.wait_for(
				&mut inner,
				std::time::Duration::from_millis(timeout_millis as u64),
			);
		} else {
			monitor.wait_condvar.wait(&mut inner);
		}

		inner.wait_count -= 1;

		// Reacquire monitor
		while inner.owner.is_some() {
			monitor.condvar.wait(&mut inner);
		}

		inner.owner = Some(thread_id);
		inner.entry_count = saved_count;

		Ok(())
	}

	pub fn notify_one(&self, obj_id: u32) {
		if let Some(monitor) = self.monitors.get(&obj_id) {
			monitor.wait_condvar.notify_one();
		}
	}

	pub fn notify_all(&self, obj_id: u32) {
		if let Some(monitor) = self.monitors.get(&obj_id) {
			monitor.wait_condvar.notify_all();
		}
	}

	pub fn clone_object(&mut self, source: &ReferenceKind) -> Result<ReferenceKind, VmError> {
		match source {
			ReferenceKind::ArrayReference(arr) => self.clone_array(arr),
			ReferenceKind::ObjectReference(obj) => {
				let class = obj.lock().class.clone();
				if !class.implements(VmSymbols::CLONEABLE) {
					return Err(VmError::InvariantError("CloneNotSupported".to_string()));
				}
				self.clone_instance(obj)
			}
		}
	}

	fn clone_array(&mut self, arr: &ArrayReference) -> Result<ReferenceKind, VmError> {
		let new_arr = match arr {
			ArrayReference::Int(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Int(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Byte(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Byte(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Short(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Short(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Long(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Long(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Float(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Float(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Double(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Double(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Char(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Char(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Boolean(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Boolean(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(),
				})))
			}
			ArrayReference::Object(a) => {
				let guard = a.lock();
				let new_id = self.next_id();
				ArrayReference::Object(Arc::new(Mutex::new(Array {
					id: new_id,
					class: guard.class.clone(),
					backing: guard.backing.clone(), // shallow clone - copies references, not objects
				})))
			}
		};

		let id = new_arr.id();
		self.objects
			.insert(id, ReferenceKind::ArrayReference(new_arr.clone()));
		Ok(ReferenceKind::ArrayReference(new_arr))
	}

	fn clone_instance(&mut self, obj: &ObjectReference) -> Result<ReferenceKind, VmError> {
		let guard = obj.lock();
		let class = guard.class.clone();

		// Create a new object with the same class
		let new_obj = self.new_object(class);

		// Copy all fields (shallow copy)
		{
			let new_guard = new_obj.lock();
			for entry in guard.fields.iter() {
				new_guard
					.fields
					.insert(entry.key().clone(), entry.value().clone());
			}
		}

		Ok(ReferenceKind::ObjectReference(new_obj))
	}
}

pub struct Monitor {
	inner: Mutex<MonitorInner>,
	condvar: Condvar,
	wait_condvar: Condvar,
}

struct MonitorInner {
	owner: Option<ThreadId>,
	entry_count: u32,
	wait_count: u32,
}

impl Monitor {
	fn new() -> Self {
		Self {
			inner: Mutex::new(MonitorInner {
				owner: None,
				entry_count: 0,
				wait_count: 0,
			}),
			condvar: Condvar::new(),
			wait_condvar: Condvar::new(),
		}
	}
}

//reflection
impl ObjectManager {
	fn get_from_klass(&self, klass: Arc<RuntimeClass>) -> ReferenceKind {
		self.get(klass.mirror())
	}
	pub fn new_constructor(
		&mut self,
		index: jint,
		modifiers: jint,
		parameters: Box<[Arc<RuntimeClass>]>,
		exceptions: Box<[Arc<RuntimeClass>]>,
		class_array_class: Arc<RuntimeClass>,
		declaring_class: Arc<RuntimeClass>,
		constructor_class: Arc<RuntimeClass>,
	) -> ObjectReference {
		let parameters = parameters
			.iter()
			.map(|c| {
				let mirror_id = c.mirror.get().unwrap();
				Some(self.get(*mirror_id))
			})
			.collect::<Box<[_]>>();
		let exceptions = exceptions
			.iter()
			.map(|c| {
				let mirror_id = c.mirror.get().unwrap();
				Some(self.get(*mirror_id))
			})
			.collect::<Box<[_]>>();
		let constructor_object = self.new_object(constructor_class);
		let parameter_types = self.new_object_array_from(class_array_class.clone(), parameters);
		let exception_types = self.new_object_array_from(class_array_class, exceptions);
		let signature = Value::NULL;
		let annotations = Value::NULL;
		let parameter_annotations = Value::NULL;
		{
			let obj = constructor_object.lock();
			obj.set_field("clazz", Value::from(self.get_from_klass(declaring_class)));
			obj.set_field("slot", Value::from(index));
			obj.set_field("parameterTypes", Value::from(parameter_types));
			obj.set_field("exceptionTypes", Value::from(exception_types));
			obj.set_field("modifiers", Value::from(modifiers));
			obj.set_field("override", Value::from(false));
			obj.set_field("signature", signature);
			obj.set_field("annotations", annotations);
			obj.set_field("parameterAnnotations", parameter_annotations);
		}
		constructor_object
	}

	pub fn new_thread_group(
		&mut self,
		thread_group_class: Arc<RuntimeClass>,
		name: ObjectReference,
		max_priority: jint,
	) -> ObjectReference {
		let tg_object = self.new_object(thread_group_class);
		{
			let obj = tg_object.lock();
			obj.set_field("name", name.into());
			obj.set_field("maxPriority", max_priority.into());
		}
		tg_object
	}
	pub fn new_holder(
		&mut self,
		holder_class: Arc<RuntimeClass>,
		group: ObjectReference,
		priority: jint,
		status: jint,
	) -> ObjectReference {
		let holder_object = self.new_object(holder_class);
		{
			let obj = holder_object.lock();
			obj.set_field("group", group.into());
			obj.set_field("priority", priority.into());
			obj.set_field("threadStatus", status.into());
		}
		holder_object
	}

	pub fn new_thread(
		&mut self,
		thread_class: Arc<RuntimeClass>,
		eetop: jlong,
		tid: jlong,
		name: ObjectReference,
		holder: ObjectReference,
	) -> ObjectReference {
		let tg_object = self.new_object(thread_class);
		{
			let obj = tg_object.lock();
			obj.set_field("eetop", eetop.into());
			obj.set_field("tid", tid.into());
			obj.set_field("name", name.into());
			obj.set_field("holder", holder.into());
		}
		tg_object
	}

	pub fn new_method_type(
		&mut self,
		method_type_class: Arc<RuntimeClass>,
		class_array_class: Arc<RuntimeClass>,
		rtype: ReferenceKind,       // Class mirror for return type
		ptypes: Vec<ReferenceKind>, // Class mirrors for param types
	) -> ObjectReference {
		let ptypes_array =
			self.new_object_array_from(class_array_class, ptypes.into_iter().map(Some).collect());
		let mt = self.new_object(method_type_class);
		{
			let obj = mt.lock();
			obj.set_field("rtype", Value::from(Some(rtype)));
			obj.set_field("ptypes", Value::from(ptypes_array));
		}
		mt
	}

	pub fn new_member_name(
		&mut self,
		member_name_class: Arc<RuntimeClass>,
		clazz: ReferenceKind,  // declaring class mirror
		name: ObjectReference, // method/field name as String
		method_type: ObjectReference,
		flags: jint, // REF_invokeStatic, etc encoded
	) -> ObjectReference {
		let mn = self.new_object(member_name_class);
		{
			let obj = mn.lock();
			obj.set_field("clazz", Value::from(Some(clazz)));
			obj.set_field("name", Value::from(name));
			obj.set_field("type", Value::from(method_type));
			obj.set_field("flags", Value::from(flags));
		}
		mn
	}

	pub fn new_direct_method_handle(
		&mut self,
		dmh_class: Arc<RuntimeClass>,
		member: ObjectReference,
		method_type: ObjectReference,
	) -> ObjectReference {
		let mh = self.new_object(dmh_class);
		{
			let obj = mh.lock();
			obj.set_field("type", Value::from(method_type));
			obj.set_field("member", Value::from(member));
		}
		mh
	}
}