jyafn 0.3.1

Computational graphs for Data Science that compile to machine code
Documentation
//! A resource that provides its functionality via extensions (i.e., external code).

use serde_derive::{Deserialize, Serialize};
use serde_with::{serde_as, DisplayFromStr};
use std::ffi::{CStr, CString};
use std::pin::Pin;
use std::sync::{Arc, Mutex};

use crate::extension::{Dumped, Extension, ExternalMethod, RawResource, ResourceSymbols};
use crate::Error;

use super::{RawResourceMethod, Resource, ResourceMethod, ResourceType};

#[serde_as]
#[derive(Debug, Clone, Serialize, Deserialize)]
struct External {
    extension: String,
    resource: String,
    #[serde_as(as = "DisplayFromStr")]
    #[serde(default)]
    version_req: semver::VersionReq,
    #[serde(default)]
    #[serde(skip_serializing)]
    #[serde(skip_deserializing)]
    ext: Arc<Mutex<Option<Arc<Extension>>>>,
}

impl External {
    fn load_extension(&self) -> Result<(), Error> {
        let extension = crate::extension::try_get(&self.extension, &self.version_req)?;

        if extension.get_resource(&self.resource).is_none() {
            return Err(format!(
                "extension {} has no resource type named {}",
                self.extension, self.resource
            )
            .into());
        }

        let mut lock = self.ext.lock().expect("poisoned");
        *lock = Some(extension.clone());

        Ok(())
    }

    fn extension(&self) -> Arc<Extension> {
        self.ext
            .lock()
            .expect("poisoned")
            .as_ref()
            .expect("extension was aready loaded")
            .clone()
    }

    fn resource(&self) -> ResourceSymbols {
        self.extension()
            .get_resource(&self.resource)
            .expect("resource not found in extension")
    }
}

#[typetag::serde]
impl ResourceType for External {
    fn from_bytes(&self, bytes: &[u8]) -> Result<Pin<Box<dyn Resource>>, Error> {
        // The _only_ way to create an `ExternalResource` is through this function. This
        // guarantees that the extension was initalized and that the resource exists.
        self.load_extension()?;

        let extension = self.extension();
        let resource = self.resource();
        let outcome = unsafe {
            // Safety: extension is correctly implemented.
            (resource.fn_from_bytes)(bytes.as_ptr(), bytes.len())
        };
        let raw_ptr = unsafe {
            // Safety: the outcome wasjust generated by the extension and never used.
            extension.outcome_to_result(outcome)?
        };

        if raw_ptr.is_null() {
            return Err(format!("loaded resource for {self:?} from bytes was null").into());
        }

        Ok(Box::pin(ExternalResource {
            r#type: self.clone(),
            ptr: RawResource(raw_ptr),
        }))
    }
}

#[derive(Debug)]
struct ExternalResource {
    r#type: External,
    ptr: RawResource,
}

// Safety: all resources must be thread-safe.
unsafe impl Send for ExternalResource {}
// Safety: all resources must be thread-safe.
unsafe impl Sync for ExternalResource {}

impl Drop for ExternalResource {
    fn drop(&mut self) {
        // This cannot panic, ever! Therefore, we prefer to leak, if necessary.
        // Some `else`s below are never reachable, but... better safe than sorry.

        let Ok(lock) = self.r#type.ext.lock() else {
            // Lock was poisoned. Leak resource!
            return;
        };
        let Some(extension) = &*lock else {
            // No extension was ever loaded. No resource to leak.
            return;
        };
        let Some(resource) = extension.get_resource(&self.r#type.resource) else {
            // No such resource in extension. Unreachable, but if it ever happens, leak!
            return;
        };

        // ... now, for the scary part!
        unsafe {
            // Safety: extension is correctly implemented.
            (resource.fn_drop)(self.ptr);
        }
    }
}

impl Resource for ExternalResource {
    fn r#type(&self) -> Arc<dyn ResourceType> {
        Arc::new(self.r#type.clone())
    }

    fn dump(&self) -> Result<Vec<u8>, Error> {
        let resource = self.r#type.resource();
        let extension = self.r#type.extension();
        unsafe {
            // Safety: extension is correctly implemented.
            let maybe_outcome = (resource.fn_dump)(self.ptr);
            if maybe_outcome.0.is_null() {
                return Err(format!("dumped resource for {:?} was null", self.r#type).into());
            }
            extension.dumped_to_vec(Dumped(extension.outcome_to_result(maybe_outcome)?))
        }
    }

    fn size(&self) -> usize {
        unsafe {
            // Safety: extension is correctly implemented.
            (self.r#type.resource().fn_size)(self.ptr)
        }
    }

    fn get_raw_ptr(&self) -> *const () {
        self.ptr.0 as *const ()
    }

    fn get_method(&self, method: &str) -> Option<ResourceMethod> {
        let c_method = CString::new(method.as_bytes()).expect("method cannot contain nul bytes");
        let extension = self.r#type.extension();
        let resource = self.r#type.resource();

        let external_method = unsafe {
            // Safety: extension is correctly implemented.
            let maybe_method = (resource.fn_get_method_def)(self.ptr, c_method.as_ptr());
            if maybe_method.is_null() {
                return None;
            }
            scopeguard::defer! {
                (extension.string.fn_drop)(maybe_method)
            }

            serde_json::from_slice::<ExternalMethod>(CStr::from_ptr(maybe_method).to_bytes())
                .expect("badly formed json from fn_get_method call")
        };

        Some(ResourceMethod {
            fn_ptr: unsafe {
                // Safety: this should have been a valid address in the extension side.
                std::mem::transmute::<usize, RawResourceMethod>(external_method.fn_ptr)
            },
            input_layout: external_method.input_layout,
            output_layout: external_method.output_layout,
        })
    }
}