Be careful about dealocating cpp structs in cpp only

barakugav · Aug 11, 2024 · e4d55eb · e4d55eb
1 parent 384ecc9
commit e4d55eb
Show file tree

Hide file tree

Showing 6 changed files with 166 additions and 67 deletions.
diff --git a/executorch-sys/cpp/executorch_rs_ext/api_utils.cpp b/executorch-sys/cpp/executorch_rs_ext/api_utils.cpp
@@ -6,57 +6,57 @@
 
 namespace executorch_rs
 {
-    namespace
-    {
-        template <typename T>
-        struct ManuallyDrop
-        {
-            union
-            {
-                T value;
-            };
-            ManuallyDrop(T &&value) : value(std::move(value)) {}
-            ~ManuallyDrop() {}
+#if defined(EXECUTORCH_RS_STD)
+    template <typename T>
+    Vec<T> crate_Vec(std::vector<T> &&vec)
+    {
+        size_t len = vec.size();
+        T *arr = new T[len];
+        std::move(vec.begin(), vec.end(), arr);
+        return Vec<T>{
+            .data = arr,
+            .len = len,
+            .cap = len,
         };
+    }
 
-        template <typename T>
-        RawVec<T> crate_RawVec(std::vector<T> &&vec)
-        {
-            auto vec2 = ManuallyDrop<std::vector<T>>(std::move(vec));
-            return RawVec<T>{
-                .data = vec2.value.data(),
-                .len = vec2.value.size(),
-                .cap = vec2.value.capacity(),
-            };
-        }
+#define VEC_DESTRUCTOR_IMPL(T, name)          \
+    void Vec_##name##_destructor(Vec<T> *vec) \
+    {                                         \
+        delete[] vec->data;                   \
+    }
 
-        static_assert(sizeof(Result_i64) == sizeof(torch::executor::Result<int64_t>), "Result_i64 size mismatch");
-        // static_assert(offsetof(Result_i64, value_) == offsetof(torch::executor::Result<int64_t>, value_), "Result_i64 value_ offset mismatch");
-        // static_assert(offsetof(Result_i64, error_) == offsetof(torch::executor::Result<int64_t>, error_), "Result_i64 error_ offset mismatch");
-        // static_assert(offsetof(Result_i64, hasValue_) == offsetof(torch::executor::Result<int64_t>, hasValue_), "Result_i64 hasValue_ offset mismatch");
-        Result_i64 crate_Result_i64(const torch::executor::Result<int64_t> &result)
-        {
-            Result_i64 result2{
-                .error_ = torch::executor::Error::Ok,
-                .hasValue_ = false,
-            };
-            memcpy(&result2, &result, sizeof(Result_i64));
-            return result2;
-        }
+    VEC_DESTRUCTOR_IMPL(char, char)
+    VEC_DESTRUCTOR_IMPL(Vec<char>, Vec_char)
+    VEC_DESTRUCTOR_IMPL(torch::executor::EValue, EValue)
+#endif
 
-        static_assert(sizeof(Result_MethodMeta) == sizeof(torch::executor::Result<torch::executor::MethodMeta>), "Result_MethodMeta size mismatch");
-        // static_assert(offsetof(Result_MethodMeta, value_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, value_), "Result_MethodMeta value_ offset mismatch");
-        // static_assert(offsetof(Result_MethodMeta, error_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, error_), "Result_MethodMeta error_ offset mismatch");
-        // static_assert(offsetof(Result_MethodMeta, hasValue_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, hasValue_), "Result_MethodMeta hasValue_ offset mismatch");
-        Result_MethodMeta crate_Result_MethodMeta(const torch::executor::Result<torch::executor::MethodMeta> &result)
-        {
-            Result_MethodMeta result2{
-                .error_ = torch::executor::Error::Ok,
-                .hasValue_ = false,
-            };
-            memcpy(&result2, &result, sizeof(Result_MethodMeta));
-            return result2;
-        }
+    static_assert(sizeof(Result_i64) == sizeof(torch::executor::Result<int64_t>), "Result_i64 size mismatch");
+    // static_assert(offsetof(Result_i64, value_) == offsetof(torch::executor::Result<int64_t>, value_), "Result_i64 value_ offset mismatch");
+    // static_assert(offsetof(Result_i64, error_) == offsetof(torch::executor::Result<int64_t>, error_), "Result_i64 error_ offset mismatch");
+    // static_assert(offsetof(Result_i64, hasValue_) == offsetof(torch::executor::Result<int64_t>, hasValue_), "Result_i64 hasValue_ offset mismatch");
+    Result_i64 crate_Result_i64(const torch::executor::Result<int64_t> &result)
+    {
+        Result_i64 result2{
+            .error_ = torch::executor::Error::Ok,
+            .hasValue_ = false,
+        };
+        memcpy(&result2, &result, sizeof(Result_i64));
+        return result2;
+    }
+
+    static_assert(sizeof(Result_MethodMeta) == sizeof(torch::executor::Result<torch::executor::MethodMeta>), "Result_MethodMeta size mismatch");
+    // static_assert(offsetof(Result_MethodMeta, value_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, value_), "Result_MethodMeta value_ offset mismatch");
+    // static_assert(offsetof(Result_MethodMeta, error_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, error_), "Result_MethodMeta error_ offset mismatch");
+    // static_assert(offsetof(Result_MethodMeta, hasValue_) == offsetof(torch::executor::Result<torch::executor::MethodMeta>, hasValue_), "Result_MethodMeta hasValue_ offset mismatch");
+    Result_MethodMeta crate_Result_MethodMeta(const torch::executor::Result<torch::executor::MethodMeta> &result)
+    {
+        Result_MethodMeta result2{
+            .error_ = torch::executor::Error::Ok,
+            .hasValue_ = false,
+        };
+        memcpy(&result2, &result, sizeof(Result_MethodMeta));
+        return result2;
     }
 
     Result_MethodMeta Program_method_meta(const torch::executor::Program *program, const char *method_name)
@@ -151,6 +151,10 @@ namespace executorch_rs
         tensor->~Tensor();
     }
 
+    torch::executor::EValue EValue_shallow_clone(torch::executor::EValue *evalue)
+    {
+        return *evalue;
+    }
     void EValue_destructor(torch::executor::EValue *evalue)
     {
         evalue->~EValue();
@@ -180,16 +184,16 @@ namespace executorch_rs
     {
         module->~Module();
     }
-    torch::executor::Result<RawVec<RawVec<char>>> Module_method_names(torch::executor::Module *module)
+    torch::executor::Result<Vec<Vec<char>>> Module_method_names(torch::executor::Module *module)
     {
         std::unordered_set<std::string> method_names = ET_UNWRAP(module->method_names());
-        std::vector<RawVec<char>> method_names_vec;
+        std::vector<Vec<char>> method_names_vec;
         for (const std::string &method_name : method_names)
         {
             std::vector<char> method_name_vec(method_name.begin(), method_name.end());
-            method_names_vec.push_back(crate_RawVec(std::move(method_name_vec)));
+            method_names_vec.push_back(crate_Vec(std::move(method_name_vec)));
         }
-        return crate_RawVec(std::move(method_names_vec));
+        return crate_Vec(std::move(method_names_vec));
     }
     torch::executor::Error Module_load_method(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name)
     {
@@ -206,12 +210,12 @@ namespace executorch_rs
         std::string method_name_str(method_name.begin(), method_name.end());
         return crate_Result_MethodMeta(module->method_meta(method_name_str));
     }
-    torch::executor::Result<RawVec<torch::executor::EValue>> Module_execute(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name, torch::executor::ArrayRef<torch::executor::EValue> inputs)
+    torch::executor::Result<Vec<torch::executor::EValue>> Module_execute(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name, torch::executor::ArrayRef<torch::executor::EValue> inputs)
     {
         std::string method_name_str(method_name.begin(), method_name.end());
         std::vector<torch::executor::EValue> inputs_vec(inputs.begin(), inputs.end());
         std::vector<torch::executor::EValue> outputs = ET_UNWRAP(module->execute(method_name_str, inputs_vec));
-        return crate_RawVec(std::move(outputs));
+        return crate_Vec(std::move(outputs));
     }
 #endif
 }
diff --git a/executorch-sys/cpp/executorch_rs_ext/api_utils.hpp b/executorch-sys/cpp/executorch_rs_ext/api_utils.hpp
@@ -25,14 +25,24 @@
 
 namespace executorch_rs
 {
+#if defined(EXECUTORCH_RS_STD)
     template <typename T>
-    struct RawVec
+    struct Vec
     {
         T *data;
         size_t len;
         size_t cap;
     };
 
+#define VEC_DESTRUCTOR_DEC(T, name) \
+    void Vec_##name##_destructor(Vec<T> *vec);
+
+    VEC_DESTRUCTOR_DEC(char, char)
+    VEC_DESTRUCTOR_DEC(Vec<char>, Vec_char)
+    VEC_DESTRUCTOR_DEC(torch::executor::EValue, EValue)
+
+#endif
+
     struct Result_i64
     {
 
@@ -84,6 +94,7 @@ namespace executorch_rs
     void *Tensor_mutable_data_ptr(const exec_aten::Tensor *tensor);
     void Tensor_destructor(exec_aten::Tensor *tensor);
 
+    torch::executor::EValue EValue_shallow_clone(torch::executor::EValue *evalue);
     void EValue_destructor(torch::executor::EValue *evalue);
     const exec_aten::ArrayRef<int64_t> BoxedEvalueList_i64_get(const torch::executor::BoxedEvalueList<int64_t> *list);
     const exec_aten::ArrayRef<exec_aten::Tensor> BoxedEvalueList_Tensor_get(const torch::executor::BoxedEvalueList<exec_aten::Tensor> *list);
@@ -93,11 +104,11 @@ namespace executorch_rs
 #if defined(EXECUTORCH_RS_MODULE)
     torch::executor::Module Module_new(torch::executor::ArrayRef<char> file_path, torch::executor::Module::MlockConfig mlock_config, torch::executor::EventTracer *event_tracer);
     void Module_destructor(torch::executor::Module *module);
-    torch::executor::Result<RawVec<RawVec<char>>> Module_method_names(torch::executor::Module *module);
+    torch::executor::Result<Vec<Vec<char>>> Module_method_names(torch::executor::Module *module);
     torch::executor::Error Module_load_method(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name);
     bool Module_is_method_loaded(const torch::executor::Module *module, torch::executor::ArrayRef<char> method_name);
     Result_MethodMeta Module_method_meta(const torch::executor::Module *module, torch::executor::ArrayRef<char> method_name);
-    torch::executor::Result<RawVec<torch::executor::EValue>> Module_execute(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name, torch::executor::ArrayRef<torch::executor::EValue> inputs);
+    torch::executor::Result<Vec<torch::executor::EValue>> Module_execute(torch::executor::Module *module, torch::executor::ArrayRef<char> method_name, torch::executor::ArrayRef<torch::executor::EValue> inputs);
 #endif
 
 }
diff --git a/executorch-sys/src/lib.rs b/executorch-sys/src/lib.rs
@@ -109,3 +109,13 @@ impl Drop for et_c::Tensor {
         unsafe { et_rs_c::Tensor_destructor(self) }
     }
 }
+
+impl<T> et_rs_c::Vec<T> {
+    pub fn as_slice(&self) -> &[T] {
+        unsafe { std::slice::from_raw_parts(self.data, self.len) }
+    }
+
+    pub fn as_mut_slice(&mut self) -> &mut [T] {
+        unsafe { std::slice::from_raw_parts_mut(self.data, self.len) }
+    }
+}
diff --git a/src/evalue.rs b/src/evalue.rs
@@ -469,6 +469,12 @@ impl Debug for EValue<'_> {
         st.finish()
     }
 }
+impl<'a> Clone for EValue<'a> {
+    fn clone(&self) -> Self {
+        let value = unsafe { et_rs_c::EValue_shallow_clone(&self.0 as *const _ as *mut _) };
+        unsafe { EValue::new(value) }
+    }
+}
 
 /// Helper class used to correlate EValues in the executor table, with the
 /// unwrapped list of the proper type. Because values in the runtime's values

diff --git a/src/module.rs b/src/module.rs
@@ -77,11 +77,13 @@ impl Module {
     ///
     /// A set of strings containing the names of the methods, or an error if the program or method failed to load.
     pub fn method_names(&mut self) -> Result<HashSet<String>> {
-        let names = unsafe { et_rs_c::Module_method_names(&mut self.0) }.rs()?;
+        let names = unsafe { et_rs_c::Module_method_names(&mut self.0) }
+            .rs()?
+            .rs();
         Ok(names
-            .rs()
-            .into_iter()
-            .map(|s| util::chars2string(s.rs()))
+            .as_slice()
+            .iter()
+            .map(|s| util::chars2string(s.to_vec()))
             .collect())
     }
 
@@ -168,10 +170,11 @@ impl Module {
         let inputs = ArrayRef::from_slice(inputs);
         let outputs =
             unsafe { et_rs_c::Module_execute(&mut self.0, method_name.0, inputs.0) }.rs()?;
-        let outputs = outputs.rs();
         // Safety: The transmute is safe because the memory layout of EValue and et_c::EValue is the same.
-        let outputs = unsafe { std::mem::transmute::<Vec<et_c::EValue>, Vec<EValue<'a>>>(outputs) };
-        Ok(outputs)
+        let outputs = unsafe {
+            std::mem::transmute::<et_rs_c::Vec<et_c::EValue>, et_rs_c::Vec<EValue<'a>>>(outputs)
+        };
+        Ok(outputs.rs().to_vec())
     }
 
     /// Executes the 'forward' method with the given input and retrieves the output.

diff --git a/src/util.rs b/src/util.rs
@@ -211,10 +211,75 @@ pub(crate) fn chars2string(chars: Vec<std::ffi::c_char>) -> String {
 }
 
 #[cfg(feature = "std")]
-impl<T> IntoRust for crate::et_rs_c::RawVec<T> {
-    type RsType = Vec<T>;
-    fn rs(self) -> Self::RsType {
-        unsafe { Vec::from_raw_parts(self.data, self.len, self.cap) }
+mod cpp_vec {
+    use super::IntoRust;
+    use crate::{et_c, et_rs_c, evalue::EValue};
+
+    pub(crate) struct CppVec<T: CppVecElm>(et_rs_c::Vec<T>);
+    impl<T: CppVecElm> CppVec<T> {
+        pub fn as_slice(&self) -> &[T] {
+            self.0.as_slice()
+        }
+
+        pub fn to_vec(&self) -> Vec<T>
+        where
+            T: Clone,
+        {
+            self.as_slice().to_vec()
+        }
+    }
+    impl<T: CppVecElm> IntoRust for et_rs_c::Vec<T> {
+        type RsType = CppVec<T>;
+        fn rs(self) -> Self::RsType {
+            CppVec(self)
+        }
+    }
+    impl IntoRust for et_rs_c::Vec<et_rs_c::Vec<core::ffi::c_char>> {
+        type RsType = CppVec<CppVec<core::ffi::c_char>>;
+        fn rs(self) -> Self::RsType {
+            // Safety: et_rs_c::Vec<T> has the same memory layout as CppVec<T>.
+            unsafe {
+                std::mem::transmute::<
+                    et_rs_c::Vec<et_rs_c::Vec<core::ffi::c_char>>,
+                    CppVec<CppVec<core::ffi::c_char>>,
+                >(self)
+            }
+        }
+    }
+    impl<T: CppVecElm> Drop for CppVec<T> {
+        fn drop(&mut self) {
+            T::drop_vec(self);
+        }
+    }
+    pub(crate) trait CppVecElm: Sized {
+        fn drop_vec(vec: &mut CppVec<Self>);
+    }
+    impl CppVecElm for core::ffi::c_char {
+        fn drop_vec(vec: &mut CppVec<Self>) {
+            unsafe { et_rs_c::Vec_char_destructor(&mut vec.0) }
+        }
+    }
+    impl CppVecElm for CppVec<core::ffi::c_char> {
+        fn drop_vec(vec: &mut CppVec<Self>) {
+            // Safety: CppVec<T> has the same memory layout as et_rs_c::Vec<T>.
+            let vec = unsafe {
+                std::mem::transmute::<
+                    &mut CppVec<CppVec<core::ffi::c_char>>,
+                    &mut et_rs_c::Vec<et_rs_c::Vec<core::ffi::c_char>>,
+                >(vec)
+            };
+            unsafe { et_rs_c::Vec_Vec_char_destructor(vec) }
+        }
+    }
+    impl<'a> CppVecElm for EValue<'a> {
+        fn drop_vec(vec: &mut CppVec<Self>) {
+            let vec = unsafe {
+                std::mem::transmute::<&mut et_rs_c::Vec<EValue<'a>>, &mut et_rs_c::Vec<et_c::EValue>>(
+                    &mut vec.0,
+                )
+            };
+            unsafe { et_rs_c::Vec_EValue_destructor(vec) }
+        }
     }
 }