JNI wrapper for Rust's BTreeMap (#76)

* fix merge conflicts

Signed-off-by: pushkarm029 <[email protected]>

* fix

Signed-off-by: pushkarm029 <[email protected]>

* remove comment

Signed-off-by: pushkarm029 <[email protected]>

* fix

Signed-off-by: pushkarm029 <[email protected]>

---------

Signed-off-by: pushkarm029 <[email protected]>

fs-storage/Cargo.toml

@@ -16,7 +16,7 @@ log = { version = "0.4.17", features = ["release_max_level_off"] }
 serde_json = "1.0.82"
 serde = { version = "1.0.138", features = ["derive"] }
 jni = { version = "0.21.1", optional = true }
-jnix = { version = "0.5.1", features = ["derive"] }
+jnix = { version = "0.5.1", features = ["derive"], optional = true }
 
 data-error = { path = "../data-error" }
 
@@ -27,4 +27,4 @@ tempdir = "0.3.7"
 
 [features]
 default = ["jni-bindings"]
-jni-bindings = ["jni"]
+jni-bindings = ["jni", "jnix"]

fs-storage/src/base_storage.rs

@@ -57,6 +57,9 @@ pub trait BaseStorage<K, V>: AsRef<BTreeMap<K, V>> {
     /// from pre-configured location in the filesystem.
     fn read_fs(&mut self) -> Result<&BTreeMap<K, V>>;
 
+    /// Get a value from the internal key-value mapping.
+    fn get(&self, id: &K) -> Option<&V>;
+
     /// Write the internal key-value mapping
     /// to pre-configured location in the filesystem.
     fn write_fs(&mut self) -> Result<()>;

fs-storage/src/btreemap_iter.rs

@@ -0,0 +1,30 @@
+use crate::base_storage::BaseStorage;
+use std::cell::RefCell;
+use std::collections::btree_map::Iter;
+use std::rc::Rc;
+
+pub struct BTreeMapIterator<'a, K, V> {
+    iter: Rc<RefCell<Iter<'a, K, V>>>,
+}
+
+impl<'a, K, V> BTreeMapIterator<'a, K, V>
+where
+    K: Ord + Clone,
+    V: Clone,
+{
+    pub fn new<S: BaseStorage<K, V>>(storage: &'a S) -> Self {
+        BTreeMapIterator {
+            iter: Rc::new(RefCell::new(storage.as_ref().iter())),
+        }
+    }
+
+    pub fn has_next(&mut self) -> bool {
+        let borrow = self.iter.borrow();
+        borrow.clone().peekable().peek().is_some()
+    }
+
+    pub fn native_next(&mut self) -> Option<(K, V)> {
+        let mut borrow = self.iter.borrow_mut();
+        borrow.next().map(|(k, v)| (k.clone(), v.clone()))
+    }
+}

fs-storage/src/file_storage.rs

@@ -242,6 +242,11 @@ where
         Ok(&self.data.entries)
     }
 
+    /// Get a value from the internal mapping
+    fn get(&self, id: &K) -> Option<&V> {
+        self.data.entries.get(id)
+    }
+
     /// Write the data to file
     ///
     /// Update the modified timestamp in file metadata to avoid OS timing issues

fs-storage/src/jni/btreemap_iter.rs

@@ -0,0 +1,75 @@
+use crate::btreemap_iter::BTreeMapIterator;
+use crate::file_storage::FileStorage;
+// This is the interface to the JVM that we'll call the majority of our
+// methods on.
+use jni::JNIEnv;
+
+// These objects are what you should use as arguments to your native
+// function. They carry extra lifetime information to prevent them escaping
+// this context and getting used after being GC'd.
+use jni::objects::{JClass, JValue};
+
+// This is just a pointer. We'll be returning it from our function. We
+// can't return one of the objects with lifetime information because the
+// lifetime checker won't let us.
+use jni::sys::{jboolean, jlong, jobject};
+
+impl BTreeMapIterator<'_, String, String> {
+    pub fn from_jlong(value: jlong) -> &'static mut Self {
+        unsafe { &mut *(value as *mut BTreeMapIterator<String, String>) }
+    }
+}
+
+#[no_mangle]
+pub extern "system" fn Java_dev_arkbuilders_core_BTreeMapIterator_create(
+    _env: JNIEnv<'_>,
+    _class: JClass,
+    file_storage_ptr: jlong,
+) -> jlong {
+    let file_storage = FileStorage::from_jlong(file_storage_ptr);
+    let iter = BTreeMapIterator::new(file_storage);
+    Box::into_raw(Box::new(iter)) as jlong
+}
+
+#[no_mangle]
+pub extern "system" fn Java_dev_arkbuilders_core_BTreeMapIterator_hasNext(
+    _env: JNIEnv<'_>,
+    _class: JClass,
+    btreemap_ptr: jlong,
+) -> jboolean {
+    let iter = BTreeMapIterator::from_jlong(btreemap_ptr);
+    iter.has_next() as jboolean
+}
+
+#[no_mangle]
+pub extern "system" fn Java_dev_arkbuilders_core_BTreeMapIterator_next(
+    mut env: JNIEnv<'_>,
+    _class: JClass,
+    btreemap_ptr: jlong,
+) -> jobject {
+    let iter = BTreeMapIterator::from_jlong(btreemap_ptr);
+    let (key, value) = iter.native_next().unwrap();
+    let key = env.new_string(key).unwrap();
+    let value = env.new_string(value).unwrap();
+    let pair = env
+        .new_object(
+            "java/util/AbstractMap$SimpleImmutableEntry",
+            "(Ljava/lang/Object;Ljava/lang/Object;)V",
+            &[JValue::Object(&key), JValue::Object(&value)],
+        )
+        .expect("Failed to create new Java Map object");
+    pair.as_raw()
+}
+
+#[no_mangle]
+pub extern "system" fn Java_dev_arkbuilders_core_BTreeMapIterator_drop(
+    _env: JNIEnv<'_>,
+    _class: JClass,
+    btreemap_ptr: jlong,
+) {
+    unsafe {
+        let _ = Box::from_raw(
+            btreemap_ptr as *mut BTreeMapIterator<String, String>,
+        );
+    }
+}

fs-storage/src/jni/file_storage.rs

@@ -14,15 +14,15 @@ use jni::objects::{JClass, JObject, JString, JValue};
 // This is just a pointer. We'll be returning it from our function. We
 // can't return one of the objects with lifetime information because the
 // lifetime checker won't let us.
-use jni::sys::{jlong, jobject};
+use jni::sys::{jlong, jobject, jstring};
 use jnix::{IntoJava, JnixEnv};
 
 use crate::base_storage::BaseStorage;
 
 use crate::file_storage::FileStorage;
 
 impl FileStorage<String, String> {
-    fn from_jlong<'a>(value: jlong) -> &'a mut Self {
+    pub fn from_jlong<'a>(value: jlong) -> &'a mut Self {
         unsafe { &mut *(value as *mut FileStorage<String, String>) }
     }
 }
@@ -172,6 +172,27 @@ pub extern "system" fn Java_dev_arkbuilders_core_FileStorage_readFS(
     linked_hash_map.as_raw()
 }
 
+#[no_mangle]
+pub extern "system" fn Java_dev_arkbuilders_core_FileStorage_get<'local>(
+    mut env: JNIEnv<'local>,
+    _class: JClass<'local>,
+    id: JString<'local>,
+    file_storage_ptr: jlong,
+) -> jstring {
+    let id: String = env
+        .get_string(&id)
+        .expect("Failed to get string from JNI")
+        .into();
+    let file_storage = FileStorage::from_jlong(file_storage_ptr);
+
+    match file_storage.get(&id) {
+        Some(value) => env
+            .new_string(value)
+            .expect("Failed to create new string")
+            .into_raw(),
+        None => JObject::null().into_raw(),
+    }
+}
 #[no_mangle]
 pub extern "system" fn Java_dev_arkbuilders_core_FileStorage_writeFS(
     mut env: JNIEnv<'_>,

fs-storage/src/jni/mod.rs

@@ -1 +1,2 @@
+pub mod btreemap_iter;
 pub mod file_storage;

fs-storage/src/lib.rs

@@ -1,8 +1,10 @@
 pub mod base_storage;
+pub mod btreemap_iter;
 pub mod file_storage;
 #[cfg(feature = "jni-bindings")]
 pub mod jni;
 pub mod monoid;
+
 mod utils;
 pub const ARK_FOLDER: &str = ".ark";
 

java/.gitignore

@@ -3,3 +3,4 @@
 
 # Ignore Gradle build output directory
 build
+*.log

java/lib/src/main/java/dev/arkbuilders/core/BTreeMapIterator.java

@@ -0,0 +1,46 @@
+package dev.arkbuilders.core;
+
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NoSuchElementException;
+
+/**
+ * Represents an iterator over a BTreeMap.
+ */
+public class BTreeMapIterator implements Iterator<Map.Entry<String, String>>, AutoCloseable {
+    private long btreemap_ptr;
+
+    private static native long create(long file_storage_ptr);
+
+    private static native boolean hasNext(long btreemap_ptr);
+
+    private static native Object next(long btreemap_ptr);
+
+    private static native void drop(long btreemap_ptr);
+
+    BTreeMapIterator(long file_storage_ptr) {
+        this.btreemap_ptr = create(file_storage_ptr);
+    }
+
+    @Override
+    public boolean hasNext() {
+        return hasNext(this.btreemap_ptr);
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public Map.Entry<String, String> next() {
+        if (!hasNext()) {
+            throw new NoSuchElementException();
+        }
+        return (Map.Entry<String, String>) next(this.btreemap_ptr);
+    }
+
+    @Override
+    public void close() {
+        if (this.btreemap_ptr != 0) {
+            drop(this.btreemap_ptr);
+            this.btreemap_ptr = 0;
+        }
+    }
+}

java/lib/src/main/java/dev/arkbuilders/core/FileStorage.java

@@ -45,6 +45,8 @@ public enum SyncStatus {
 
     private static native Object readFS(long file_storage_ptr);
 
+    private static native String get(String id, long file_storage_ptr);
+
     private static native void writeFS(long file_storage_ptr);
 
     private static native void erase(long file_storage_ptr);
@@ -107,6 +109,16 @@ public Object readFS() {
         return readFS(this.fileStoragePtr);
     }
 
+    /**
+     * Get the value of a key from the internal mapping.
+     *
+     * @param id The key.
+     * @return The value.
+     */
+    public String get(String id) {
+        return get(id, this.fileStoragePtr);
+    }
+
     /**
      * Write the data to file.
      * 
@@ -133,4 +145,13 @@ public void erase() {
     public void merge(FileStorage other) {
         merge(this.fileStoragePtr, other.fileStoragePtr);
     }
+
+    /**
+     * Create a new iterator for the BTreeMap
+     * 
+     * @return The iterator
+     */
+    public BTreeMapIterator iterator() {
+        return new BTreeMapIterator(this.fileStoragePtr);
+    }
 }

java/lib/src/test/java/dev/arkbuilders/core/FileStorageTest.java

@@ -6,6 +6,7 @@
 import java.io.File;
 import java.nio.file.Path;
 import java.util.LinkedHashMap;
+import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.*;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -114,6 +115,42 @@ public void testFileStorageMainScenario() {
         assertFalse(file.exists());
     }
 
+    @Test
+    public void testFileStorageGet() {
+        Path storagePath = tempDir.resolve("test.txt");
+        FileStorage fileStorage = new FileStorage("test", storagePath.toString());
+
+        fileStorage.set("key", "value");
+        fileStorage.set("key", "value1");
+        fileStorage.set("key1", "value");
+
+        assertEquals("value1", fileStorage.get("key"));
+        assertEquals("value", fileStorage.get("key1"));
+    }
+
+    @Test
+    public void testBTreeMapIterator(){
+        Path storagePath = tempDir.resolve("test.txt");
+        FileStorage fileStorage = new FileStorage("test", storagePath.toString());
+
+        fileStorage.set("key", "value");
+        fileStorage.set("key", "value1");
+        fileStorage.set("key1", "value");
+
+        fileStorage.writeFS();
+
+        @SuppressWarnings("unchecked")
+        LinkedHashMap<String, String> data = (LinkedHashMap<String, String>) fileStorage.readFS();
+
+        BTreeMapIterator bTreeMapIterator = fileStorage.iterator();
+        Map<String, String> iteratorData = new LinkedHashMap<>();
+        while(bTreeMapIterator.hasNext()){
+            Map.Entry<String, String> entry = bTreeMapIterator.next();
+            iteratorData.put(entry.getKey(), entry.getValue());
+        }
+        assertEquals(data, iteratorData);
+    }
+
     @Test
     public void testRemoveException() {
         Path storagePath = tempDir.resolve("test.txt");