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")
}
}