From 3aabfc668da028138cfd3bbe94dd80f189b0c25c Mon Sep 17 00:00:00 2001
From: Marijn Suijten <marijn@traverseresearch.nl>
Date: Tue, 10 Nov 2020 11:44:35 +0100
Subject: [PATCH] Free allocated memory using the appropriate API
 (free/CoTaskMemFree)

On Windows the right API needs to be used to free this pointer in the
same way as it was allocated by DXC (using CoTaskMemAlloc). On
non-Windows platforms this call is [delegated] to `free` from libc.

[delegated]: https://github.com/microsoft/DirectXShaderCompiler/blob/a8d9780046cb64a1cea842fa6fc28a250e3e2c09/include/dxc/Support/WinAdapter.h#L46
---
 Cargo.toml                  |  5 +++-
 src/intellisense/wrapper.rs | 48 +++++++++++++++++++++----------------
 src/os.rs                   |  8 +++++++
 3 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2020aed..4bceb27 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,7 +31,10 @@ widestring = "0.4.2"
 thiserror = "1.0"
 
 [target.'cfg(windows)'.dependencies]
-winapi = { version = "0.3.9", features = ["wtypes", "oleauto"] }
+winapi = { version = "0.3.9", features = ["wtypes", "oleauto", "combaseapi"] }
+
+[target.'cfg(unix)'.dependencies]
+libc = "0.2"
 
 [dev-dependencies]
 rspirv = "0.6.0"
diff --git a/src/intellisense/wrapper.rs b/src/intellisense/wrapper.rs
index 89a44f3..f240118 100644
--- a/src/intellisense/wrapper.rs
+++ b/src/intellisense/wrapper.rs
@@ -1,5 +1,5 @@
 use crate::intellisense::ffi::*;
-use crate::os::{BSTR, HRESULT, LPSTR};
+use crate::os::{CoTaskMemFree, BSTR, HRESULT, LPSTR};
 use crate::utils::HassleError;
 use crate::wrapper::Dxc;
 use com_rs::ComPtr;
@@ -178,16 +178,19 @@ impl DxcCursor {
             check_hr!(
                 self.inner
                     .get_children(skip, max_count, &mut result_length, &mut result),
-                // get_children allocates a buffer to pass the result in.  Vec will
-                // free it on Drop.
-                Vec::from_raw_parts(result, result_length as usize, result_length as usize)
-                    .into_iter()
-                    .map(|ptr| {
-                        let mut childcursor = ComPtr::<IDxcCursor>::new();
-                        *childcursor.as_mut_ptr() = ptr;
-                        DxcCursor::new(childcursor)
-                    })
-                    .collect::<Vec<_>>()
+                {
+                    // get_children allocates a buffer to pass the result in.
+                    let child_cursors = std::slice::from_raw_parts(result, result_length as usize)
+                        .iter()
+                        .map(|&ptr| {
+                            let mut childcursor = ComPtr::<IDxcCursor>::new();
+                            *childcursor.as_mut_ptr() = ptr;
+                            DxcCursor::new(childcursor)
+                        })
+                        .collect::<Vec<_>>();
+                    CoTaskMemFree(result as *mut _);
+                    child_cursors
+                }
             )
         }
     }
@@ -359,16 +362,19 @@ impl DxcCursor {
                     &mut result_length,
                     &mut result
                 ),
-                // find_references_in_file allocates a buffer to pass the result in.
-                // Vec will free it on Drop.
-                Vec::from_raw_parts(result, result_length as usize, result_length as usize)
-                    .into_iter()
-                    .map(|ptr| {
-                        let mut childcursor = ComPtr::<IDxcCursor>::new();
-                        *childcursor.as_mut_ptr() = ptr;
-                        DxcCursor::new(childcursor)
-                    })
-                    .collect::<Vec<_>>()
+                {
+                    // find_references_in_file allocates a buffer to pass the result in.
+                    let child_cursors = std::slice::from_raw_parts(result, result_length as usize)
+                        .iter()
+                        .map(|&ptr| {
+                            let mut childcursor = ComPtr::<IDxcCursor>::new();
+                            *childcursor.as_mut_ptr() = ptr;
+                            DxcCursor::new(childcursor)
+                        })
+                        .collect::<Vec<_>>();
+                    CoTaskMemFree(result as *mut _);
+                    child_cursors
+                }
             )
         }
     }
diff --git a/src/os.rs b/src/os.rs
index bda56b8..b27a01e 100644
--- a/src/os.rs
+++ b/src/os.rs
@@ -4,6 +4,8 @@ mod os_defs {
         ntdef::{HRESULT, LPCSTR, LPCWSTR, LPSTR, LPWSTR, WCHAR},
         wtypes::BSTR,
     };
+
+    pub use winapi::um::combaseapi::CoTaskMemFree;
 }
 
 #[cfg(not(windows))]
@@ -18,6 +20,12 @@ mod os_defs {
     pub type BSTR = *mut OLECHAR;
     pub type LPBSTR = *mut BSTR;
     pub type HRESULT = i32;
+
+    #[allow(non_snake_case)]
+    pub unsafe fn CoTaskMemFree(p: *mut libc::c_void) {
+        // https://github.com/microsoft/DirectXShaderCompiler/blob/a8d9780046cb64a1cea842fa6fc28a250e3e2c09/include/dxc/Support/WinAdapter.h#L46
+        libc::free(p)
+    }
 }
 
 pub use os_defs::*;