hash_map.rs

use alloc::{format, vec, vec::Vec};

use jvm::{Array, ClassInstanceRef, Jvm, Result};
use jvm_class_proto::{JavaFieldProto, JavaMethodProto};
use jvm_types::{ClassAccessFlags, FieldAccessFlags, MethodAccessFlags};

use crate::{RuntimeClassProto, RuntimeContext, classes::java::lang::Object};

use super::HashMapEntry;

const DEFAULT_INITIAL_CAPACITY: i32 = 16;
const DEFAULT_LOAD_FACTOR: f32 = 0.75;

// class java.util.HashMap
pub struct HashMap;

impl HashMap {
    pub fn as_proto() -> RuntimeClassProto {
        RuntimeClassProto {
            name: "java/util/HashMap",
            parent_class: Some("java/util/AbstractMap"),
            interfaces: vec!["java/lang/Cloneable", "java/io/Serializable"],
            methods: vec![
                JavaMethodProto::new("<init>", "()V", Self::init, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("<init>", "(I)V", Self::init_with_capacity, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("<init>", "(IF)V", Self::init_with_capacity_and_load_factor, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("<init>", "(Ljava/util/Map;)V", Self::init_from_map, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("size", "()I", Self::size, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("isEmpty", "()Z", Self::is_empty, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("containsKey", "(Ljava/lang/Object;)Z", Self::contains_key, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("containsValue", "(Ljava/lang/Object;)Z", Self::contains_value, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("get", "(Ljava/lang/Object;)Ljava/lang/Object;", Self::get, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new(
                    "put",
                    "(Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;",
                    Self::put,
                    MethodAccessFlags::PUBLIC,
                ),
                JavaMethodProto::new("putAll", "(Ljava/util/Map;)V", Self::put_all, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new(
                    "remove",
                    "(Ljava/lang/Object;)Ljava/lang/Object;",
                    Self::remove,
                    MethodAccessFlags::PUBLIC,
                ),
                JavaMethodProto::new("clear", "()V", Self::clear, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("keySet", "()Ljava/util/Set;", Self::key_set, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("values", "()Ljava/util/Collection;", Self::values, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("entrySet", "()Ljava/util/Set;", Self::entry_set, MethodAccessFlags::PUBLIC),
                JavaMethodProto::new("initializeMap", "()V", Self::initialize_map, MethodAccessFlags::empty()),
                JavaMethodProto::new(
                    "storeNewEntry",
                    "(ILjava/lang/Object;Ljava/lang/Object;I)V",
                    Self::store_new_entry,
                    MethodAccessFlags::empty(),
                ),
                JavaMethodProto::new(
                    "insertNewEntry",
                    "(ILjava/lang/Object;Ljava/lang/Object;I)V",
                    Self::insert_new_entry,
                    MethodAccessFlags::empty(),
                ),
                JavaMethodProto::new("keyIterator", "()Ljava/util/Iterator;", Self::key_iterator, MethodAccessFlags::empty()),
                JavaMethodProto::new(
                    "valueIterator",
                    "()Ljava/util/Iterator;",
                    Self::value_iterator,
                    MethodAccessFlags::empty(),
                ),
                JavaMethodProto::new(
                    "entryIterator",
                    "()Ljava/util/Iterator;",
                    Self::entry_iterator,
                    MethodAccessFlags::empty(),
                ),
            ],
            fields: vec![
                JavaFieldProto::new("table", "[Ljava/util/HashMap$Entry;", FieldAccessFlags::TRANSIENT),
                JavaFieldProto::new("size", "I", FieldAccessFlags::TRANSIENT),
                JavaFieldProto::new("threshold", "I", FieldAccessFlags::empty()),
                JavaFieldProto::new("loadFactor", "F", FieldAccessFlags::FINAL),
                JavaFieldProto::new("modCount", "I", FieldAccessFlags::TRANSIENT),
            ],
            access_flags: ClassAccessFlags::PUBLIC,
        }
    }

    async fn init(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<()> {
        tracing::debug!("java.util.HashMap::<init>({this:?})");

        let _: () = jvm
            .invoke_special(
                &this,
                "java/util/HashMap",
                "<init>",
                "(IF)V",
                (DEFAULT_INITIAL_CAPACITY, DEFAULT_LOAD_FACTOR),
            )
            .await?;

        Ok(())
    }

    async fn key_set(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::keySet({this:?})");

        let key_set = jvm.new_class("java/util/HashMap$KeySet", "(Ljava/util/HashMap;)V", (this,)).await?;

        Ok(key_set.into())
    }

    async fn values(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::values({this:?})");

        let values = jvm.new_class("java/util/HashMap$Values", "(Ljava/util/HashMap;)V", (this,)).await?;

        Ok(values.into())
    }

    async fn entry_set(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::entrySet({this:?})");

        let entry_set = jvm.new_class("java/util/HashMap$EntrySet", "(Ljava/util/HashMap;)V", (this,)).await?;

        Ok(entry_set.into())
    }

    async fn init_with_capacity(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>, capacity: i32) -> Result<()> {
        tracing::debug!("java.util.HashMap::<init>({this:?}, {capacity:?})");

        jvm.invoke_special(&this, "java/util/HashMap", "<init>", "(IF)V", (capacity, DEFAULT_LOAD_FACTOR))
            .await
    }

    async fn init_with_capacity_and_load_factor(
        jvm: &Jvm,
        _: &mut RuntimeContext,
        mut this: ClassInstanceRef<Self>,
        capacity: i32,
        load_factor: f32,
    ) -> Result<()> {
        tracing::debug!("java.util.HashMap::<init>({this:?}, {capacity:?}, {load_factor:?})");

        if capacity < 0 {
            return Err(jvm
                .exception("java/lang/IllegalArgumentException", &format!("Illegal Capacity: {capacity}"))
                .await);
        }
        if load_factor <= 0.0 || load_factor.is_nan() {
            return Err(jvm
                .exception("java/lang/IllegalArgumentException", &format!("Illegal Load: {load_factor}"))
                .await);
        }

        let _: () = jvm.invoke_special(&this, "java/util/AbstractMap", "<init>", "()V", ()).await?;

        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.instantiate_array("Ljava/util/HashMap$Entry;", capacity as usize).await?.into();
        jvm.put_field(&mut this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;", table)
            .await?;
        jvm.put_field(&mut this, "java/util/HashMap", "size", "I", 0).await?;
        jvm.put_field(
            &mut this,
            "java/util/HashMap",
            "threshold",
            "I",
            Self::threshold_for_capacity(capacity, load_factor),
        )
        .await?;
        jvm.put_field(&mut this, "java/util/HashMap", "loadFactor", "F", load_factor).await?;
        jvm.put_field(&mut this, "java/util/HashMap", "modCount", "I", 0).await?;
        let _: () = jvm.invoke_virtual(&this, "java/util/HashMap", "initializeMap", "()V", ()).await?;

        Ok(())
    }

    async fn init_from_map(jvm: &Jvm, _: &mut RuntimeContext, mut this: ClassInstanceRef<Self>, map: ClassInstanceRef<Object>) -> Result<()> {
        tracing::debug!("java.util.HashMap::<init>({this:?}, {map:?})");

        if map.is_null() {
            return Err(jvm.exception("java/lang/NullPointerException", "map").await);
        }
        let size: i32 = jvm.invoke_virtual(&map, &map.class_definition().name(), "size", "()I", ()).await?;
        let capacity = size.saturating_mul(2).max(DEFAULT_INITIAL_CAPACITY);
        let _: () = jvm
            .invoke_special(&this, "java/util/HashMap", "<init>", "(IF)V", (capacity, DEFAULT_LOAD_FACTOR))
            .await?;
        let entry_set: ClassInstanceRef<Object> = jvm
            .invoke_virtual(&map, &map.class_definition().name(), "entrySet", "()Ljava/util/Set;", ())
            .await?;
        let entries: ClassInstanceRef<Array<Object>> = jvm
            .invoke_virtual(&entry_set, &entry_set.class_definition().name(), "toArray", "()[Ljava/lang/Object;", ())
            .await?;
        let count = jvm.array_length(&entries).await?;
        for entry in jvm.load_array::<ClassInstanceRef<Object>>(&entries, 0, count).await? {
            let key: ClassInstanceRef<Object> = jvm
                .invoke_virtual(&entry, &entry.class_definition().name(), "getKey", "()Ljava/lang/Object;", ())
                .await?;
            let value: ClassInstanceRef<Object> = jvm
                .invoke_virtual(&entry, &entry.class_definition().name(), "getValue", "()Ljava/lang/Object;", ())
                .await?;
            Self::put_for_create(jvm, &mut this, key, value).await?;
        }

        Ok(())
    }

    async fn size(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<i32> {
        tracing::debug!("java.util.HashMap::size({this:?})");

        jvm.get_field(&this, "java/util/HashMap", "size", "I").await
    }

    async fn is_empty(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<bool> {
        tracing::debug!("java.util.HashMap::isEmpty({this:?})");

        let size: i32 = jvm.get_field(&this, "java/util/HashMap", "size", "I").await?;

        Ok(size == 0)
    }

    async fn contains_key(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>, key: ClassInstanceRef<Object>) -> Result<bool> {
        tracing::debug!("java.util.HashMap::containsKey({this:?}, {key:?})");

        let entry = Self::find_entry(jvm, &this, &key).await?;

        Ok(!entry.is_null())
    }

    async fn contains_value(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>, value: ClassInstanceRef<Object>) -> Result<bool> {
        tracing::debug!("java.util.HashMap::containsValue({this:?}, {value:?})");

        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(&this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        for bucket_index in 0..table_len {
            let mut entry = Self::load_bucket(jvm, &table, bucket_index).await?;
            while !entry.is_null() {
                let entry_value: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "value", "Ljava/lang/Object;").await?;
                if Self::object_equals(jvm, &value, &entry_value).await? {
                    return Ok(true);
                }

                entry = jvm
                    .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                    .await?;
            }
        }

        Ok(false)
    }

    async fn get(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>, key: ClassInstanceRef<Object>) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::get({this:?}, {key:?})");

        let entry = Self::find_entry(jvm, &this, &key).await?;
        if entry.is_null() {
            return Ok(None.into());
        }

        jvm.get_field(&entry, "java/util/HashMap$Entry", "value", "Ljava/lang/Object;").await
    }

    async fn put(
        jvm: &Jvm,
        _: &mut RuntimeContext,
        mut this: ClassInstanceRef<Self>,
        key: ClassInstanceRef<Object>,
        value: ClassInstanceRef<Object>,
    ) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::put({this:?}, {key:?}, {value:?})");

        let key_hash = Self::object_hash_or_zero(jvm, &key).await?;
        Self::ensure_table_for_insert(jvm, &mut this).await?;

        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(&this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        let bucket_index = Self::bucket_index(key_hash, table_len);

        let mut entry = Self::load_bucket(jvm, &table, bucket_index).await?;
        while !entry.is_null() {
            let entry_hash: i32 = jvm.get_field(&entry, "java/util/HashMap$Entry", "hash", "I").await?;
            if entry_hash == key_hash {
                let entry_key: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "key", "Ljava/lang/Object;").await?;
                if Self::keys_equal(jvm, &key, &entry_key).await? {
                    let old_value: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "value", "Ljava/lang/Object;").await?;
                    jvm.put_field(&mut entry, "java/util/HashMap$Entry", "value", "Ljava/lang/Object;", value)
                        .await?;
                    let _: () = jvm
                        .invoke_virtual(&entry, "java/util/HashMap$Entry", "onAccess", "(Ljava/util/HashMap;)V", (this.clone(),))
                        .await?;
                    return Ok(old_value);
                }
            }

            entry = jvm
                .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                .await?;
        }

        let _: () = jvm
            .invoke_virtual(
                &this,
                "java/util/HashMap",
                "insertNewEntry",
                "(ILjava/lang/Object;Ljava/lang/Object;I)V",
                (key_hash, key, value, bucket_index as i32),
            )
            .await?;

        Ok(None.into())
    }

    async fn put_all(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>, map: ClassInstanceRef<Object>) -> Result<()> {
        tracing::debug!("java.util.HashMap::putAll({this:?}, {map:?})");

        jvm.invoke_special(&this, "java/util/AbstractMap", "putAll", "(Ljava/util/Map;)V", (map,))
            .await
    }

    async fn remove(
        jvm: &Jvm,
        _: &mut RuntimeContext,
        mut this: ClassInstanceRef<Self>,
        key: ClassInstanceRef<Object>,
    ) -> Result<ClassInstanceRef<Object>> {
        tracing::debug!("java.util.HashMap::remove({this:?}, {key:?})");

        let key_hash = Self::object_hash_or_zero(jvm, &key).await?;
        let mut table: ClassInstanceRef<Array<HashMapEntry>> =
            jvm.get_field(&this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        if table_len == 0 {
            return Ok(None.into());
        }
        let bucket_index = Self::bucket_index(key_hash, table_len);

        let mut previous: ClassInstanceRef<HashMapEntry> = None.into();
        let mut entry = Self::load_bucket(jvm, &table, bucket_index).await?;
        while !entry.is_null() {
            let entry_hash: i32 = jvm.get_field(&entry, "java/util/HashMap$Entry", "hash", "I").await?;
            if entry_hash == key_hash {
                let entry_key: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "key", "Ljava/lang/Object;").await?;
                if Self::keys_equal(jvm, &key, &entry_key).await? {
                    let next: ClassInstanceRef<HashMapEntry> = jvm
                        .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                        .await?;
                    if previous.is_null() {
                        jvm.store_array(&mut table, bucket_index, core::iter::once(next.clone())).await?;
                    } else {
                        jvm.put_field(
                            &mut previous,
                            "java/util/HashMap$Entry",
                            "next",
                            "Ljava/util/HashMap$Entry;",
                            next.clone(),
                        )
                        .await?;
                    }

                    let old_value: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "value", "Ljava/lang/Object;").await?;
                    let _: () = jvm
                        .invoke_virtual(&entry, "java/util/HashMap$Entry", "onRemoval", "(Ljava/util/HashMap;)V", (this.clone(),))
                        .await?;
                    let null_entry: ClassInstanceRef<HashMapEntry> = None.into();
                    jvm.put_field(&mut entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;", null_entry)
                        .await?;

                    let size: i32 = jvm.get_field(&this, "java/util/HashMap", "size", "I").await?;
                    jvm.put_field(&mut this, "java/util/HashMap", "size", "I", size - 1).await?;
                    let mod_count: i32 = jvm.get_field(&this, "java/util/HashMap", "modCount", "I").await?;
                    jvm.put_field(&mut this, "java/util/HashMap", "modCount", "I", mod_count.wrapping_add(1))
                        .await?;

                    return Ok(old_value);
                }
            }

            previous = entry;
            entry = jvm
                .get_field(&previous, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                .await?;
        }

        Ok(None.into())
    }

    async fn clear(jvm: &Jvm, _: &mut RuntimeContext, mut this: ClassInstanceRef<Self>) -> Result<()> {
        tracing::debug!("java.util.HashMap::clear({this:?})");

        let mut table: ClassInstanceRef<Array<HashMapEntry>> =
            jvm.get_field(&this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        if table_len > 0 {
            let nulls: Vec<ClassInstanceRef<HashMapEntry>> = (0..table_len).map(|_| None.into()).collect();
            jvm.store_array(&mut table, 0, nulls).await?;
        }
        jvm.put_field(&mut this, "java/util/HashMap", "size", "I", 0).await?;
        let mod_count: i32 = jvm.get_field(&this, "java/util/HashMap", "modCount", "I").await?;
        jvm.put_field(&mut this, "java/util/HashMap", "modCount", "I", mod_count.wrapping_add(1))
            .await?;

        Ok(())
    }

    async fn initialize_map(_: &Jvm, _: &mut RuntimeContext, _: ClassInstanceRef<Self>) -> Result<()> {
        Ok(())
    }

    async fn store_new_entry(
        jvm: &Jvm,
        _: &mut RuntimeContext,
        mut this: ClassInstanceRef<Self>,
        hash: i32,
        key: ClassInstanceRef<Object>,
        value: ClassInstanceRef<Object>,
        bucket_index: i32,
    ) -> Result<()> {
        let mut table: ClassInstanceRef<Array<HashMapEntry>> =
            jvm.get_field(&this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let existing = Self::load_bucket(jvm, &table, bucket_index as usize).await?;
        let entry: ClassInstanceRef<HashMapEntry> = jvm
            .new_class(
                "java/util/HashMap$Entry",
                "(ILjava/lang/Object;Ljava/lang/Object;Ljava/util/HashMap$Entry;)V",
                (hash, key, value, existing),
            )
            .await?
            .into();
        jvm.store_array(&mut table, bucket_index as usize, core::iter::once(entry)).await?;
        let size: i32 = jvm.get_field(&this, "java/util/HashMap", "size", "I").await?;
        jvm.put_field(&mut this, "java/util/HashMap", "size", "I", size + 1).await
    }

    async fn insert_new_entry(
        jvm: &Jvm,
        _: &mut RuntimeContext,
        mut this: ClassInstanceRef<Self>,
        hash: i32,
        key: ClassInstanceRef<Object>,
        value: ClassInstanceRef<Object>,
        bucket_index: i32,
    ) -> Result<()> {
        let size: i32 = jvm.get_field(&this, "java/util/HashMap", "size", "I").await?;
        let mod_count: i32 = jvm.get_field(&this, "java/util/HashMap", "modCount", "I").await?;
        jvm.put_field(&mut this, "java/util/HashMap", "modCount", "I", mod_count.wrapping_add(1))
            .await?;
        let _: () = jvm
            .invoke_virtual(
                &this,
                "java/util/HashMap",
                "storeNewEntry",
                "(ILjava/lang/Object;Ljava/lang/Object;I)V",
                (hash, key, value, bucket_index),
            )
            .await?;
        let threshold: i32 = jvm.get_field(&this, "java/util/HashMap", "threshold", "I").await?;
        if size >= threshold {
            Self::rehash(jvm, &mut this).await?;
        }

        Ok(())
    }

    async fn key_iterator(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        let snapshot = Self::snapshot_entries(jvm, &this).await?;

        Ok(jvm
            .new_class(
                "java/util/HashMap$KeyIterator",
                "(Ljava/util/HashMap;[Ljava/lang/Object;)V",
                (this, snapshot),
            )
            .await?
            .into())
    }

    async fn value_iterator(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        let snapshot = Self::snapshot_entries(jvm, &this).await?;

        Ok(jvm
            .new_class(
                "java/util/HashMap$ValueIterator",
                "(Ljava/util/HashMap;[Ljava/lang/Object;)V",
                (this, snapshot),
            )
            .await?
            .into())
    }

    async fn entry_iterator(jvm: &Jvm, _: &mut RuntimeContext, this: ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Object>> {
        let snapshot = Self::snapshot_entries(jvm, &this).await?;

        Ok(jvm
            .new_class(
                "java/util/HashMap$EntryIterator",
                "(Ljava/util/HashMap;[Ljava/lang/Object;)V",
                (this, snapshot),
            )
            .await?
            .into())
    }

    pub(super) async fn find_entry(
        jvm: &Jvm,
        this: &ClassInstanceRef<Self>,
        key: &ClassInstanceRef<Object>,
    ) -> Result<ClassInstanceRef<HashMapEntry>> {
        let key_hash = Self::object_hash_or_zero(jvm, key).await?;
        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        if table_len == 0 {
            return Ok(None.into());
        }
        let bucket_index = Self::bucket_index(key_hash, table_len);

        let mut entry = Self::load_bucket(jvm, &table, bucket_index).await?;
        while !entry.is_null() {
            let entry_hash: i32 = jvm.get_field(&entry, "java/util/HashMap$Entry", "hash", "I").await?;
            if entry_hash == key_hash {
                let entry_key: ClassInstanceRef<Object> = jvm.get_field(&entry, "java/util/HashMap$Entry", "key", "Ljava/lang/Object;").await?;
                if Self::keys_equal(jvm, key, &entry_key).await? {
                    return Ok(entry);
                }
            }

            entry = jvm
                .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                .await?;
        }

        Ok(None.into())
    }

    pub(super) async fn put_for_create(
        jvm: &Jvm,
        this: &mut ClassInstanceRef<Self>,
        key: ClassInstanceRef<Object>,
        value: ClassInstanceRef<Object>,
    ) -> Result<()> {
        let hash = Self::object_hash_or_zero(jvm, &key).await?;
        Self::ensure_table_for_insert(jvm, this).await?;
        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let bucket_index = Self::bucket_index(hash, jvm.array_length(&table).await?);
        jvm.invoke_virtual(
            this,
            "java/util/HashMap",
            "storeNewEntry",
            "(ILjava/lang/Object;Ljava/lang/Object;I)V",
            (hash, key, value, bucket_index as i32),
        )
        .await
    }

    async fn ensure_table_for_insert(jvm: &Jvm, this: &mut ClassInstanceRef<Self>) -> Result<()> {
        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        if jvm.array_length(&table).await? > 0 {
            return Ok(());
        }

        let new_capacity = 1;
        let new_table: ClassInstanceRef<Array<HashMapEntry>> = jvm.instantiate_array("Ljava/util/HashMap$Entry;", new_capacity).await?.into();
        jvm.put_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;", new_table)
            .await?;
        let load_factor: f32 = jvm.get_field(this, "java/util/HashMap", "loadFactor", "F").await?;
        jvm.put_field(
            this,
            "java/util/HashMap",
            "threshold",
            "I",
            Self::threshold_for_capacity(new_capacity as i32, load_factor),
        )
        .await?;

        Ok(())
    }

    pub(super) async fn rehash(jvm: &Jvm, this: &mut ClassInstanceRef<Self>) -> Result<()> {
        let old_table: ClassInstanceRef<Array<HashMapEntry>> =
            jvm.get_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let old_capacity = jvm.array_length(&old_table).await?;
        let new_capacity = old_capacity * 2 + 1;

        let mut new_table: ClassInstanceRef<Array<HashMapEntry>> = jvm.instantiate_array("Ljava/util/HashMap$Entry;", new_capacity).await?.into();
        for bucket_index in 0..old_capacity {
            let mut entry = Self::load_bucket(jvm, &old_table, bucket_index).await?;
            while !entry.is_null() {
                let next: ClassInstanceRef<HashMapEntry> = jvm
                    .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                    .await?;
                let entry_hash: i32 = jvm.get_field(&entry, "java/util/HashMap$Entry", "hash", "I").await?;
                let new_index = Self::bucket_index(entry_hash, new_capacity);

                let existing = Self::load_bucket(jvm, &new_table, new_index).await?;
                jvm.put_field(&mut entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;", existing)
                    .await?;
                jvm.store_array(&mut new_table, new_index, core::iter::once(entry)).await?;

                entry = next;
            }
        }

        jvm.put_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;", new_table)
            .await?;
        let load_factor: f32 = jvm.get_field(this, "java/util/HashMap", "loadFactor", "F").await?;
        jvm.put_field(
            this,
            "java/util/HashMap",
            "threshold",
            "I",
            Self::threshold_for_capacity(new_capacity as i32, load_factor),
        )
        .await?;

        Ok(())
    }

    async fn load_bucket(jvm: &Jvm, table: &ClassInstanceRef<Array<HashMapEntry>>, bucket_index: usize) -> Result<ClassInstanceRef<HashMapEntry>> {
        Ok(jvm.load_array(table, bucket_index, 1).await?.remove(0))
    }

    async fn object_equals(jvm: &Jvm, left: &ClassInstanceRef<Object>, right: &ClassInstanceRef<Object>) -> Result<bool> {
        if left.is_null() {
            return Ok(right.is_null());
        }

        if right.is_null() {
            return Ok(false);
        }

        jvm.invoke_virtual(left, "java/lang/Object", "equals", "(Ljava/lang/Object;)Z", (right.clone(),))
            .await
    }

    async fn keys_equal(jvm: &Jvm, query_key: &ClassInstanceRef<Object>, stored_key: &ClassInstanceRef<Object>) -> Result<bool> {
        if query_key.is_null() {
            return Ok(stored_key.is_null());
        }

        if stored_key.is_null() {
            return Ok(false);
        }

        jvm.invoke_virtual(query_key, "java/lang/Object", "equals", "(Ljava/lang/Object;)Z", (stored_key.clone(),))
            .await
    }

    async fn object_hash_or_zero(jvm: &Jvm, value: &ClassInstanceRef<Object>) -> Result<i32> {
        if value.is_null() {
            return Ok(0);
        }

        jvm.invoke_virtual(value, "java/lang/Object", "hashCode", "()I", ()).await
    }

    fn bucket_index(hash: i32, table_len: usize) -> usize {
        (hash & 0x7FFFFFFF) as usize % table_len
    }

    async fn snapshot_entries(jvm: &Jvm, this: &ClassInstanceRef<Self>) -> Result<ClassInstanceRef<Array<Object>>> {
        let table: ClassInstanceRef<Array<HashMapEntry>> = jvm.get_field(this, "java/util/HashMap", "table", "[Ljava/util/HashMap$Entry;").await?;
        let table_len = jvm.array_length(&table).await?;
        let size: i32 = jvm.get_field(this, "java/util/HashMap", "size", "I").await?;
        let mut elements: Vec<ClassInstanceRef<Object>> = Vec::with_capacity(size as usize);

        for bucket_index in 0..table_len {
            let mut entry = Self::load_bucket(jvm, &table, bucket_index).await?;
            while !entry.is_null() {
                elements.push(ClassInstanceRef::new(entry.clone().instance));

                entry = jvm
                    .get_field(&entry, "java/util/HashMap$Entry", "next", "Ljava/util/HashMap$Entry;")
                    .await?;
            }
        }

        let mut snapshot: ClassInstanceRef<Array<Object>> = jvm.instantiate_array("Ljava/lang/Object;", elements.len()).await?.into();
        if !elements.is_empty() {
            jvm.store_array(&mut snapshot, 0, elements).await?;
        }

        Ok(snapshot)
    }

    fn threshold_for_capacity(capacity: i32, load_factor: f32) -> i32 {
        if capacity <= 0 { 0 } else { (capacity as f32 * load_factor) as i32 }
    }
}