feat: expose HostMemoryOrVec via API

trtexec-rs/src/main.rs

@@ -10,14 +10,13 @@ use rustnn::load_graph_from_path;
 use std::ffi::{c_void, OsString};
 use std::fs::File;
 use std::io::{BufRead, Read, Write};
-use std::ops::Deref;
 use std::sync::atomic::AtomicU32;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::time::Instant;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::{prelude::*, EnvFilter};
-use trtx::host_memory::HostMemory;
+use trtx::host_memory::HostMemoryOrVec;
 use trtx::{Builder, Logger, OnnxParser, ProfilingVerbosity};
 use trtx::{LayerInformationFormat, Runtime};
 
@@ -75,39 +74,6 @@ fn digest_hex(digest: &md5::Digest) -> String {
     format!("{digest:x}")
 }
 
-enum HostMemoryOrVec<'memory> {
-    HostMemory(HostMemory<'memory>),
-    Vec(Vec<u8>),
-}
-
-impl<'memory> AsRef<[u8]> for HostMemoryOrVec<'memory> {
-    fn as_ref(&self) -> &[u8] {
-        match self {
-            HostMemoryOrVec::HostMemory(host_memory) => host_memory.as_ref(),
-            HostMemoryOrVec::Vec(items) => items.as_ref(),
-        }
-    }
-}
-
-impl<'memory> Deref for HostMemoryOrVec<'memory> {
-    type Target = [u8];
-
-    fn deref(&self) -> &Self::Target {
-        self.as_ref()
-    }
-}
-
-impl<'buffer> From<HostMemory<'buffer>> for HostMemoryOrVec<'buffer> {
-    fn from(value: HostMemory<'buffer>) -> Self {
-        HostMemoryOrVec::HostMemory(value)
-    }
-}
-impl From<Vec<u8>> for HostMemoryOrVec<'_> {
-    fn from(value: Vec<u8>) -> Self {
-        HostMemoryOrVec::Vec(value)
-    }
-}
-
 fn main() -> Result<()> {
     let args = Args::parse();
     if let Some(shell) = args.shell_completion {

trtx/src/host_memory.rs

@@ -67,3 +67,70 @@ impl<'builder> Deref for HostMemory<'builder> {
         self.as_ref()
     }
 }
+
+pub enum HostMemoryOrVec<'memory> {
+    HostMemory(HostMemory<'memory>),
+    Vec(Vec<u8>),
+}
+
+impl<'memory> HostMemoryOrVec<'memory> {
+    /// Returns `true` if the host memory or vec is [`HostMemory`].
+    ///
+    /// [`HostMemory`]: HostMemoryOrVec::HostMemory
+    #[must_use]
+    pub fn is_host_memory(&self) -> bool {
+        matches!(self, Self::HostMemory(..))
+    }
+
+    pub fn as_host_memory(&self) -> Option<&HostMemory<'memory>> {
+        if let Self::HostMemory(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+
+    /// Returns `true` if the host memory or vec is [`Vec`].
+    ///
+    /// [`Vec`]: HostMemoryOrVec::Vec
+    #[must_use]
+    pub fn is_vec(&self) -> bool {
+        matches!(self, Self::Vec(..))
+    }
+
+    pub fn as_vec(&self) -> Option<&Vec<u8>> {
+        if let Self::Vec(v) = self {
+            Some(v)
+        } else {
+            None
+        }
+    }
+}
+
+impl<'memory> AsRef<[u8]> for HostMemoryOrVec<'memory> {
+    fn as_ref(&self) -> &[u8] {
+        match self {
+            HostMemoryOrVec::HostMemory(host_memory) => host_memory.as_ref(),
+            HostMemoryOrVec::Vec(items) => items.as_ref(),
+        }
+    }
+}
+
+impl<'memory> Deref for HostMemoryOrVec<'memory> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.as_ref()
+    }
+}
+
+impl<'buffer> From<HostMemory<'buffer>> for HostMemoryOrVec<'buffer> {
+    fn from(value: HostMemory<'buffer>) -> Self {
+        HostMemoryOrVec::HostMemory(value)
+    }
+}
+impl From<Vec<u8>> for HostMemoryOrVec<'_> {
+    fn from(value: Vec<u8>) -> Self {
+        HostMemoryOrVec::Vec(value)
+    }
+}