Dynamic custom jet loader

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 resolver = "2"
-members = ["cli", "core", "service"]
+members = ["cli", "core", "jet_plugins", "plugin_test", "service"]
 
 [workspace.dependencies]
 hal-simplicity = "0.2.0"

core/Cargo.toml

@@ -13,3 +13,4 @@ hex-literal = "1.1.0"
 hal-simplicity = { workspace = true }
 thiserror = { workspace = true }
 hex = { workspace = true }
+dlopen2 = "0.8.2"

core/src/jets/custom_jet.rs

@@ -0,0 +1,305 @@
+use crate::jets::{environments::ElementsUnchainedEnv, jet_dyn::CustomJetApi};
+use dlopen2::wrapper::Container;
+use hal_simplicity::simplicity::jet::Jet;
+use std::sync::LazyLock;
+
+pub(crate) static JET_DLL: LazyLock<Option<Container<CustomJetApi>>> = LazyLock::new(|| {
+    std::env::var("JET_DLL_PATH")
+        .ok()
+        .and_then(|path| unsafe { Container::<CustomJetApi>::load(&path) }.ok())
+});
+/// `simplicity_unchained_core::jets::jet_dyn::decode` copies input bits into a separate buffer to transfer them via FFI.
+/// This constant controls how many bits will be read.
+const MAX_JET_BIT_LEN: u32 = 30;
+
+#[repr(C)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct CustomJet {
+    pub index: usize,
+}
+
+impl CustomJet {
+    pub fn all() -> &'static [Self] {
+        JET_DLL.as_ref().expect("DLL is not loaded").all_jets()
+    }
+
+    pub unsafe fn to_base_jet<T>(&self) -> Option<T> {
+        unsafe {
+            JET_DLL
+                .as_ref()
+                .expect("DLL is not loaded")
+                .to_base_jet(*self)
+        }
+    }
+
+    pub unsafe fn from_base_jet<T>(jet: &T) -> Self {
+        unsafe {
+            JET_DLL
+                .as_ref()
+                .expect("DLL is not loaded")
+                .from_base_jet(jet)
+        }
+    }
+}
+
+impl std::str::FromStr for CustomJet {
+    type Err = hal_simplicity::simplicity::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match JET_DLL.as_ref().expect("DLL is not loaded").from_str(s) {
+            Ok(_jet) => Ok(_jet),
+            Err(err) => Err(err),
+        }
+    }
+}
+
+impl PartialOrd for CustomJet {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .partial_cmp(*self, *other)
+    }
+}
+
+impl Ord for CustomJet {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .cmp(*self, *other)
+    }
+}
+
+impl core::hash::Hash for CustomJet {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .hash(*self, state);
+    }
+}
+
+impl std::fmt::Debug for CustomJet {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .debug_fmt(*self, f)
+    }
+}
+
+impl std::fmt::Display for CustomJet {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .display_fmt(*self, f)
+    }
+}
+
+impl Jet for CustomJet {
+    type Environment = ElementsUnchainedEnv;
+    type CJetEnvironment = ElementsUnchainedEnv;
+
+    fn c_jet_env(env: &Self::Environment) -> &Self::CJetEnvironment {
+        env
+    }
+
+    fn cmr(&self) -> hal_simplicity::simplicity::Cmr {
+        JET_DLL.as_ref().expect("DLL is not loaded").cmr(*self)
+    }
+
+    fn source_ty(&self) -> hal_simplicity::simplicity::jet::type_name::TypeName {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .source_ty(*self)
+    }
+
+    fn target_ty(&self) -> hal_simplicity::simplicity::jet::type_name::TypeName {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .target_ty(*self)
+    }
+
+    fn encode<W: std::io::Write>(
+        &self,
+        w: &mut hal_simplicity::simplicity::BitWriter<W>,
+    ) -> std::io::Result<usize> {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .encode(*self, w)
+    }
+
+    /// # Safety
+    ///
+    /// Due to the lack of a `Clone` bound on `I`, the underlying implementation uses
+    /// `ptr::read` to create bitwise copies of the iterator. This is unsafe and may cause
+    /// undefined behavior if `I` contains types that manage unique resources.
+    /// This works correctly for common slice-based iterators like `Copied<Iter<u8>>`.
+    ///
+    /// See <https://github.com/BlockstreamResearch/rust-simplicity/issues/342> for details.
+    fn decode<I: Iterator<Item = u8>>(
+        bits: &mut hal_simplicity::simplicity::BitIter<I>,
+    ) -> Result<Self, hal_simplicity::simplicity::decode::Error> {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .decode(bits, MAX_JET_BIT_LEN)
+    }
+
+    fn c_jet_ptr(
+        &self,
+    ) -> &dyn Fn(
+        &mut hal_simplicity::simplicity::ffi::CFrameItem,
+        hal_simplicity::simplicity::ffi::CFrameItem,
+        &Self::CJetEnvironment,
+    ) -> bool {
+        JET_DLL
+            .as_ref()
+            .expect("DLL is not loaded")
+            .c_jet_ptr(*self)
+    }
+
+    fn cost(&self) -> hal_simplicity::simplicity::Cost {
+        JET_DLL.as_ref().expect("DLL is not loaded").cost(*self)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use core::panic;
+    use std::{
+        hash::{Hash, Hasher},
+        str::FromStr,
+    };
+
+    use hal_simplicity::simplicity::{BitIter, BitWriter};
+
+    use super::*;
+
+    #[test]
+    fn test_dll_from_str() {
+        panic!("Rebuild `plugin_test` using code from `lib1.rs`");
+        unsafe {
+            std::env::set_var("JET_DLL_PATH", "../target/debug/libplugin_test.so");
+        }
+        let jet = CustomJet::from_str("custom_jet_1");
+
+        assert!(jet.is_ok());
+        assert_eq!(jet.unwrap().to_string(), "custom_jet_1")
+    }
+
+    #[test]
+    fn test_dll_hash() {
+        panic!("Rebuild `plugin_test` using code from `lib1.rs`");
+        unsafe {
+            std::env::set_var("JET_DLL_PATH", "../target/debug/libplugin_test.so");
+        }
+        let jet = CustomJet::from_str("custom_jet_1").expect("Failed to load jet");
+
+        let mut hasher = std::collections::hash_map::DefaultHasher::new();
+        jet.hash(&mut hasher);
+
+        let _ = hasher.finish();
+    }
+
+    #[test]
+    fn test_dll_encode() {
+        panic!("Rebuild `plugin_test` using code from `lib1.rs`");
+        unsafe {
+            std::env::set_var("JET_DLL_PATH", "../target/debug/libplugin_test.so");
+        }
+        let custom_jet = CustomJet::from_str("custom_jet_1").expect("Failed to load jet");
+        let elements_jet = CustomJet::from_str("input_asset").expect("Failed to load jet");
+
+        let mut buffer = Vec::new();
+        let mut w = BitWriter::new(&mut buffer);
+
+        let res1 = custom_jet.encode(&mut w);
+        let res2 = elements_jet.encode(&mut w);
+
+        assert!(res1.is_ok());
+        assert!(res2.is_ok());
+
+        assert_eq!(res1.unwrap(), 20);
+        assert_eq!(res2.unwrap(), 19);
+
+        w.flush_all().unwrap();
+
+        let mut bit_iter = BitIter::from(buffer);
+
+        let mut read_code = |read_bits: u8| -> Option<u64> {
+            let mut res = 0;
+            let mut cursor = 1 << (read_bits - 1);
+
+            for _ in 0..read_bits {
+                match bit_iter.next() {
+                    None => return None,
+                    Some(true) => res |= cursor,
+                    Some(false) => {}
+                }
+                cursor >>= 1;
+            }
+            Some(res)
+        };
+
+        let custom_jet_code = read_code(20);
+        let elements_jet_code = read_code(19);
+
+        assert!(custom_jet_code.is_some());
+        assert!(elements_jet_code.is_some());
+
+        assert_eq!(custom_jet_code.unwrap(), 1047454);
+        assert_eq!(elements_jet_code.unwrap(), 462369);
+    }
+
+    #[test]
+    fn test_dll_decode() {
+        panic!("Rebuild `plugin_test` using code from `lib1.rs`");
+        unsafe {
+            std::env::set_var("JET_DLL_PATH", "../target/debug/libplugin_test.so");
+        }
+        let custom_jet = CustomJet::from_str("custom_jet_1").expect("Failed to load jet");
+        let elements_jet = CustomJet::from_str("input_asset").expect("Failed to load jet");
+
+        let mut buffer = Vec::new();
+        let mut w = BitWriter::new(&mut buffer);
+
+        let res1 = custom_jet.encode(&mut w);
+        let res2 = elements_jet.encode(&mut w);
+
+        assert!(res1.is_ok());
+        assert!(res2.is_ok());
+
+        w.flush_all().unwrap();
+
+        let mut bit_iter = BitIter::from(buffer.as_ref());
+
+        let custom_decoded = CustomJet::decode(&mut bit_iter);
+        assert!(custom_decoded.is_ok());
+
+        assert_eq!(custom_decoded.unwrap().to_string(), "custom_jet_1");
+
+        let elements_decoded = CustomJet::decode(&mut bit_iter);
+        assert!(elements_decoded.is_ok());
+
+        assert_eq!(elements_decoded.unwrap().to_string(), "input_asset");
+    }
+
+    #[test]
+    fn test_dll_c_jet_ptr() {
+        panic!("Rebuild `plugin_test` using code from `lib1.rs`");
+        unsafe {
+            std::env::set_var("JET_DLL_PATH", "../target/debug/libplugin_test.so");
+        }
+        let custom_jet = CustomJet::from_str("custom_jet_1").expect("Failed to load jet");
+        let elements_jet = CustomJet::from_str("input_asset").expect("Failed to load jet");
+
+        let _ = custom_jet.c_jet_ptr();
+        let _ = elements_jet.c_jet_ptr();
+    }
+}