From 121c16592476754800b00a3b51595f4799944d04 Mon Sep 17 00:00:00 2001
From: Mason Reed <mason@vector35.com>
Date: Thu, 1 Aug 2024 14:45:39 -0400
Subject: Pass length to free register list callback

Allows language bindings like rust to free register lists sanely
---
 rust/src/callingconvention.rs | 69 +++++++++++++++++++++----------------------
 1 file changed, 34 insertions(+), 35 deletions(-)

(limited to 'rust/src/callingconvention.rs')

diff --git a/rust/src/callingconvention.rs b/rust/src/callingconvention.rs
index a009b439..915bdcc2 100644
--- a/rust/src/callingconvention.rs
+++ b/rust/src/callingconvention.rs
@@ -81,34 +81,13 @@ where
         })
     }
 
-    fn alloc_register_list<I: Iterator<Item = u32> + ExactSizeIterator>(
-        items: I,
-        count: &mut usize,
-    ) -> *mut u32 {
-        let len = items.len();
-        *count = len;
-
-        if len == 0 {
-            return ptr::null_mut();
-        }
-
-        let res: Box<[_]> = [len as u32].into_iter().chain(items).collect();
-        debug_assert!(res.len() == len + 1);
-
-        // it's free on the function below: `cb_free_register_list`
-        let raw = Box::leak(res);
-        &mut raw[1]
-    }
-
-    extern "C" fn cb_free_register_list(_ctxt: *mut c_void, regs: *mut u32) {
+    extern "C" fn cb_free_register_list(_ctxt: *mut c_void, regs: *mut u32, count: usize) {
         ffi_wrap!("CallingConvention::free_register_list", unsafe {
             if regs.is_null() {
                 return;
             }
-
-            let actual_start = regs.offset(-1);
-            let len = (*actual_start) + 1;
-            let _regs = Box::from_raw(ptr::slice_from_raw_parts_mut(actual_start, len as usize));
+            
+            let _regs = Box::from_raw(ptr::slice_from_raw_parts_mut(regs, count));
         })
     }
 
@@ -118,9 +97,13 @@ where
     {
         ffi_wrap!("CallingConvention::caller_saved_registers", unsafe {
             let ctxt = &*(ctxt as *mut CustomCallingConventionContext<C>);
-            let regs = ctxt.cc.caller_saved_registers();
+            let mut regs = ctxt.cc.caller_saved_registers();
 
-            alloc_register_list(regs.iter().map(|r| r.id()), &mut *count)
+            // SAFETY: `count` is an out parameter
+            unsafe { *count = regs.len() };
+            let regs_ptr = regs.as_mut_ptr();
+            mem::forget(regs);
+            regs_ptr as *mut _
         })
     }
 
@@ -130,9 +113,13 @@ where
     {
         ffi_wrap!("CallingConvention::callee_saved_registers", unsafe {
             let ctxt = &*(ctxt as *mut CustomCallingConventionContext<C>);
-            let regs = ctxt.cc.callee_saved_registers();
-
-            alloc_register_list(regs.iter().map(|r| r.id()), &mut *count)
+            let mut regs = ctxt.cc.callee_saved_registers();
+            
+            // SAFETY: `count` is an out parameter
+            unsafe { *count = regs.len() };
+            let regs_ptr = regs.as_mut_ptr();
+            mem::forget(regs);
+            regs_ptr as *mut _
         })
     }
 
@@ -142,9 +129,13 @@ where
     {
         ffi_wrap!("CallingConvention::int_arg_registers", unsafe {
             let ctxt = &*(ctxt as *mut CustomCallingConventionContext<C>);
-            let regs = ctxt.cc.int_arg_registers();
+            let mut regs = ctxt.cc.int_arg_registers();
 
-            alloc_register_list(regs.iter().map(|r| r.id()), &mut *count)
+            // SAFETY: `count` is an out parameter
+            unsafe { *count = regs.len() };
+            let regs_ptr = regs.as_mut_ptr();
+            mem::forget(regs);
+            regs_ptr as *mut _
         })
     }
 
@@ -154,9 +145,13 @@ where
     {
         ffi_wrap!("CallingConvention::float_arg_registers", unsafe {
             let ctxt = &*(ctxt as *mut CustomCallingConventionContext<C>);
-            let regs = ctxt.cc.float_arg_registers();
+            let mut regs = ctxt.cc.float_arg_registers();
 
-            alloc_register_list(regs.iter().map(|r| r.id()), &mut *count)
+            // SAFETY: `count` is an out parameter
+            unsafe { *count = regs.len() };
+            let regs_ptr = regs.as_mut_ptr();
+            mem::forget(regs);
+            regs_ptr as *mut _
         })
     }
 
@@ -272,9 +267,13 @@ where
     {
         ffi_wrap!("CallingConvention::implicitly_defined_registers", unsafe {
             let ctxt = &*(ctxt as *mut CustomCallingConventionContext<C>);
-            let regs = ctxt.cc.implicitly_defined_registers();
+            let mut regs = ctxt.cc.implicitly_defined_registers();
 
-            alloc_register_list(regs.iter().map(|r| r.id()), &mut *count)
+            // SAFETY: `count` is an out parameter
+            unsafe { *count = regs.len() };
+            let regs_ptr = regs.as_mut_ptr();
+            mem::forget(regs);
+            regs_ptr as *mut _
         })
     }
 
-- 
cgit v1.3.1