Add JNI-style global references

Author: dlunch

java_runtime/src/classes/java/lang/thread.rs

@@ -3,7 +3,7 @@ use core::time::Duration;
 
 use java_class_proto::{JavaFieldProto, JavaMethodProto};
 use java_constants::{FieldAccessFlags, MethodAccessFlags};
-use jvm::{ClassInstanceRef, Jvm, Result, runtime::JavaLangString};
+use jvm::{ClassInstanceRef, GlobalRef, Jvm, Result, runtime::JavaLangString};
 
 use crate::{
     RuntimeClassProto, RuntimeContext, SpawnCallback,
@@ -191,7 +191,7 @@ impl Thread {
         struct ThreadStartProxy {
             jvm: Jvm,
             thread_id: i32,
-            this: ClassInstanceRef<Thread>,
+            this: GlobalRef<Thread>,
         }
 
         #[async_trait::async_trait]
@@ -226,7 +226,7 @@ impl Thread {
                     }
                 }
 
-                let mut this = self.this.clone();
+                let mut this = (*self.this).clone();
                 let cleanup = if let Err(error) = self.jvm.monitor_enter(&self.this).await {
                     Err(error)
                 } else {
@@ -253,12 +253,16 @@ impl Thread {
 
         let id: i32 = jvm.invoke_virtual(&this, "hashCode", "()I", ()).await?;
 
+        let this = match jvm.new_global_ref(&this) {
+            Some(this) => this,
+            None => return Err(jvm.exception("java/lang/NullPointerException", "thread is null").await),
+        };
         context.spawn(
             jvm,
             Box::new(ThreadStartProxy {
                 jvm: jvm.clone(),
                 thread_id: id,
-                this: this.clone(),
+                this,
             }),
         );
 

jvm/src/garbage_collector.rs

@@ -8,6 +8,7 @@ use crate::{ClassDefinition, ClassInstance, Field, JavaValue, Jvm, class_loader:
 pub fn determine_garbage(
     jvm: &Jvm,
     threads: &BTreeMap<u64, JvmThread>,
+    global_references: &BTreeMap<u64, Box<dyn ClassInstance>>,
     all_class_instances: &HashSet<Box<dyn ClassInstance>>,
     classes: &BTreeMap<String, Class>,
     interned_strings: &[Box<dyn ClassInstance>],
@@ -30,6 +31,10 @@ pub fn determine_garbage(
         find_reachable_objects(jvm, x, &mut reachable_objects);
     });
 
+    global_references.values().for_each(|object| {
+        find_reachable_objects(jvm, object, &mut reachable_objects);
+    });
+
     interned_strings.iter().for_each(|x| {
         find_reachable_objects(jvm, x, &mut reachable_objects);
     });

jvm/src/global_ref.rs

@@ -0,0 +1,31 @@
+use alloc::{boxed::Box, collections::BTreeMap, sync::Arc};
+use core::{ops::Deref, sync::atomic::AtomicU64};
+
+use parking_lot::RwLock;
+
+use crate::{ClassInstance, ClassInstanceRef};
+
+pub(crate) struct GlobalReferences {
+    pub(crate) next_id: AtomicU64,
+    pub(crate) objects: RwLock<BTreeMap<u64, Box<dyn ClassInstance>>>,
+}
+
+pub struct GlobalRef<T> {
+    pub(crate) references: Arc<GlobalReferences>,
+    pub(crate) id: u64,
+    pub(crate) reference: ClassInstanceRef<T>,
+}
+
+impl<T> Deref for GlobalRef<T> {
+    type Target = ClassInstanceRef<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.reference
+    }
+}
+
+impl<T> Drop for GlobalRef<T> {
+    fn drop(&mut self) {
+        self.references.objects.write().remove(&self.id);
+    }
+}

jvm/src/jvm.rs

@@ -4,7 +4,7 @@ use alloc::{borrow::ToOwned, boxed::Box, collections::BTreeMap, format, string::
 use core::{
     fmt::Debug,
     iter,
-    sync::atomic::{AtomicBool, Ordering},
+    sync::atomic::{AtomicBool, AtomicU64, Ordering},
 };
 
 use dyn_clone::clone_box;
@@ -17,13 +17,14 @@ use crate::{
     Result,
     array_class_instance::{ArrayRawBuffer, ArrayRawBufferMut},
     class_definition::ClassDefinition,
-    class_instance::ClassInstance,
+    class_instance::{ClassInstance, ClassInstanceRef},
     class_loader::{
         BootstrapClassLoader, BootstrapClassLoaderWrapper, Class, ClassLoaderWrapper, InitState, InitializationAction, JavaClassLoaderWrapper,
     },
     error::JavaError,
     field::Field,
     garbage_collector::determine_garbage,
+    global_ref::{GlobalRef, GlobalReferences},
     invoke_arg::InvokeArg,
     method::Method,
     monitor::{Monitor, MonitorWait, MonitorWaitTimeout},
@@ -36,6 +37,7 @@ use crate::{
 struct JvmInner {
     classes: RwLock<BTreeMap<String, Class>>,
     threads: RwLock<BTreeMap<u64, JvmThread>>,
+    global_references: Arc<GlobalReferences>,
     all_objects: RwLock<HashSet<Box<dyn ClassInstance>>>,
     string_pool: RwLock<BTreeMap<Vec<u16>, Box<dyn ClassInstance>>>,
     monitors: RwLock<BTreeMap<usize, Arc<Monitor>>>,
@@ -59,6 +61,10 @@ impl Jvm {
             inner: Arc::new(JvmInner {
                 classes: RwLock::new(BTreeMap::new()),
                 threads: RwLock::new(BTreeMap::new()),
+                global_references: Arc::new(GlobalReferences {
+                    next_id: AtomicU64::new(0),
+                    objects: RwLock::new(BTreeMap::new()),
+                }),
                 all_objects: RwLock::new(HashSet::new()),
                 string_pool: RwLock::new(BTreeMap::new()),
                 monitors: RwLock::new(BTreeMap::new()),
@@ -107,6 +113,16 @@ impl Jvm {
 
         jvm.inner.bootstrapping.store(false, Ordering::Relaxed);
 
+        let thread_id = (jvm.inner.get_current_thread_id)();
+        jvm.inner
+            .threads
+            .write()
+            .get_mut(&thread_id)
+            .unwrap()
+            .top_frame_mut()
+            .local_variables_mut()
+            .clear();
+
         Ok(jvm)
     }
 
@@ -188,7 +204,19 @@ impl Jvm {
 
             self.ensure_initialized(&declaring_class).await?;
 
-            Ok(declaring_class.definition.get_static_field(&*field)?.into())
+            let value = declaring_class.definition.get_static_field(&*field)?;
+            if let JavaValue::Object(Some(instance)) = &value {
+                let thread_id = (self.inner.get_current_thread_id)();
+                self.inner
+                    .threads
+                    .write()
+                    .get_mut(&thread_id)
+                    .unwrap()
+                    .top_frame_mut()
+                    .local_variables_mut()
+                    .push(instance.clone());
+            }
+            Ok(value.into())
         } else {
             Err(self
                 .exception("java/lang/NoSuchFieldError", &format!("{class_name}.{name}:{descriptor}"))
@@ -230,7 +258,19 @@ impl Jvm {
         let field = self.find_field(&*instance.class_definition(), name, descriptor)?;
 
         if let Some(field) = field {
-            Ok(instance.get_field(&*field)?.into())
+            let value = instance.get_field(&*field)?;
+            if let JavaValue::Object(Some(instance)) = &value {
+                let thread_id = (self.inner.get_current_thread_id)();
+                self.inner
+                    .threads
+                    .write()
+                    .get_mut(&thread_id)
+                    .unwrap()
+                    .top_frame_mut()
+                    .local_variables_mut()
+                    .push(instance.clone());
+            }
+            Ok(value.into())
         } else {
             Err(self
                 .exception(
@@ -407,6 +447,15 @@ impl Jvm {
         if let Some(array) = array {
             let values = array.load(offset, count)?;
 
+            let thread_id = (self.inner.get_current_thread_id)();
+            let mut threads = self.inner.threads.write();
+            let local_variables = threads.get_mut(&thread_id).unwrap().top_frame_mut().local_variables_mut();
+            values.iter().for_each(|value| {
+                if let JavaValue::Object(Some(instance)) = value {
+                    local_variables.push(instance.clone());
+                }
+            });
+
             Ok(iter::IntoIterator::into_iter(values).map(|x| x.into()).collect::<Vec<_>>())
         } else {
             Err(self.exception("java/lang/IllegalArgumentException", "Not an array").await)
@@ -761,11 +810,12 @@ impl Jvm {
 
         let garbage = {
             let threads = self.inner.threads.read();
+            let global_references = self.inner.global_references.objects.read();
             let all_objects = self.inner.all_objects.read();
             let classes = self.inner.classes.read();
             let interned_strings = self.interned_strings();
 
-            determine_garbage(self, &threads, &all_objects, &classes, &interned_strings)
+            determine_garbage(self, &threads, &global_references, &all_objects, &classes, &interned_strings)
         };
 
         let garbage_count = garbage.len();
@@ -888,6 +938,18 @@ impl Jvm {
         Ok(())
     }
 
+    pub fn new_global_ref<T>(&self, reference: &ClassInstanceRef<T>) -> Option<GlobalRef<T>> {
+        let instance = reference.instance.as_ref()?.clone();
+        let id = self.inner.global_references.next_id.fetch_add(1, Ordering::Relaxed);
+        self.inner.global_references.objects.write().insert(id, instance);
+
+        Some(GlobalRef {
+            references: self.inner.global_references.clone(),
+            id,
+            reference: reference.clone(),
+        })
+    }
+
     pub fn detach_thread(&self) -> Result<()> {
         let thread_id = (self.inner.get_current_thread_id)();
         self.inner.threads.write().remove(&thread_id);
@@ -1043,13 +1105,25 @@ impl Jvm {
             .write()
             .get_mut(&thread_id)
             .unwrap()
-            .push_java_frame(class, class_instance, &method_str);
+            .push_java_frame(class, class_instance, &method_str, &args);
 
         let result = method.run(self, args).await;
 
         tracing::trace!("Execute result: {result:?}");
 
-        self.inner.threads.write().get_mut(&thread_id).unwrap().pop_frame();
+        let returned_reference = match &result {
+            Ok(JavaValue::Object(Some(instance))) => Some(instance.clone()),
+            Err(JavaError::JavaException(exception)) => Some(exception.clone()),
+            _ => None,
+        };
+        {
+            let mut threads = self.inner.threads.write();
+            let thread = threads.get_mut(&thread_id).unwrap();
+            thread.pop_frame();
+            if let Some(returned_reference) = returned_reference {
+                thread.top_frame_mut().local_variables_mut().push(returned_reference);
+            }
+        }
 
         if let Some(object) = &synchronized_object
             && let Err(error) = self.monitor_exit(object).await

jvm/src/lib.rs

@@ -10,6 +10,7 @@ mod class_loader;
 mod error;
 mod field;
 mod garbage_collector;
+mod global_ref;
 mod invoke_arg;
 mod jvm;
 mod method;
@@ -38,6 +39,7 @@ pub use self::{
     class_loader::BootstrapClassLoader,
     error::JavaError,
     field::Field,
+    global_ref::GlobalRef,
     jvm::Jvm,
     method::Method,
     monitor::{MonitorWait, MonitorWaitTimeout},

jvm/src/thread.rs

@@ -4,7 +4,7 @@ use alloc::{
     vec::Vec,
 };
 
-use crate::{ClassInstance, class_loader::Class};
+use crate::{ClassInstance, JavaValue, class_loader::Class};
 
 pub enum StackFrame {
     Java(JavaStackFrame),
@@ -49,12 +49,18 @@ impl JvmThread {
         self.java_thread = Some(java_thread);
     }
 
-    pub fn push_java_frame(&mut self, class: &Class, class_instance: Option<Box<dyn ClassInstance>>, method: &str) {
+    pub fn push_java_frame(&mut self, class: &Class, class_instance: Option<Box<dyn ClassInstance>>, method: &str, args: &[JavaValue]) {
         self.stack.push(StackFrame::Java(JavaStackFrame {
             class: class.clone(),
             class_instance,
             method: method.to_string(),
-            local_variables: Vec::new(),
+            local_variables: args
+                .iter()
+                .filter_map(|arg| match arg {
+                    JavaValue::Object(Some(instance)) => Some(instance.clone()),
+                    _ => None,
+                })
+                .collect(),
         }));
     }