dyn-loader 0.1.0

//! # dyn-loader
//!
//! Dynamic library loader with dyn-fat-pointer-bridge for loading trait objects
//! from `.so`/`.dylib` plugin files.
//!
//! Architecture (adapted from my-agent's plugin-loader + safe-arc + rust-plugin-loader):
//!
//! - **DynLib**: Wraps `libloading::Library` with `Arc` for shared ownership.
//! - **AbiDynFatPtr**: ABI-stable representation of a Rust fat pointer
//!   (data ptr + vtable ptr), `#[repr(C)]` for C ABI compatibility.
//! - **AbiStableDynRef**: Fat pointer + retain/release function pointers,
//!   enabling safe cross-boundary Arc-like reference counting.
//! - **DynPlugin<T>**: Loaded plugin holding a `SafeArcDyn<T>` that can be
//!   dereferenced to `&T` to call trait methods on the loaded object.
//!
//! ## Usage
//!
//! ```ignore
//! // In the plugin .so:
//! #[no_mangle]
//! pub extern "C" fn core_ast_transform_entry() -> AbiStableDynRef {
//!     SafeArcDyn::from_arc(Arc::new(MyTransform) as Arc<dyn Transform>).into_abi()
//! }
//!
//! // In the host:
//! let plugin = DynPlugin::<dyn Transform>::load("libmy_transform.so", b"core_ast_transform_entry\0")?;
//! let transform: &dyn Transform = plugin.as_ref();
//! ```

use std::ffi::{c_char, c_void};
use std::marker::PhantomData;
use std::path::Path;
use std::sync::Arc;

use anyhow::{Context, Result};
use libloading::Library;

// ---------------------------------------------------------------------------
// DynLib — Arc-shared dynamic library
// ---------------------------------------------------------------------------

#[derive(Clone)]
pub struct DynLib {
    library: Arc<Library>,
    path: std::path::PathBuf,
}

impl DynLib {
    /// Load a dynamic library from the given path.
    ///
    /// # Safety
    ///
    /// The target file must be a valid dynamic library for the current process.
    pub unsafe fn load(path: &Path) -> Result<Self> {
        let library = unsafe { Library::new(path) }
            .with_context(|| format!("failed to load dynamic library: {}", path.display()))?;
        Ok(Self {
            library: Arc::new(library),
            path: path.to_path_buf(),
        })
    }
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Load a symbol from the library.
    ///
    /// # Safety
    ///
    /// The caller must ensure `T` matches the actual exported symbol type.
    pub unsafe fn symbol<T: Copy>(&self, name: &[u8]) -> Result<T> {
        let sym = unsafe { self.library.get::<T>(name) }.with_context(|| {
            format!(
                "symbol '{}' not found in {}",
                display_symbol(name),
                self.path.display()
            )
        })?;
        Ok(*sym)
    }

    /// Try to load a symbol; returns `None` if not found.
    ///
    /// # Safety
    ///
    /// The caller must ensure `T` matches the actual exported symbol type.
    pub unsafe fn try_symbol<T: Copy>(&self, name: &[u8]) -> Option<T> {
        unsafe { self.library.get::<T>(name) }.ok().map(|s| *s)
    }
}

// ---------------------------------------------------------------------------
// AbiDynFatPtr — ABI-stable fat pointer
// ---------------------------------------------------------------------------

/// ABI-stable representation of a Rust dyn trait fat pointer.
/// Two words: data pointer + vtable pointer.
#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AbiDynFatPtr {
    pub data: *const c_void,
    pub vtable: *const c_void,
}

impl AbiDynFatPtr {
    pub const fn null() -> Self {
        Self {
            data: std::ptr::null(),
            vtable: std::ptr::null(),
        }
    }

    pub fn is_null(self) -> bool {
        self.data.is_null() || self.vtable.is_null()
    }
}

// SAFETY: AbiDynFatPtr is just two raw pointers; the actual thread-safety
// is governed by the enclosing SafeArcDyn<T> where T: Send + Sync.
unsafe impl Send for AbiDynFatPtr {}
unsafe impl Sync for AbiDynFatPtr {}

// ---------------------------------------------------------------------------
// AbiStableDynRef — fat pointer + retain/release for Arc-like semantics
// ---------------------------------------------------------------------------

pub type RetainFn = unsafe extern "C" fn(AbiDynFatPtr);
pub type ReleaseFn = unsafe extern "C" fn(AbiDynFatPtr);

/// ABI-stable reference to a dyn trait object, with retain/release for
/// safe reference counting across dynamic library boundaries.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct AbiStableDynRef {
    pub object: AbiDynFatPtr,
    pub retain: RetainFn,
    pub release: ReleaseFn,
}

impl AbiStableDynRef {
    pub const fn null() -> Self {
        Self {
            object: AbiDynFatPtr::null(),
            retain: retain_noop,
            release: release_noop,
        }
    }

    pub fn is_null(self) -> bool {
        self.object.is_null()
    }
}

// SAFETY: The retain/release functions manage the Arc ref-count;
// thread-safety follows T: Send + Sync.
unsafe impl Send for AbiStableDynRef {}
unsafe impl Sync for AbiStableDynRef {}

unsafe extern "C" fn retain_noop(_: AbiDynFatPtr) {}
unsafe extern "C" fn release_noop(_: AbiDynFatPtr) {}

// ---------------------------------------------------------------------------
// Pack / unpack fat pointers
// ---------------------------------------------------------------------------

/// Pack a `*const T` (where T: ?Sized) into an ABI-stable fat pointer.
///
/// # Safety
///
/// `ptr` must be a valid fat pointer (e.g., `*const dyn Trait`).
pub unsafe fn pack_fat_ptr<T: ?Sized>(ptr: *const T) -> AbiDynFatPtr {
    // Fat pointers are exactly 2 words: data + metadata (vtable for dyn Trait)
    unsafe { std::mem::transmute_copy(&ptr) }
}

/// Unpack an ABI-stable fat pointer back to `*const T`.
///
/// # Safety
///
/// `ptr` must have been created from a compatible `*const T`.
pub unsafe fn unpack_fat_ptr<T: ?Sized>(ptr: AbiDynFatPtr) -> *const T {
    unsafe { std::mem::transmute_copy(&ptr) }
}

// ---------------------------------------------------------------------------
// Arc retain/release for dyn trait objects
// ---------------------------------------------------------------------------

unsafe extern "C" fn retain_arc<T: ?Sized>(ptr: AbiDynFatPtr) {
    let raw: *const T = unsafe { unpack_fat_ptr(ptr) };
    unsafe { Arc::increment_strong_count(raw) };
}

unsafe extern "C" fn release_arc<T: ?Sized>(ptr: AbiDynFatPtr) {
    let raw: *const T = unsafe { unpack_fat_ptr(ptr) };
    unsafe { drop(Arc::from_raw(raw)) };
}

// ---------------------------------------------------------------------------
// SafeArcDyn — safe wrapper around AbiStableDynRef
// ---------------------------------------------------------------------------

/// A safe, cloneable, droppable reference to a dyn trait object loaded
/// from a dynamic library. Uses Arc-like retain/release for memory safety.
#[repr(transparent)]
pub struct SafeArcDyn<T: ?Sized> {
    raw: AbiStableDynRef,
    _marker: PhantomData<*const T>,
}

impl<T: ?Sized> SafeArcDyn<T> {
    /// Create from an `Arc<T>`. The Arc's reference count is managed
    /// via the retain/release function pointers.
    pub fn from_arc(value: Arc<T>) -> Self {
        let raw = Arc::into_raw(value);
        Self {
            raw: AbiStableDynRef {
                object: unsafe { pack_fat_ptr(raw) },
                retain: retain_arc::<T>,
                release: release_arc::<T>,
            },
            _marker: PhantomData,
        }
    }

    /// Get the ABI-stable representation (for exporting from a plugin).
    pub fn into_abi(self) -> AbiStableDynRef {
        let raw = self.raw;
        std::mem::forget(self); // Don't drop — caller takes ownership
        raw
    }

    /// Reconstruct from an ABI-stable representation (for loading in host).
    ///
    /// # Safety
    ///
    /// `raw` must have been created by `SafeArcDyn::<T>` or compatible code.
    pub unsafe fn from_abi(raw: AbiStableDynRef) -> Self {
        Self {
            raw,
            _marker: PhantomData,
        }
    }

    /// Get a reference to the trait object.
    ///
    /// # Safety
    ///
    /// The stored fat pointer must be valid for the lifetime of this reference.
    pub unsafe fn trait_ref(&self) -> &T {
        unsafe { &*unpack_fat_ptr::<T>(self.raw.object) }
    }
}

// SAFETY: SafeArcDyn<T> is Arc-like; safe to Send/Sync when T is.
unsafe impl<T: ?Sized + Send> Send for SafeArcDyn<T> {}
unsafe impl<T: ?Sized + Sync> Sync for SafeArcDyn<T> {}

impl<T: ?Sized> Clone for SafeArcDyn<T> {
    fn clone(&self) -> Self {
        unsafe { (self.raw.retain)(self.raw.object) };
        Self {
            raw: self.raw,
            _marker: PhantomData,
        }
    }
}

impl<T: ?Sized> Drop for SafeArcDyn<T> {
    fn drop(&mut self) {
        unsafe { (self.raw.release)(self.raw.object) };
    }
}

// ---------------------------------------------------------------------------
// DynPlugin — loaded plugin with trait object access
// ---------------------------------------------------------------------------

/// A loaded dynamic library plugin that exposes a trait object via
/// the dyn-fat-pointer-bridge pattern.
pub struct DynPlugin<T: ?Sized> {
    _lib: DynLib,
    plugin: SafeArcDyn<T>,
}

/// Type of the entry point function that plugins must export.
pub type PluginEntryPoint = unsafe extern "C" fn() -> AbiStableDynRef;

impl<T: ?Sized> DynPlugin<T> {
    /// Load a plugin from a dynamic library file.
    ///
    /// The library must export a function with the given symbol name
    /// that returns an `AbiStableDynRef` created via `SafeArcDyn::into_abi()`.
    ///
    /// # Safety
    ///
    /// - The target file must be a valid dynamic library for the current process.
    /// - The entry point must return a valid `AbiStableDynRef` for trait `T`.
    pub unsafe fn load(path: &Path, entry_symbol: &[u8]) -> Result<Self> {
        let lib = unsafe { DynLib::load(path) }?;
        unsafe { Self::from_lib(lib, entry_symbol) }
    }
    /// Load from an already-loaded `DynLib`.
    ///
    /// # Safety
    ///
    /// The entry point must return a valid `AbiStableDynRef` for trait `T`.
    pub unsafe fn from_lib(lib: DynLib, entry_symbol: &[u8]) -> Result<Self> {
        let entry: PluginEntryPoint = unsafe { lib.symbol(entry_symbol) }
            .with_context(|| format!("entry point '{}' not found", display_symbol(entry_symbol)))?;
        let abi_ref = unsafe { entry() };
        if abi_ref.is_null() {
            anyhow::bail!("entry point returned null AbiStableDynRef");
        }
        let plugin = unsafe { SafeArcDyn::<T>::from_abi(abi_ref) };
        Ok(Self { _lib: lib, plugin })
    }

    /// Get a reference to the loaded trait object.
    pub fn trait_ref(&self) -> &T {
        unsafe { self.plugin.trait_ref() }
    }

    /// Get a cloned SafeArcDyn (for sharing across threads).
    pub fn clone_handle(&self) -> SafeArcDyn<T> {
        self.plugin.clone()
    }
}

// SAFETY: DynPlugin<T> owns a DynLib (Arc<Library>) + SafeArcDyn<T>;
// safe to Send/Sync when T is.
unsafe impl<T: ?Sized + Send> Send for DynPlugin<T> {}
unsafe impl<T: ?Sized + Sync> Sync for DynPlugin<T> {}

// ---------------------------------------------------------------------------
// VTablePlugin — simpler vtable-based loading (Copy types only)
// ---------------------------------------------------------------------------

/// Load a Copy-sized vtable/descriptor struct from a dynamic library.
pub struct VTablePlugin<T: Copy> {
    _lib: DynLib,
    vtable: T,
}

impl<T: Copy> VTablePlugin<T> {
    /// Load a vtable struct from a dynamic library.
    ///
    /// # Safety
    ///
    /// - The target file must be a valid dynamic library.
    /// - The symbol must refer to a static vtable-compatible value of type `T`.
    pub unsafe fn load(path: &Path, symbol: &[u8]) -> Result<Self> {
        let lib = unsafe { DynLib::load(path) }?;
        unsafe { Self::from_lib(lib, symbol) }
    }

    /// # Safety
    ///
    /// The symbol must refer to a static vtable-compatible value of type `T`.
    pub unsafe fn from_lib(lib: DynLib, symbol: &[u8]) -> Result<Self> {
        let getter: unsafe extern "C" fn() -> *const T = unsafe { lib.symbol(symbol) }?;
        let vtable = unsafe { getter() };
        let vtable = unsafe { vtable.as_ref() }
            .copied()
            .ok_or_else(|| anyhow::anyhow!("vtable getter returned null"))?;
        Ok(Self { _lib: lib, vtable })
    }

    pub fn vtable(&self) -> &T {
        &self.vtable
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn display_symbol(symbol: &[u8]) -> String {
    let end = symbol.iter().position(|&b| b == 0).unwrap_or(symbol.len());
    String::from_utf8_lossy(&symbol[..end]).into_owned()
}

/// Quick check: does the file look like a plugin with the given entry symbol?
///
/// # Safety
///
/// Probes an arbitrary dynamic library for a specific exported symbol.
pub unsafe fn looks_like_plugin(path: &Path, entry_symbol: &[u8]) -> bool {
    match unsafe { DynLib::load(path) } {
        Ok(lib) => unsafe { lib.try_symbol::<PluginEntryPoint>(entry_symbol) }.is_some(),
        Err(_) => false,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    trait Demo: Send + Sync {
        fn value(&self) -> i32;
    }

    struct DemoValue(i32);

    impl Demo for DemoValue {
        fn value(&self) -> i32 {
            self.0
        }
    }

    #[test]
    fn safe_arc_dyn_round_trip() {
        let arc: Arc<dyn Demo> = Arc::new(DemoValue(42));
        let wrapped = SafeArcDyn::from_arc(arc);
        let abi = wrapped.into_abi();
        let restored = unsafe { SafeArcDyn::<dyn Demo>::from_abi(abi) };
        assert_eq!(unsafe { restored.trait_ref() }.value(), 42);
    }

    #[test]
    fn safe_arc_dyn_clone() {
        let arc: Arc<dyn Demo> = Arc::new(DemoValue(7));
        let wrapped = SafeArcDyn::from_arc(arc);
        let cloned = wrapped.clone();
        assert_eq!(unsafe { wrapped.trait_ref() }.value(), 7);
        assert_eq!(unsafe { cloned.trait_ref() }.value(), 7);
        drop(wrapped);
        assert_eq!(unsafe { cloned.trait_ref() }.value(), 7);
    }
}

// ---------------------------------------------------------------------------
// C++ math module VTable ABI — matches cpp/math_module/include/math_vtable.h
// ---------------------------------------------------------------------------

/// Opaque handle for a C++ MathSession
pub type MathSession = *mut c_void;

/// Descriptor for a generated math function (matches C++ GeneratedFunction)
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct GeneratedFunction {
    pub name: *const c_char,
    pub signature: *const c_char,
    pub arg_count: u32,
    pub id: u32,
}

/// Vtable struct matching C++ MathModuleVtable layout exactly
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct MathModuleVtable {
    // Session lifecycle
    pub create_session: unsafe extern "C" fn() -> MathSession,
    pub destroy_session: unsafe extern "C" fn(MathSession),

    // Math operations (lazy — only "generated" if called)
    pub add_i32: unsafe extern "C" fn(MathSession, i32, i32) -> i32,
    pub add_i64: unsafe extern "C" fn(MathSession, i64, i64) -> i64,
    pub add_f64: unsafe extern "C" fn(MathSession, f64, f64) -> f64,
    pub mul_i32: unsafe extern "C" fn(MathSession, i32, i32) -> i32,
    pub mul_f64: unsafe extern "C" fn(MathSession, f64, f64) -> f64,
    pub pi: unsafe extern "C" fn(MathSession) -> f64,
    pub tau: unsafe extern "C" fn(MathSession) -> f64,

    // Code-generation introspection
    pub generated_count: unsafe extern "C" fn(MathSession) -> u32,
    pub generated_at: unsafe extern "C" fn(MathSession, u32) -> GeneratedFunction,

    // Module info
    pub module_name: unsafe extern "C" fn() -> *const c_char,
    pub module_version: unsafe extern "C" fn() -> u32,
}

// SAFETY: MathModuleVtable contains only function pointers and is safe to Send/Sync
unsafe impl Send for MathModuleVtable {}
unsafe impl Sync for MathModuleVtable {}

#[cfg(test)]
mod cpp_math_tests {
    use super::*;
    use std::ffi::CStr;
    use std::path::PathBuf;

    fn math_module_path() -> PathBuf {
        // Relative from crate root (rust/crates/dyn-loader) to cpp build output
        let mut p = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        p.push("../../cpp/build/math_module");
        p.push("libmath_module.so");
        p
    }

    #[test]
    fn load_cpp_math_module_via_vtable() {
        let path = math_module_path();
        if !path.exists() {
            eprintln!("SKIP: {} not found — build cpp/ first", path.display());
            return;
        }

        let plugin =
            unsafe { VTablePlugin::<MathModuleVtable>::load(&path, b"math_module_get_vtable\0") }
                .expect("failed to load math_module");
        let vt = plugin.vtable();

        // Module info
        unsafe {
            let name = CStr::from_ptr((vt.module_name)());
            assert_eq!(name.to_str().unwrap(), "core-ast-math");
            assert_eq!((vt.module_version)(), 1);
        }

        // Create session
        let session = unsafe { (vt.create_session)() };
        assert!(!session.is_null());

        // No functions generated yet
        assert_eq!(unsafe { (vt.generated_count)(session) }, 0);

        // Call add_i32 — marks it as "used"
        let result = unsafe { (vt.add_i32)(session, 10, 20) };
        assert_eq!(result, 30);
        assert_eq!(unsafe { (vt.generated_count)(session) }, 1);

        // Call mul_f64 — marks it as "used"
        let result = unsafe { (vt.mul_f64)(session, 3.0, 7.0) };
        assert!((result - 21.0).abs() < 1e-10);
        assert_eq!(unsafe { (vt.generated_count)(session) }, 2);

        // Call pi
        let pi_val = unsafe { (vt.pi)(session) };
        assert!((pi_val - std::f64::consts::PI).abs() < 1e-10);
        assert_eq!(unsafe { (vt.generated_count)(session) }, 3);

        // Introspect generated functions
        let func0 = unsafe { (vt.generated_at)(session, 0) };
        let func0_name = unsafe { CStr::from_ptr(func0.name) }.to_str().unwrap();
        assert_eq!(func0_name, "add_i32");

        let func1 = unsafe { (vt.generated_at)(session, 1) };
        let func1_name = unsafe { CStr::from_ptr(func1.name) }.to_str().unwrap();
        assert_eq!(func1_name, "mul_f64");

        let func2 = unsafe { (vt.generated_at)(session, 2) };
        let func2_name = unsafe { CStr::from_ptr(func2.name) }.to_str().unwrap();
        assert_eq!(func2_name, "pi");

        // Call add_i32 again — should NOT add duplicate
        let _ = unsafe { (vt.add_i32)(session, 1, 2) };
        assert_eq!(unsafe { (vt.generated_count)(session) }, 3);

        // Destroy session
        unsafe { (vt.destroy_session)(session) };
    }
}