Skip to main content

hurray_ffi/
buffer_list.rs

1//! Opaque list of buffer handles for multi-buffer tensors (ADR-030).
2//!
3//! A tensor whose descriptor references more than one buffer — per-channel,
4//! NF4 or MXFP quantization, sparse layouts, block-paged, composite — needs all
5//! of its buffers to travel together. [`HurrayBufferList`] is that carrier: an
6//! ordered, owning collection of [`HurrayBuffer`] handles.
7//!
8//! ## Ownership
9//!
10//! The list **owns** every handle pushed into it. [`hurray_buffer_list_get`]
11//! returns a *borrowed* pointer — the caller MUST NOT destroy it.
12//! [`hurray_buffer_list_destroy`] destroys every owned handle exactly once and
13//! then frees the list.
14//!
15//! ## Order
16//!
17//! Element `i` of the list is the buffer at index `i` of the tensor
18//! descriptor's buffer table (ADR-030 § 3). Buffer indices appearing in
19//! quantization descriptors (`scale_buffer_index`, `zero_point_buffer_index`),
20//! layout descriptors, and composite members therefore index this list
21//! directly.
22
23use crate::{
24    buffer::{hurray_buffer_destroy, HurrayBuffer},
25    panic::{catch, null_check},
26    status::{HurrayStatus, HURRAY_ERR_INDEX_OUT_OF_BOUNDS, HURRAY_ERR_INTERNAL, HURRAY_OK},
27};
28
29// ── Debug sentinel constants ──────────────────────────────────────────────────
30
31/// Sentinel placed in `HurrayBufferList.sentinel` on allocation (debug builds only).
32#[cfg(debug_assertions)]
33const SENTINEL_LIVE: u64 = 0x1157_0000_11FE_0000;
34
35/// Sentinel written before the list allocation is freed, to detect a double
36/// destroy in debug builds.
37#[cfg(debug_assertions)]
38const SENTINEL_DEAD: u64 = 0x1157_DEAD_1157_DEAD;
39
40// ── HurrayBufferList ──────────────────────────────────────────────────────────
41
42/// Opaque handle to an ordered, owning list of [`HurrayBuffer`] handles.
43///
44/// Like [`HurrayBuffer`], this struct is **not** `#[repr(C)]`; its layout is an
45/// implementation detail and callers MUST treat it as a black-box pointer.
46/// Keeping it opaque is what allows [`hurray_buffer_list_get`] to bounds-check
47/// and to hand back a borrowed handle — a bare `HurrayBuffer**` array could do
48/// neither.
49pub struct HurrayBufferList {
50    /// Double-destroy sentinel (debug builds only).
51    #[cfg(debug_assertions)]
52    sentinel: u64,
53    /// Owned handles in descriptor-buffer-table order. A slot is nulled as soon
54    /// as its handle is destroyed, so a re-entrant destroy cannot free it twice.
55    buffers: Vec<*mut HurrayBuffer>,
56}
57
58// SAFETY: mirrors the `HurrayBuffer` contract — the C ABI is single-thread-at-a-time
59// per handle, so `Send` only permits moving the box across threads during
60// construction and destruction, never concurrent access.
61unsafe impl Send for HurrayBufferList {}
62
63// ── Constructor ───────────────────────────────────────────────────────────────
64
65/// Creates an empty [`HurrayBufferList`].
66///
67/// `capacity` is a hint only; the list grows as needed. Pass the tensor's
68/// buffer count to avoid reallocation.
69///
70/// The caller owns the returned list and MUST destroy it exactly once with
71/// [`hurray_buffer_list_destroy`].
72///
73/// # Safety
74///
75/// `out_list` MUST be a valid, non-null, writable pointer.
76///
77/// # Examples
78///
79/// ```
80/// use hurray_ffi::buffer_list::{hurray_buffer_list_destroy, hurray_buffer_list_new};
81/// use hurray_ffi::{HurrayBufferList, HURRAY_OK};
82///
83/// let mut list: *mut HurrayBufferList = std::ptr::null_mut();
84/// assert_eq!(unsafe { hurray_buffer_list_new(2, &mut list) }, HURRAY_OK);
85/// assert!(!list.is_null());
86///
87/// assert_eq!(unsafe { hurray_buffer_list_destroy(&mut list) }, HURRAY_OK);
88/// assert!(list.is_null()); // destroy nulls the caller's pointer
89/// ```
90#[no_mangle]
91pub unsafe extern "C" fn hurray_buffer_list_new(
92    capacity: u64,
93    out_list: *mut *mut HurrayBufferList,
94) -> HurrayStatus {
95    catch(|| {
96        null_check!(out_list);
97
98        // Cap the pre-allocation: `capacity` is an untrusted hint, and a bogus
99        // value must not let a caller request an enormous allocation up front.
100        let hint = capacity.min(64) as usize;
101
102        let boxed = Box::new(HurrayBufferList {
103            #[cfg(debug_assertions)]
104            sentinel: SENTINEL_LIVE,
105            buffers: Vec::with_capacity(hint),
106        });
107        // SAFETY: just allocated; ownership transfers to the caller, who must
108        // call hurray_buffer_list_destroy exactly once.
109        *out_list = Box::into_raw(boxed);
110        HURRAY_OK
111    })
112}
113
114// ── Push ──────────────────────────────────────────────────────────────────────
115
116/// Appends `buffer` to `list`, transferring ownership of the handle to the list.
117///
118/// On success the caller MUST NOT destroy `buffer` — destroying the list
119/// destroys it. On failure ownership stays with the caller.
120///
121/// Buffers MUST be pushed in descriptor buffer-table order: the first push is
122/// buffer index `0`, the second is index `1`, and so on.
123///
124/// # Safety
125///
126/// - `list` MUST be a live handle from [`hurray_buffer_list_new`].
127/// - `buffer` MUST be a live handle from `hurray_buffer_from_ptr` that has not
128///   been destroyed and is not already owned by another list.
129///
130/// # Examples
131///
132/// ```
133/// use hurray_ffi::buffer::hurray_buffer_from_ptr;
134/// use hurray_ffi::buffer_list::{
135///     hurray_buffer_list_destroy, hurray_buffer_list_len, hurray_buffer_list_new,
136///     hurray_buffer_list_push,
137/// };
138/// use hurray_ffi::{HurrayBuffer, HurrayBufferList, HURRAY_OK};
139///
140/// #[repr(align(64))]
141/// struct Aligned([u8; 64]);
142/// let mut data = Aligned([0u8; 64]);
143///
144/// let mut buffer: *mut HurrayBuffer = std::ptr::null_mut();
145/// assert_eq!(
146///     unsafe {
147///         hurray_buffer_from_ptr(
148///             data.0.as_mut_ptr().cast(), 64, 64, 0, 0, 0,
149///             None, std::ptr::null_mut(), &mut buffer,
150///         )
151///     },
152///     HURRAY_OK
153/// );
154///
155/// let mut list: *mut HurrayBufferList = std::ptr::null_mut();
156/// unsafe { hurray_buffer_list_new(1, &mut list) };
157/// assert_eq!(unsafe { hurray_buffer_list_push(list, buffer) }, HURRAY_OK);
158///
159/// let mut len: u64 = 0;
160/// unsafe { hurray_buffer_list_len(list, &mut len) };
161/// assert_eq!(len, 1);
162///
163/// // Destroying the list destroys the pushed buffer too.
164/// unsafe { hurray_buffer_list_destroy(&mut list) };
165/// ```
166#[no_mangle]
167pub unsafe extern "C" fn hurray_buffer_list_push(
168    list: *mut HurrayBufferList,
169    buffer: *mut HurrayBuffer,
170) -> HurrayStatus {
171    catch(|| {
172        null_check!(list, buffer);
173
174        #[cfg(debug_assertions)]
175        // SAFETY: list is non-null and points to a live HurrayBufferList.
176        if (*list).sentinel == SENTINEL_DEAD {
177            return HURRAY_ERR_INTERNAL;
178        }
179
180        // SAFETY: list is non-null and points to a live HurrayBufferList.
181        (*list).buffers.push(buffer);
182        HURRAY_OK
183    })
184}
185
186// ── Accessors ─────────────────────────────────────────────────────────────────
187
188/// Reads the number of buffers in `list`.
189///
190/// # Safety
191///
192/// - `list` MUST be a live handle from [`hurray_buffer_list_new`].
193/// - `out_len` MUST be a valid, non-null, writable pointer.
194///
195/// # Examples
196///
197/// ```
198/// use hurray_ffi::buffer_list::{
199///     hurray_buffer_list_destroy, hurray_buffer_list_len, hurray_buffer_list_new,
200/// };
201/// use hurray_ffi::{HurrayBufferList, HURRAY_OK};
202///
203/// let mut list: *mut HurrayBufferList = std::ptr::null_mut();
204/// unsafe { hurray_buffer_list_new(0, &mut list) };
205///
206/// let mut len: u64 = 7;
207/// assert_eq!(unsafe { hurray_buffer_list_len(list, &mut len) }, HURRAY_OK);
208/// assert_eq!(len, 0);
209///
210/// unsafe { hurray_buffer_list_destroy(&mut list) };
211/// ```
212#[no_mangle]
213pub unsafe extern "C" fn hurray_buffer_list_len(
214    list: *const HurrayBufferList,
215    out_len: *mut u64,
216) -> HurrayStatus {
217    catch(|| {
218        null_check!(list, out_len);
219
220        // SAFETY: list is non-null and points to a live HurrayBufferList.
221        *out_len = (*list).buffers.len() as u64;
222        HURRAY_OK
223    })
224}
225
226/// Borrows the [`HurrayBuffer`] at `index`.
227///
228/// The returned handle is **borrowed**: ownership stays with the list, and the
229/// caller MUST NOT call `hurray_buffer_destroy` on it. It stays valid until the
230/// list is destroyed.
231///
232/// Returns [`HURRAY_ERR_INDEX_OUT_OF_BOUNDS`] if `index` is not less than the
233/// list length.
234///
235/// # Safety
236///
237/// - `list` MUST be a live handle from [`hurray_buffer_list_new`].
238/// - `out_buffer` MUST be a valid, non-null, writable pointer.
239///
240/// # Examples
241///
242/// ```
243/// use hurray_ffi::buffer_list::{
244///     hurray_buffer_list_destroy, hurray_buffer_list_get, hurray_buffer_list_new,
245/// };
246/// use hurray_ffi::status::HURRAY_ERR_INDEX_OUT_OF_BOUNDS;
247/// use hurray_ffi::{HurrayBuffer, HurrayBufferList};
248///
249/// let mut list: *mut HurrayBufferList = std::ptr::null_mut();
250/// unsafe { hurray_buffer_list_new(0, &mut list) };
251///
252/// // An empty list has no index 0.
253/// let mut got: *mut HurrayBuffer = std::ptr::null_mut();
254/// assert_eq!(
255///     unsafe { hurray_buffer_list_get(list, 0, &mut got) },
256///     HURRAY_ERR_INDEX_OUT_OF_BOUNDS
257/// );
258///
259/// unsafe { hurray_buffer_list_destroy(&mut list) };
260/// ```
261#[no_mangle]
262pub unsafe extern "C" fn hurray_buffer_list_get(
263    list: *const HurrayBufferList,
264    index: u64,
265    out_buffer: *mut *mut HurrayBuffer,
266) -> HurrayStatus {
267    catch(|| {
268        null_check!(list, out_buffer);
269
270        // SAFETY: list is non-null and points to a live HurrayBufferList.
271        let buffers = &(*list).buffers;
272        let Ok(idx) = usize::try_from(index) else {
273            return HURRAY_ERR_INDEX_OUT_OF_BOUNDS;
274        };
275        match buffers.get(idx) {
276            // A nulled slot means the list is mid-destroy; treat it as absent
277            // rather than handing back a dangling handle.
278            Some(&b) if !b.is_null() => {
279                *out_buffer = b;
280                HURRAY_OK
281            }
282            _ => HURRAY_ERR_INDEX_OUT_OF_BOUNDS,
283        }
284    })
285}
286
287// ── Destructor ────────────────────────────────────────────────────────────────
288
289/// Destroys `*list` and every [`HurrayBuffer`] it owns, then writes null through
290/// `list`.
291///
292/// Nulling the caller's pointer is the sound half of Arrow's "release marks the
293/// structure released" discipline: the list allocation itself is freed here, so
294/// a marker written inside it could not be read back, but the caller's own
295/// variable can be invalidated.
296///
297/// Each owned slot is nulled as its handle is destroyed, so a release callback
298/// that panics or re-enters cannot cause a double free. A panicking callback
299/// leaks the remainder of the list rather than corrupting it.
300///
301/// Passing a pointer to a null pointer is a no-op and returns [`HURRAY_OK`],
302/// which makes cleanup paths idempotent.
303///
304/// # Safety
305///
306/// - `list` MUST be a valid, non-null, writable pointer to a `*mut HurrayBufferList`.
307/// - `*list` MUST be a live handle from [`hurray_buffer_list_new`], or null.
308///
309/// # Examples
310///
311/// ```
312/// use hurray_ffi::buffer_list::{hurray_buffer_list_destroy, hurray_buffer_list_new};
313/// use hurray_ffi::{HurrayBufferList, HURRAY_OK};
314///
315/// let mut list: *mut HurrayBufferList = std::ptr::null_mut();
316/// unsafe { hurray_buffer_list_new(0, &mut list) };
317///
318/// assert_eq!(unsafe { hurray_buffer_list_destroy(&mut list) }, HURRAY_OK);
319/// assert!(list.is_null());
320///
321/// // Idempotent: destroying an already-nulled pointer is a no-op.
322/// assert_eq!(unsafe { hurray_buffer_list_destroy(&mut list) }, HURRAY_OK);
323/// ```
324#[no_mangle]
325pub unsafe extern "C" fn hurray_buffer_list_destroy(
326    list: *mut *mut HurrayBufferList,
327) -> HurrayStatus {
328    catch(|| {
329        null_check!(list);
330
331        // SAFETY: list is a valid, writable pointer to a handle slot.
332        let ptr = *list;
333        if ptr.is_null() {
334            return HURRAY_OK;
335        }
336
337        #[cfg(debug_assertions)]
338        {
339            // SAFETY: ptr is non-null and points to a live HurrayBufferList.
340            if (*ptr).sentinel == SENTINEL_DEAD {
341                return HURRAY_ERR_INTERNAL;
342            }
343            // Written through the raw pointer before Box::from_raw so the freed
344            // slot carries the dead marker; writing after would only update the
345            // moved-out local copy.
346            (*ptr).sentinel = SENTINEL_DEAD;
347        }
348
349        // Destroy each owned handle, nulling its slot first so a re-entrant or
350        // panicking release callback can never see a destroyable handle twice.
351        // SAFETY: ptr is non-null and points to a live HurrayBufferList.
352        for slot in (*ptr).buffers.iter_mut() {
353            let buffer = std::mem::replace(slot, std::ptr::null_mut());
354            if !buffer.is_null() {
355                // SAFETY: the list owns this handle, it is non-null, and this is
356                // its first and only destroy.
357                hurray_buffer_destroy(buffer);
358            }
359        }
360
361        // SAFETY: created by hurray_buffer_list_new via Box::into_raw; this is
362        // the first and only destroy.
363        drop(Box::from_raw(ptr));
364
365        // The caller's variable is now observably dead.
366        *list = std::ptr::null_mut();
367        HURRAY_OK
368    })
369}
370
371// ── Tests ─────────────────────────────────────────────────────────────────────
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376    use crate::buffer::{hurray_buffer_byte_size, hurray_buffer_from_ptr};
377    use crate::status::HURRAY_ERR_NULL_POINTER;
378    use std::ffi::c_void;
379    use std::sync::atomic::{AtomicUsize, Ordering};
380
381    #[repr(align(64))]
382    struct Aligned([u8; 128]);
383
384    /// Counts releases through the caller-supplied context pointer rather than a
385    /// global: cargo runs tests in parallel, so a shared counter would be raced by
386    /// any other test that destroys a buffer.
387    unsafe extern "C" fn counting_release(_data: *mut c_void, ctx: *mut c_void) {
388        if !ctx.is_null() {
389            (*(ctx as *const AtomicUsize)).fetch_add(1, Ordering::SeqCst);
390        }
391    }
392
393    fn make_buffer(data: &mut Aligned, size: u64, releases: &AtomicUsize) -> *mut HurrayBuffer {
394        let mut buffer: *mut HurrayBuffer = std::ptr::null_mut();
395        let status = unsafe {
396            hurray_buffer_from_ptr(
397                data.0.as_mut_ptr().cast(),
398                size,
399                64,
400                0,
401                0,
402                0,
403                Some(counting_release),
404                releases as *const AtomicUsize as *mut c_void,
405                &mut buffer,
406            )
407        };
408        assert_eq!(status, HURRAY_OK);
409        buffer
410    }
411
412    #[test]
413    fn new_and_destroy_nulls_the_caller_pointer() {
414        let mut list: *mut HurrayBufferList = std::ptr::null_mut();
415        assert_eq!(unsafe { hurray_buffer_list_new(4, &mut list) }, HURRAY_OK);
416        assert!(!list.is_null());
417        assert_eq!(unsafe { hurray_buffer_list_destroy(&mut list) }, HURRAY_OK);
418        assert!(list.is_null());
419    }
420
421    #[test]
422    fn destroy_is_idempotent_on_a_null_slot() {
423        let mut list: *mut HurrayBufferList = std::ptr::null_mut();
424        assert_eq!(unsafe { hurray_buffer_list_destroy(&mut list) }, HURRAY_OK);
425    }
426
427    #[test]
428    fn null_arguments_are_rejected() {
429        assert_eq!(
430            unsafe { hurray_buffer_list_new(0, std::ptr::null_mut()) },
431            HURRAY_ERR_NULL_POINTER
432        );
433        assert_eq!(
434            unsafe { hurray_buffer_list_destroy(std::ptr::null_mut()) },
435            HURRAY_ERR_NULL_POINTER
436        );
437        let mut len = 0u64;
438        assert_eq!(
439            unsafe { hurray_buffer_list_len(std::ptr::null(), &mut len) },
440            HURRAY_ERR_NULL_POINTER
441        );
442    }
443
444    #[test]
445    fn push_then_len_and_get_round_trip_in_order() {
446        let mut a = Aligned([1u8; 128]);
447        let mut b = Aligned([2u8; 128]);
448        let releases = AtomicUsize::new(0);
449        let buf_a = make_buffer(&mut a, 128, &releases);
450        let buf_b = make_buffer(&mut b, 64, &releases);
451
452        let mut list: *mut HurrayBufferList = std::ptr::null_mut();
453        unsafe { hurray_buffer_list_new(2, &mut list) };
454        assert_eq!(unsafe { hurray_buffer_list_push(list, buf_a) }, HURRAY_OK);
455        assert_eq!(unsafe { hurray_buffer_list_push(list, buf_b) }, HURRAY_OK);
456
457        let mut len = 0u64;
458        unsafe { hurray_buffer_list_len(list, &mut len) };
459        assert_eq!(len, 2);
460
461        // Order is push order, i.e. descriptor buffer-table order: the 128-byte
462        // buffer is index 0 and the 64-byte one is index 1.
463        for (index, expected) in [(0u64, 128u64), (1, 64)] {
464            let mut got: *mut HurrayBuffer = std::ptr::null_mut();
465            assert_eq!(
466                unsafe { hurray_buffer_list_get(list, index, &mut got) },
467                HURRAY_OK
468            );
469            let mut size = 0u64;
470            unsafe { hurray_buffer_byte_size(got, &mut size) };
471            assert_eq!(size, expected);
472        }
473
474        unsafe { hurray_buffer_list_destroy(&mut list) };
475    }
476
477    #[test]
478    fn get_rejects_an_out_of_range_index() {
479        let mut a = Aligned([0u8; 128]);
480        let releases = AtomicUsize::new(0);
481        let buf = make_buffer(&mut a, 128, &releases);
482        let mut list: *mut HurrayBufferList = std::ptr::null_mut();
483        unsafe { hurray_buffer_list_new(1, &mut list) };
484        unsafe { hurray_buffer_list_push(list, buf) };
485
486        let mut got: *mut HurrayBuffer = std::ptr::null_mut();
487        assert_eq!(
488            unsafe { hurray_buffer_list_get(list, 1, &mut got) },
489            HURRAY_ERR_INDEX_OUT_OF_BOUNDS
490        );
491        assert_eq!(
492            unsafe { hurray_buffer_list_get(list, u64::MAX, &mut got) },
493            HURRAY_ERR_INDEX_OUT_OF_BOUNDS
494        );
495
496        unsafe { hurray_buffer_list_destroy(&mut list) };
497    }
498
499    #[test]
500    fn destroying_the_list_releases_every_owned_buffer_exactly_once() {
501        let releases = AtomicUsize::new(0);
502        let mut a = Aligned([0u8; 128]);
503        let mut b = Aligned([0u8; 128]);
504        let mut c = Aligned([0u8; 128]);
505
506        let mut list: *mut HurrayBufferList = std::ptr::null_mut();
507        unsafe { hurray_buffer_list_new(3, &mut list) };
508        for data in [&mut a, &mut b, &mut c] {
509            let buf = make_buffer(data, 128, &releases);
510            unsafe { hurray_buffer_list_push(list, buf) };
511        }
512
513        unsafe { hurray_buffer_list_destroy(&mut list) };
514        assert_eq!(releases.load(Ordering::SeqCst), 3);
515        assert!(list.is_null());
516    }
517}