Collect class instances to be used in gc

Author: dlunch

Cargo.lock

@@ -13,6 +13,7 @@ async-trait = { version = "^0.1", default-features = false }
 bytemuck = { version = "^1.20", default-features = false, features = ["extern_crate_alloc"] }
 dyn-clone = { version = "^1.0", default-features = false }
 dyn-hash = { version = "^0.2", default-features = false }
+hashbrown = { version = "^0.15", features = ["default-hasher"], default-features = false }
 nom = { version = "^7.1", default-features = false, features = ["alloc"] }
 parking_lot = { version = "^0.12", default-features = false }
 tracing = { version = "^0.1", default-features = false, features = ["attributes"] }

Cargo.toml

@@ -8,14 +8,14 @@ license.workspace = true
 async-trait = { workspace = true }
 bytemuck = { workspace = true }
 dyn-clone = { workspace = true }
+hashbrown = { workspace = true }
 parking_lot = { workspace = true }
 tracing = { workspace = true }
 
 async-lock = { version = "^3.4", default-features = false }
 chrono = { version = "^0.4", default-features = false }
 encoding_rs = { version = "^0.8", features = ["alloc"], default-features = false }
 event-listener = { version = "^5.3", default-features = false }
-hashbrown = { version = "^0.15", features = ["default-hasher"], default-features = false }
 zip = { version = "^2.2", features = ["deflate"], default-features = false }
 url = { version = "^2.5", default-features = false }
 

java_runtime/Cargo.toml

@@ -11,6 +11,7 @@ async-trait = { workspace = true }
 bytemuck = { workspace = true }
 dyn-clone = { workspace = true }
 dyn-hash = { workspace = true }
+hashbrown = { workspace = true }
 nom = { workspace = true }
 parking_lot = { workspace = true }
 tracing = { workspace = true }

jvm/Cargo.toml

@@ -28,6 +28,13 @@ pub trait ClassInstance: Sync + Send + AsAny + Debug + DynHash + DynClone + 'sta
 clone_trait_object!(ClassInstance);
 hash_trait_object!(ClassInstance);
 
+impl Eq for dyn ClassInstance {}
+impl PartialEq for dyn ClassInstance {
+    fn eq(&self, other: &Self) -> bool {
+        self.equals(other).unwrap()
+    }
+}
+
 // array wrapper for ClassInstanceRef
 pub struct Array<T>(PhantomData<T>);
 

jvm/src/class_instance.rs

@@ -10,6 +10,7 @@ use core::{
 
 use bytemuck::cast_slice;
 use dyn_clone::clone_box;
+use hashbrown::HashSet;
 use parking_lot::RwLock;
 
 use java_constants::MethodAccessFlags;
@@ -33,6 +34,7 @@ use crate::{
 struct JvmInner {
     classes: RwLock<BTreeMap<String, Class>>,
     threads: RwLock<BTreeMap<u64, JvmThread>>,
+    all_objects: RwLock<HashSet<Box<dyn ClassInstance>>>,
     get_current_thread_id: Box<dyn Fn() -> u64 + Sync + Send>,
     bootstrap_class_loader: Box<dyn BootstrapClassLoader>,
     bootstrapping: AtomicBool,
@@ -53,6 +55,7 @@ impl Jvm {
             inner: Arc::new(JvmInner {
                 classes: RwLock::new(BTreeMap::new()),
                 threads: RwLock::new(BTreeMap::new()),
+                all_objects: RwLock::new(HashSet::new()),
                 get_current_thread_id: Box::new(get_current_thread_id),
                 bootstrap_class_loader: Box::new(bootstrap_class_loader),
                 bootstrapping: AtomicBool::new(true),
@@ -108,6 +111,13 @@ impl Jvm {
 
         let instance = class.definition.instantiate()?;
 
+        let thread_id = (self.inner.get_current_thread_id)();
+        let mut threads = self.inner.threads.write();
+        let thread = threads.get_mut(&thread_id).unwrap();
+
+        thread.top_frame_mut().local_variables().push(instance.clone());
+        self.inner.all_objects.write().insert(instance.clone());
+
         Ok(instance)
     }
 
@@ -131,6 +141,14 @@ impl Jvm {
         let array_class = class.as_array_class_definition().unwrap();
 
         let instance = array_class.instantiate_array(length)?;
+
+        let thread_id = (self.inner.get_current_thread_id)();
+        let mut threads = self.inner.threads.write();
+        let thread = threads.get_mut(&thread_id).unwrap();
+
+        thread.top_frame_mut().local_variables().push(instance.clone());
+        self.inner.all_objects.write().insert(instance.clone());
+
         Ok(instance)
     }
 
@@ -423,10 +441,10 @@ impl Jvm {
         }
     }
 
-    // temporary until we have working gc
     pub fn destroy(&self, instance: Box<dyn ClassInstance>) -> Result<()> {
         tracing::debug!("Destroy {}", instance.class_definition().name());
 
+        self.inner.all_objects.write().remove(&instance);
         instance.destroy();
 
         Ok(())
@@ -496,7 +514,7 @@ impl Jvm {
         let threads = self.inner.threads.read();
         let thread = threads.get(&thread_id).unwrap();
 
-        Ok(thread.top_frame().map(|x| (x.class.clone(), x.class_instance.clone())))
+        Ok(thread.top_java_frame().map(|x| (x.class.clone(), x.class_instance.clone())))
     }
 
     pub async fn register_class(
@@ -553,18 +571,18 @@ impl Jvm {
         let threads = self.inner.threads.read();
         let thread = threads.get(&thread_id).unwrap();
 
-        let mut result = Vec::with_capacity(thread.stack.len());
-
-        for item in thread.stack.iter().rev() {
-            // skip exception classes
-            if self.is_inherited_from(&*item.class.definition, "java/lang/Throwable") {
-                continue;
-            }
-
-            result.push(format!("{}.{}", item.class.definition.name(), item.method));
-        }
-
-        result
+        thread
+            .iter_java_frame()
+            .rev()
+            .filter_map(|x| {
+                // skip exception classes
+                if self.is_inherited_from(&*x.class.definition, "java/lang/Throwable") {
+                    None
+                } else {
+                    Some(format!("{}.{}", x.class.definition.name(), x.method))
+                }
+            })
+            .collect()
     }
 
     async fn register_class_internal(&self, class: Class, class_loader_wrapper: Option<&dyn ClassLoaderWrapper>) -> Result<()> {
@@ -624,6 +642,7 @@ impl Jvm {
     pub fn attach_thread(&self) -> Result<()> {
         let thread_id = (self.inner.get_current_thread_id)();
         self.inner.threads.write().insert(thread_id, JvmThread::new());
+        self.push_native_frame();
 
         Ok(())
     }
@@ -635,6 +654,17 @@ impl Jvm {
         Ok(())
     }
 
+    // TODO we need safe, ergonomic api..
+    pub fn push_native_frame(&self) {
+        let thread_id = (self.inner.get_current_thread_id)();
+        self.inner.threads.write().get_mut(&thread_id).unwrap().push_native_frame();
+    }
+
+    pub fn pop_frame(&self) {
+        let thread_id = (self.inner.get_current_thread_id)();
+        self.inner.threads.write().get_mut(&thread_id).unwrap().pop_frame();
+    }
+
     pub async fn current_class_loader(&self) -> Result<Box<dyn ClassInstance>> {
         let calling_class = self.find_calling_class()?;
 
@@ -703,7 +733,7 @@ impl Jvm {
             .write()
             .get_mut(&thread_id)
             .unwrap()
-            .push_frame(class, class_instance, &method_str);
+            .push_java_frame(class, class_instance, &method_str);
 
         let result = method.run(self, args).await;
 

jvm/src/jvm.rs

@@ -6,34 +6,72 @@ use alloc::{
 
 use crate::{class_loader::Class, ClassInstance};
 
+pub enum StackFrame {
+    Java(JavaStackFrame),
+    Native(NativeStackFrame),
+}
+
+impl StackFrame {
+    pub fn local_variables(&mut self) -> &mut Vec<Box<dyn ClassInstance>> {
+        match self {
+            StackFrame::Java(java_frame) => &mut java_frame.local_variables,
+            StackFrame::Native(native_frame) => &mut native_frame.local_variables,
+        }
+    }
+}
+
 pub struct JvmThread {
-    pub stack: Vec<JvmStackFrame>,
+    stack: Vec<StackFrame>,
 }
 
 impl JvmThread {
     pub fn new() -> Self {
         Self { stack: Vec::new() }
     }
 
-    pub fn push_frame(&mut self, class: &Class, class_instance: Option<Box<dyn ClassInstance>>, method: &str) {
-        self.stack.push(JvmStackFrame {
+    pub fn push_java_frame(&mut self, class: &Class, class_instance: Option<Box<dyn ClassInstance>>, method: &str) {
+        self.stack.push(StackFrame::Java(JavaStackFrame {
             class: class.clone(),
             class_instance,
             method: method.to_string(),
-        });
+            local_variables: Vec::new(),
+        }));
+    }
+
+    pub fn push_native_frame(&mut self) {
+        self.stack.push(StackFrame::Native(NativeStackFrame { local_variables: Vec::new() }));
     }
 
-    pub fn pop_frame(&mut self) -> Option<JvmStackFrame> {
+    pub fn pop_frame(&mut self) -> Option<StackFrame> {
         self.stack.pop()
     }
 
-    pub fn top_frame(&self) -> Option<&JvmStackFrame> {
-        self.stack.last()
+    pub fn top_frame_mut(&mut self) -> &mut StackFrame {
+        self.stack.last_mut().unwrap()
+    }
+
+    pub fn top_java_frame(&self) -> Option<&JavaStackFrame> {
+        self.stack.iter().rev().find_map(|frame| match frame {
+            StackFrame::Java(java_frame) => Some(java_frame),
+            _ => None,
+        })
+    }
+
+    pub fn iter_java_frame(&self) -> impl DoubleEndedIterator<Item = &JavaStackFrame> {
+        self.stack.iter().filter_map(|frame| match frame {
+            StackFrame::Java(java_frame) => Some(java_frame),
+            _ => None,
+        })
     }
 }
 
-pub struct JvmStackFrame {
+pub struct JavaStackFrame {
     pub class: Class,
     pub class_instance: Option<Box<dyn ClassInstance>>,
     pub method: String,
+    pub local_variables: Vec<Box<dyn ClassInstance>>,
+}
+
+pub struct NativeStackFrame {
+    pub local_variables: Vec<Box<dyn ClassInstance>>,
 }

jvm/src/thread.rs

@@ -119,7 +119,11 @@ impl ArrayClassInstance for ArrayClassInstanceImpl {
     fn destroy(self: Box<Self>) {}
 
     fn equals(&self, other: &dyn ClassInstance) -> Result<bool> {
-        let other = other.as_any().downcast_ref::<ArrayClassInstanceImpl>().unwrap();
+        let other = other.as_any().downcast_ref::<ArrayClassInstanceImpl>();
+        if other.is_none() {
+            return Ok(false);
+        }
+        let other = other.unwrap();
 
         Ok(Arc::ptr_eq(&self.inner, &other.inner))
     }

jvm_rust/src/array_class_instance.rs

@@ -40,7 +40,11 @@ impl ClassInstance for ClassInstanceImpl {
     }
 
     fn equals(&self, other: &dyn ClassInstance) -> Result<bool> {
-        let other = other.as_any().downcast_ref::<ClassInstanceImpl>().unwrap();
+        let other = other.as_any().downcast_ref::<ClassInstanceImpl>();
+        if other.is_none() {
+            return Ok(false);
+        }
+        let other = other.unwrap();
 
         Ok(Arc::ptr_eq(&self.inner, &other.inner))
     }