Get protobuf message passing working

java/patchbukkit/src/main/java/org/patchbukkit/bridge/BridgeUtils.java

@@ -0,0 +1,13 @@
+package org.patchbukkit.bridge;
+
+import java.util.UUID;
+
+public class BridgeUtils {
+    public static UUID convertUuid(patchbukkit.common.UUID uuid) {
+        return UUID.fromString(uuid.getValue());
+    }
+
+    public static patchbukkit.common.UUID convertUuid(UUID uuid) {
+        return patchbukkit.common.UUID.newBuilder().setValue(uuid.toString()).build();
+    }
+}

java/patchbukkit/src/main/java/org/patchbukkit/bridge/NativePatchBukkit.java

@@ -17,8 +17,6 @@ public class NativePatchBukkit {
     private static MethodHandle sendMessageNative;
     private static MethodHandle registerEventNative;
     private static MethodHandle callEventNative;
-    private static MethodHandle getAbilitiesNative;
-    private static MethodHandle setAbilitiesNative;
     private static MethodHandle getLocationNative;
     private static MethodHandle getWorldNative;
     private static MethodHandle freeStringNative;
@@ -188,8 +186,6 @@ public class NativePatchBukkit {
         long sendMessageAddr,
         long registerEventAddr,
         long callEventAddr,
-        long getAbilitiesAddr,
-        long setAbilitiesAddr,
         long getLocationAddr,
         long freeStringAddr,
         long getWorldAddr,
@@ -225,25 +221,6 @@ public class NativePatchBukkit {
             )
         );
 
-        // bool rust_get_abilities(const char* uuid, AbilitiesFFI* out)
-        getAbilitiesNative = LINKER.downcallHandle(
-            MemorySegment.ofAddress(getAbilitiesAddr),
-            FunctionDescriptor.of(
-                ValueLayout.JAVA_BOOLEAN, // return type
-                ValueLayout.ADDRESS, // uuid string
-                ValueLayout.ADDRESS // out pointer to AbilitiesFFI
-            )
-        );
-
-        // bool rust_set_abilities(const char* uuid, AbilitiesFFI* abilities)
-        setAbilitiesNative = LINKER.downcallHandle(
-            MemorySegment.ofAddress(setAbilitiesAddr),
-            FunctionDescriptor.of(
-                ValueLayout.JAVA_BOOLEAN, // return type
-                ValueLayout.ADDRESS, // uuid string
-                ValueLayout.ADDRESS // pointer to AbilitiesFFI
-            )
-        );
 
         // bool rust_get_location(const char* uuid, Vec3FFI* out)
         getLocationNative = LINKER.downcallHandle(
@@ -339,67 +316,6 @@ public class NativePatchBukkit {
         }
     }
 
-    /**
-     * Get a player's abilities by UUID.
-     *
-     * @param uuid The player's UUID
-     * @return The player's abilities, or null if player not found
-     */
-    public static Abilities getAbilities(UUID uuid) {
-        try (Arena arena = Arena.ofConfined()) {
-            MemorySegment uuidStr = arena.allocateFrom(uuid.toString());
-            MemorySegment outStruct = arena.allocate(ABILITIES_LAYOUT);
-
-            boolean success = (boolean) getAbilitiesNative.invokeExact(
-                uuidStr,
-                outStruct
-            );
-
-            if (!success) {
-                return null;
-            }
-
-            return new Abilities(
-                (boolean) INVULNERABLE.get(outStruct, 0L),
-                (boolean) FLYING.get(outStruct, 0L),
-                (boolean) ALLOW_FLYING.get(outStruct, 0L),
-                (boolean) CREATIVE.get(outStruct, 0L),
-                (boolean) ALLOW_MODIFY_WORLD.get(outStruct, 0L),
-                (float) FLY_SPEED.get(outStruct, 0L),
-                (float) WALK_SPEED.get(outStruct, 0L)
-            );
-        } catch (Throwable t) {
-            throw new RuntimeException("Failed to call native getAbilities", t);
-        }
-    }
-
-    public static boolean setAbilities(UUID uuid, Abilities abilities) {
-        try (Arena arena = Arena.ofConfined()) {
-            MemorySegment uuidStr = arena.allocateFrom(uuid.toString());
-            MemorySegment abilitiesStruct = arena.allocate(ABILITIES_LAYOUT);
-
-            // Populate the struct with values from the Abilities record
-            INVULNERABLE.set(abilitiesStruct, 0L, abilities.invulnerable());
-            FLYING.set(abilitiesStruct, 0L, abilities.flying());
-            ALLOW_FLYING.set(abilitiesStruct, 0L, abilities.allowFlying());
-            CREATIVE.set(abilitiesStruct, 0L, abilities.creative());
-            ALLOW_MODIFY_WORLD.set(
-                abilitiesStruct,
-                0L,
-                abilities.allowModifyWorld()
-            );
-            FLY_SPEED.set(abilitiesStruct, 0L, abilities.flySpeed());
-            WALK_SPEED.set(abilitiesStruct, 0L, abilities.walkSpeed());
-
-            return (boolean) setAbilitiesNative.invokeExact(
-                uuidStr,
-                abilitiesStruct
-            );
-        } catch (Throwable t) {
-            throw new RuntimeException("Failed to call native setAbilities", t);
-        }
-    }
-
     /**
      * Get a player's location by UUID.
      *

java/patchbukkit/src/main/java/org/patchbukkit/entity/PatchBukkitPlayer.java

@@ -62,6 +62,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Range;
 import org.jetbrains.annotations.Unmodifiable;
 import org.jetbrains.annotations.UnmodifiableView;
+import org.patchbukkit.bridge.BridgeUtils;
 import org.patchbukkit.bridge.NativePatchBukkit;
 import org.patchbukkit.registry.PatchBukkitSound;
 
@@ -78,6 +79,8 @@ import net.kyori.adventure.bossbar.BossBar;
 import net.kyori.adventure.text.Component;
 import net.kyori.adventure.util.TriState;
 import net.md_5.bungee.api.chat.BaseComponent;
+import patchbukkit.bridge.NativeBridgeFfi;
+import patchbukkit.common.SetAbilitiesRequest;
 
 @SuppressWarnings({ "deprecation", "removal" })
 public class PatchBukkitPlayer
@@ -1177,16 +1180,15 @@ public class PatchBukkitPlayer
 
     @Override
     public boolean getAllowFlight() {
-        return NativePatchBukkit.getAbilities(this.uuid).allowFlying();
+        return NativeBridgeFfi.getAbilities(BridgeUtils.convertUuid(this.uuid)).getAllowFlying();
     }
 
     @Override
     public void setAllowFlight(boolean flight) {
-        var abilities = NativePatchBukkit.getAbilities(this.uuid).withAllowFlying(flight);
+        var abilities = NativeBridgeFfi.getAbilities(BridgeUtils.convertUuid(this.uuid)).toBuilder();
-        if (flight == false) {
+        abilities.setAllowFlying(flight);
-            abilities = abilities.withFlying(false);
+        if (abilities.getFlying()) abilities.setFlying(false);
-        }
+        NativeBridgeFfi.setAbilities(SetAbilitiesRequest.newBuilder().setAbilities(abilities).setUuid(BridgeUtils.convertUuid(this.uuid)).build());
-        NativePatchBukkit.setAbilities(this.uuid, abilities);
     }
 
     @Override
@@ -1279,8 +1281,9 @@ public class PatchBukkitPlayer
 
     @Override
     public void setFlying(boolean value) {
-        var abilities = NativePatchBukkit.getAbilities(this.getUniqueId()).withFlying(value);
+        var playerUuid = BridgeUtils.convertUuid(this.getUniqueId());
-        NativePatchBukkit.setAbilities(this.getUniqueId(), abilities);
+        var abilities = NativeBridgeFfi.getAbilities(playerUuid).toBuilder().setFlying(value).build();
+        NativeBridgeFfi.setAbilities(SetAbilitiesRequest.newBuilder().setAbilities(abilities).setUuid(playerUuid).build());
     }
 
     @Override

java/patchbukkit/src/main/java/org/patchbukkit/registry/PatchBukkitTag.java

@@ -8,6 +8,7 @@ import org.bukkit.Keyed;
 import org.bukkit.NamespacedKey;
 import org.bukkit.Registry;
 import org.jetbrains.annotations.Unmodifiable;
+import patchbukkit.bridge.NativeBridgeFfi;
 
 import java.util.*;
 import java.util.stream.Collectors;

rust/build/java.rs

@@ -88,7 +88,9 @@ pub fn setup_java(base: PathBuf) {
         ("org.eclipse.sisu", "org.eclipse.sisu.inject", "0.9.0.M2"),
         ("org.apache.commons", "commons-lang3", "3.12.0"),
         ("net.kyori", "option", "1.1.0"),
+        // These are deps of patchbukkit itself
         ("io.papermc.paper", "paper-api", "1.21.11-R0.1-SNAPSHOT"),
+        ("com.google.protobuf", "protobuf-java", "4.33.5"),
     ];
 
     let jvm = JvmBuilder::new()

rust/build/protobufs.rs

@@ -1,10 +1,30 @@
 use glob::glob;
 use prost_build::{Config, Service, ServiceGenerator};
+use std::io::Write;
 use std::path::PathBuf;
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+struct ServiceInfo {
+    java_class: String,
+    methods: Vec<MethodInfo>,
+    module_path: String,
+}
+
+#[derive(Clone)]
+struct MethodInfo {
+    fn_name: String,
+}
+
+#[derive(Default)]
+struct SharedState {
+    services: Vec<ServiceInfo>,
+}
 
 pub struct FfiServiceGenerator {
     impl_module: String,
     proto_module: String,
+    state: Arc<Mutex<SharedState>>,
 }
 
 impl FfiServiceGenerator {
@@ -12,6 +32,7 @@ impl FfiServiceGenerator {
         Self {
             impl_module: impl_module.into(),
             proto_module: proto_module.into(),
+            state: Arc::new(Mutex::new(SharedState::default())),
         }
     }
 
@@ -34,14 +55,42 @@ impl FfiServiceGenerator {
 
         format!("{}::{module_path}::{type_name}", self.proto_module)
     }
+
+    fn proto_package_to_java_class(&self, package: &str, service_name: &str) -> String {
+        format!("{}.{}Ffi", package, service_name)
+    }
+
+    fn package_to_module_path(&self, package: &str) -> String {
+        package
+            .split('.')
+            .map(|s| to_snake_case(s))
+            .collect::<Vec<_>>()
+            .join("::")
+    }
+
+    fn get_state(&self) -> Arc<Mutex<SharedState>> {
+        Arc::clone(&self.state)
+    }
 }
 
 impl ServiceGenerator for FfiServiceGenerator {
     fn generate(&mut self, service: Service, buf: &mut String) {
+        let module_path = self.package_to_module_path(&service.package);
+
+        let mut service_info = ServiceInfo {
+            java_class: self.proto_package_to_java_class(&service.package, &service.name),
+            methods: Vec::new(),
+            module_path,
+        };
+
         for method in &service.methods {
             let fn_name = format!("ffi_{}_{}", to_snake_case(&service.name), &method.name);
             let input_type = self.proto_type_to_rust(&method.input_proto_type);
 
+            service_info.methods.push(MethodInfo {
+                fn_name: fn_name.clone(),
+            });
+
             buf.push_str(&format!(
                 r#"
 #[unsafe(no_mangle)]
@@ -70,6 +119,8 @@ pub unsafe extern "C" fn {fn_name}(
                 self.impl_module
             ));
         }
+
+        self.state.lock().unwrap().services.push(service_info);
     }
 }
 
@@ -104,10 +155,77 @@ pub fn setup_protobufs(base: PathBuf) {
         .filter_map(Result::ok)
         .collect();
 
+    let generator = FfiServiceGenerator::new("crate::java::native_callbacks", "crate::proto");
+    let state = generator.get_state();
+
     let mut config = Config::new();
-    config.service_generator(Box::new(FfiServiceGenerator::new(
+    config.service_generator(Box::new(generator));
-        "crate::java::native_callbacks",
-        "crate::proto",
-    )));
     config.compile_protos(&paths, &[proto_path]).unwrap();
+
+    // Generate the shared initialization file
+    let out_dir = std::env::var("OUT_DIR").unwrap();
+    let init_path = PathBuf::from(&out_dir).join("ffi_init.rs");
+
+    let state = state.lock().unwrap();
+    let mut file = std::fs::File::create(&init_path).unwrap();
+
+    writeln!(
+        file,
+        r#"
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn ffi_free_bytes(ptr: *mut u8, len: usize) {{
+    if !ptr.is_null() && len > 0 {{
+        unsafe {{
+            drop(Vec::from_raw_parts(ptr, len, len));
+        }}
+    }}
+}}
+
+pub fn initialize_ffi_callbacks(jvm: &j4rs::Jvm) -> anyhow::Result<()> {{
+    use j4rs::InvocationArg;"#
+    )
+    .unwrap();
+
+    for service in &state.services {
+        writeln!(
+            file,
+            r#"
+    // Initialize {}
+    jvm.invoke_static(
+        "{}",
+        "init",
+        &["#,
+            service.java_class, service.java_class
+        )
+        .unwrap();
+
+        for method in &service.methods {
+            writeln!(file,
+                "            InvocationArg::try_from(crate::proto::{}::{} as *const () as i64)?.into_primitive()?,",
+                service.module_path, method.fn_name
+            ).unwrap();
+        }
+
+        writeln!(
+            file,
+            r#"        ],
+    )?;
+
+    jvm.invoke_static(
+        "{}",
+        "initFree",
+        &[InvocationArg::try_from(ffi_free_bytes as *const () as i64)?.into_primitive()?],
+    )?;"#,
+            service.java_class
+        )
+        .unwrap();
+    }
+
+    writeln!(
+        file,
+        r#"
+    Ok(())
+}}"#
+    )
+    .unwrap();
 }

rust/src/java/native_callbacks/mod.rs

@@ -5,7 +5,7 @@ use j4rs::{InvocationArg, Jvm};
 use pumpkin::plugin::Context;
 use tokio::sync::mpsc;
 
-use crate::{java::jvm::commands::JvmCommand, proto};
+use crate::{java::jvm::commands::JvmCommand, proto::initialize_ffi_callbacks};
 
 mod abilities;
 pub use abilities::*;
@@ -68,10 +68,6 @@ pub fn initialize_callbacks(jvm: &Jvm) -> Result<()> {
     let send_message_addr = message::rust_send_message as *const () as i64;
     let register_event_addr = events::rust_register_event as *const () as i64;
     let call_event_addr = events::rust_call_event as *const () as i64;
-    let get_abilities_addr =
-        proto::patchbukkit::bridge::ffi_native_bridge_get_abilities as *const () as i64;
-    let set_abilities_addr =
-        proto::patchbukkit::bridge::ffi_native_bridge_set_abilities as *const () as i64;
     let get_location_addr = location::rust_get_location as *const () as i64;
     let free_string_addr = memory::rust_free_string as *const () as i64;
     let get_world_addr = world::rust_get_world as *const () as i64;
@@ -87,8 +83,6 @@ pub fn initialize_callbacks(jvm: &Jvm) -> Result<()> {
             InvocationArg::try_from(send_message_addr)?.into_primitive()?,
             InvocationArg::try_from(register_event_addr)?.into_primitive()?,
             InvocationArg::try_from(call_event_addr)?.into_primitive()?,
-            InvocationArg::try_from(get_abilities_addr)?.into_primitive()?,
-            InvocationArg::try_from(set_abilities_addr)?.into_primitive()?,
             InvocationArg::try_from(get_location_addr)?.into_primitive()?,
             InvocationArg::try_from(free_string_addr)?.into_primitive()?,
             InvocationArg::try_from(get_world_addr)?.into_primitive()?,
@@ -98,5 +92,7 @@ pub fn initialize_callbacks(jvm: &Jvm) -> Result<()> {
         ],
     )?;
 
+    initialize_ffi_callbacks(jvm)?;
+
     Ok(())
 }

rust/src/proto/mod.rs

@@ -11,3 +11,5 @@ pub mod patchbukkit {
         include!(concat!(env!("OUT_DIR"), "/patchbukkit.events.rs"));
     }
 }
+
+include!(concat!(env!("OUT_DIR"), "/ffi_init.rs"));