Merge pull request #6 from SunDoge/boxed_slice

reduce memory usage

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "dlpark"
-version = "0.3.0"
+version = "0.4.0"
 edition = "2021"
 authors = ["SunDoge"]
 license = "Apache-2.0"
@@ -13,8 +13,8 @@ readme = "README.md"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-half = { version = "2.3.1", optional = true }
-pyo3 = { version = "0.19.0", optional = true }
+half = { version = "2.3", optional = true }
+pyo3 = { version = "0.19", optional = true }
 
 [workspace]
 members = ["examples/from_numpy", "examples/with_pyo3"]

src/manager_ctx.rs

@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::ptr::NonNull;
 
 use crate::tensor::traits::{IntoDLPack, TensorView};
@@ -16,28 +15,31 @@ unsafe extern "C" fn deleter_fn<T>(dl_managed_tensor: *mut ffi::DLManagedTensor)
 /// The lifetime should be 'static since we don't managed its memory.
 /// Otherwise, we should copy the data and convert its type to i64 and managed it ourselves.
 #[derive(Debug)]
-pub struct CowIntArray(Cow<'static, [i64]>);
+pub enum CowIntArray {
+    Owned(Box<[i64]>),
+    Borrowed(&'static [i64]),
+}
 
 impl CowIntArray {
-    pub fn from_owned(v: Vec<i64>) -> Self {
-        Self(Cow::Owned(v))
+    pub fn from_owned(v: Box<[i64]>) -> Self {
+        Self::Owned(v)
     }
 
     pub fn from_borrowed(v: &'static [i64]) -> Self {
-        Self(Cow::Borrowed(v))
+        Self::Borrowed(v)
     }
 
     pub fn as_ptr(&self) -> *mut i64 {
-        match self.0 {
-            Cow::Borrowed(v) => v.as_ptr() as *mut i64,
-            Cow::Owned(ref v) => v.as_ptr() as *mut i64,
+        match self {
+            Self::Borrowed(v) => v.as_ptr() as *mut i64,
+            Self::Owned(ref v) => v.as_ptr() as *mut i64,
         }
     }
 
     fn len(&self) -> usize {
-        match self.0 {
-            Cow::Borrowed(v) => v.len(),
-            Cow::Owned(ref v) => v.len(),
+        match self {
+            Self::Borrowed(v) => v.len(),
+            Self::Owned(ref v) => v.len(),
         }
     }
 

src/tensor/impls.rs

@@ -32,11 +32,11 @@ macro_rules! impl_for_rust_type {
             }
 
             fn shape(&self) -> CowIntArray {
-                CowIntArray::from_owned(vec![])
+                CowIntArray::from_owned(Box::new([]))
             }
 
             fn strides(&self) -> Option<CowIntArray> {
-                Some(CowIntArray::from_owned(vec![]))
+                Some(CowIntArray::from_owned(Box::new([])))
             }
         }
     };
@@ -85,11 +85,11 @@ where
     }
 
     fn shape(&self) -> CowIntArray {
-        CowIntArray::from_owned(vec![self.len() as i64])
+        CowIntArray::from_owned(Box::new([self.len() as i64]))
     }
 
     fn strides(&self) -> Option<CowIntArray> {
-        Some(CowIntArray::from_owned(vec![1]))
+        Some(CowIntArray::from_owned(Box::new([1])))
     }
 }
 
@@ -114,11 +114,11 @@ where
     }
 
     fn shape(&self) -> CowIntArray {
-        CowIntArray::from_owned(vec![self.len() as i64])
+        CowIntArray::from_owned(Box::new([self.len() as i64]))
     }
 
     fn strides(&self) -> Option<CowIntArray> {
-        Some(CowIntArray::from_owned(vec![1]))
+        Some(CowIntArray::from_owned(Box::new([1])))
     }
 }
 
@@ -143,11 +143,11 @@ where
     }
 
     fn shape(&self) -> CowIntArray {
-        CowIntArray::from_owned(vec![self.len() as i64])
+        CowIntArray::from_owned(Box::new([self.len() as i64]))
     }
 
     fn strides(&self) -> Option<CowIntArray> {
-        Some(CowIntArray::from_owned(vec![1]))
+        Some(CowIntArray::from_owned(Box::new([1])))
     }
 }
 

src/tensor/traits.rs

@@ -76,7 +76,7 @@ pub trait ToTensor {
 
     fn calculate_contiguous_strides(&self) -> CowIntArray {
         let strides = calculate_contiguous_strides(self.shape().as_slice());
-        CowIntArray::from_owned(strides)
+        CowIntArray::from_owned(strides.into_boxed_slice())
     }
 }