ort2 0.1.2

onnxruntime wrapper c/c++ api
Documentation
use std::{ffi::c_void, marker::PhantomData, ptr::null_mut};

use ort2_sys as ffi;
use tracing::trace;

use crate::{api::ok, error::Result, memory::MemoryInfo, session::Session};

pub trait AllocatorTrait {
    fn inner(&self) -> *mut ffi::OrtAllocator;

    #[allow(clippy::not_unsafe_ptr_arg_deref)]
    fn free(&self, ptr: *mut c_void) -> Result<()> {
        ok!(AllocatorFree, self.inner(), ptr)?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct Allocator<'a> {
    inner: *mut ffi::OrtAllocator,
    marker: PhantomData<&'a Session>,
}

impl<'a> Allocator<'a> {
    pub fn new(session: &'a Session, mem_info: &MemoryInfo) -> Result<Self> {
        let mut inner = null_mut();
        ok!(
            CreateAllocator,
            session.inner(),
            mem_info.inner(),
            &mut inner
        )?;
        Ok(Self {
            inner,
            marker: PhantomData,
        })
    }
}

impl Drop for Allocator<'_> {
    fn drop(&mut self) {
        trace!(?self, "dropping");
    }
}

impl AllocatorTrait for Allocator<'_> {
    fn inner(&self) -> *mut ffi::OrtAllocator {
        self.inner
    }
}

#[derive(Debug)]
pub struct DefaultAllocator {
    inner: *mut ffi::OrtAllocator,
}

impl AllocatorTrait for DefaultAllocator {
    fn inner(&self) -> *mut ffi::OrtAllocator {
        self.inner
    }
}

impl DefaultAllocator {
    pub fn new() -> Result<Self> {
        let mut inner = null_mut();
        ok!(GetAllocatorWithDefaultOptions, &mut inner)?;
        Ok(DefaultAllocator { inner })
    }
}

impl Default for DefaultAllocator {
    fn default() -> Self {
        Self::new().expect("failed to get default allocator")
    }
}