Skip to main content

hurray_ffi/
item.rs

1//! A read item: one tensor, or a composite and everything under it.
2//!
3//! `hurray-io` hands a driver a sum type — `StreamItem` and `FileItem` are each a tensor
4//! or a composite whose members are items in turn — and C has no sum type to receive it
5//! in. This module is the answer ADR-041 § Open Questions asked for: an item crosses as
6//! the sequence the **format itself** uses, a head followed by its members, with each
7//! node declaring how many of the nodes after it are its own
8//! ([ADR-027](../../../docs/adr/ADR-027-composite-tensors.md) § Binding).
9//!
10//! So walking a composite from C is a loop over an array, and the nesting is in the
11//! `member_count` column:
12//!
13//! ```text
14//!   index  kind       name    member_count
15//!   0      composite  "w"     2              <- the head, nodes[0]
16//!   1      tensor     "w.0"   0
17//!   2      composite  "w.1"   2              <- a nested head: its members follow it
18//!   3      tensor     "w.1.0" 0
19//!   4      tensor     "w.1.1" 0
20//! ```
21//!
22//! Preorder rather than a tree of pointers, for a reason that outlives taste: a pointer
23//! tree crossing an ABI can be cyclic, and a walker on either side loops forever on one.
24//! An array with member counts cannot describe a cycle at all — every child sits after
25//! its parent.
26//!
27//! Both surfaces produce this: [`hurray_reader_next_item`][crate::hurray_reader_next_item]
28//! from the stream core's assembler, and
29//! [`hurray_file_reader_read_composite`][crate::hurray_file_reader_read_composite] from
30//! the file core's walk. A composite is the same shape whichever one carried it.
31
32use hurray_core::TensorDescriptor;
33use hurray_io::file::FileItem;
34use hurray_io::stream::StreamItem;
35
36use crate::{
37    panic::{catch, null_check},
38    status::{HurrayStatus, HURRAY_ERR_INDEX_OUT_OF_BOUNDS, HURRAY_ERR_INTERNAL, HURRAY_OK},
39    tensor::{hurray_tensor_destroy, tensor_handle, HurrayTensor},
40};
41
42// ── Node kinds ────────────────────────────────────────────────────────────────
43
44/// A plain tensor: it owns buffers and no node after it belongs to it.
45///
46/// Unlike the `HURRAY_KV_*` tags, these two are **not** wire values: the wire has no
47/// tensor-or-composite tag, it has a composite *layout* (tag `0x0B`). The kind is here so
48/// that walking the array needs no descriptor decode, and so that a head with
49/// `member_count == 0` — which the layout permits — is still distinguishable from a
50/// tensor.
51pub const HURRAY_ITEM_TENSOR: u8 = 0;
52
53/// A composite head: it owns no buffers, and the next `member_count` sibling-level nodes
54/// are its members.
55pub const HURRAY_ITEM_COMPOSITE: u8 = 1;
56
57// ── Nodes ─────────────────────────────────────────────────────────────────────
58
59/// One node of an item, as a C caller sees it.
60///
61/// A flat tagged struct rather than a union, for the reason
62/// [`HurrayKvValue`][crate::HurrayKvValue] is one: every binding generator downstream of
63/// this header reads a struct correctly and few read a union well.
64///
65/// Every pointer in it **borrows** from the [`HurrayItem`] it came from and is valid
66/// until [`hurray_item_destroy`]. In particular `tensor` MUST NOT be passed to
67/// [`hurray_tensor_destroy`] — the item owns it.
68#[repr(C)]
69#[derive(Clone, Copy)]
70pub struct HurrayItemNode {
71    /// [`HURRAY_ITEM_TENSOR`] or [`HURRAY_ITEM_COMPOSITE`].
72    pub kind: u8,
73    /// First byte of the node's name, or null when it has none.
74    ///
75    /// A file names every tensor, heads and members alike, because each gets its own
76    /// footer-index entry. A stream names none: the names are not on the wire, so there
77    /// is nothing to borrow and this is null.
78    pub name: *const u8,
79    /// Byte length of `name`; `0` when `name` is null. UTF-8, not NUL-terminated.
80    pub name_len: usize,
81    /// The node's tensor handle, borrowed: its descriptor and its buffers.
82    ///
83    /// A composite head is the buffer-less tensor it is on the wire, so this is non-null
84    /// for every node and its buffer list is empty for a head.
85    pub tensor: *mut HurrayTensor,
86    /// How many of the nodes following this one are its direct members. `0` for a tensor.
87    pub member_count: u32,
88}
89
90impl Default for HurrayItemNode {
91    /// A zeroed node, which is what a C caller gets from `{0}`.
92    fn default() -> Self {
93        Self {
94            kind: HURRAY_ITEM_TENSOR,
95            name: std::ptr::null(),
96            name_len: 0,
97            tensor: std::ptr::null_mut(),
98            member_count: 0,
99        }
100    }
101}
102
103/// One node, with the parts the item owns.
104struct OwnedNode {
105    kind: u8,
106    /// Owned, so the pointer handed to C stays valid for the item's life.
107    name: Option<String>,
108    /// Owned; destroyed by [`hurray_item_destroy`].
109    tensor: *mut HurrayTensor,
110    member_count: u32,
111}
112
113// ── HurrayItem ────────────────────────────────────────────────────────────────
114
115/// Opaque handle to one read item: a tensor, or a composite and its whole subtree.
116///
117/// Not `#[repr(C)]`; the layout is an implementation detail. The handle **owns** every
118/// node's tensor and name, and one [`hurray_item_destroy`] releases all of them.
119///
120/// A plain tensor is an item of exactly one node, so a caller that does not care about
121/// composites can read `hurray_item_node(item, 0, &node)` and ignore the rest.
122pub struct HurrayItem {
123    /// Double-destroy sentinel (debug builds only).
124    #[cfg(debug_assertions)]
125    sentinel: u64,
126    /// Preorder: a parent precedes its members, members are in wire order.
127    nodes: Vec<OwnedNode>,
128}
129
130/// Sentinel placed in `HurrayItem.sentinel` on allocation (debug builds only).
131#[cfg(debug_assertions)]
132const SENTINEL_LIVE: u64 = 0x17E4_0000_17E4_0000;
133
134/// Sentinel written before the allocation is freed, to detect a double destroy in debug
135/// builds.
136#[cfg(debug_assertions)]
137const SENTINEL_DEAD: u64 = 0x17E4_DEAD_17E4_DEAD;
138
139// SAFETY: mirrors the HurrayTensor contract — the C ABI is single-thread-at-a-time per
140// handle, so Send only permits moving the box across threads during construction and
141// destruction, never concurrent access.
142unsafe impl Send for HurrayItem {}
143
144// ── Accessors ─────────────────────────────────────────────────────────────────
145
146/// The item's own kind: [`HURRAY_ITEM_TENSOR`] or [`HURRAY_ITEM_COMPOSITE`].
147///
148/// The kind of node `0`, which a caller that only wants to know whether a composite
149/// arrived can ask for without walking.
150///
151/// # Safety
152///
153/// `handle` and `out_kind` MUST be valid, non-null pointers.
154#[no_mangle]
155pub unsafe extern "C" fn hurray_item_kind(
156    handle: *const HurrayItem,
157    out_kind: *mut u8,
158) -> HurrayStatus {
159    catch(|| {
160        null_check!(handle, out_kind);
161        // SAFETY: handle is non-null and points to a live HurrayItem, which the C ABI's
162        // one-handle-one-thread rule says nothing else is touching.
163        let item = &*handle;
164        // Every constructor pushes the root first, so an item always has a node.
165        let Some(root) = item.nodes.first() else {
166            return HURRAY_ERR_INTERNAL;
167        };
168        *out_kind = root.kind;
169        HURRAY_OK
170    })
171}
172
173/// How many nodes the item has in total, counting the root and every nested member.
174///
175/// `1` for a plain tensor.
176///
177/// # Safety
178///
179/// `handle` and `out_count` MUST be valid, non-null pointers.
180#[no_mangle]
181pub unsafe extern "C" fn hurray_item_node_count(
182    handle: *const HurrayItem,
183    out_count: *mut u64,
184) -> HurrayStatus {
185    catch(|| {
186        null_check!(handle, out_count);
187        // SAFETY: handle is non-null and points to a live HurrayItem.
188        let item = &*handle;
189        *out_count = item.nodes.len() as u64;
190        HURRAY_OK
191    })
192}
193
194/// Copies node `index` into `out_node`.
195///
196/// Nodes are in preorder: node `0` is the root, and a node's `member_count` direct members
197/// begin at the next index, each followed immediately by its own subtree. The pointers
198/// inside the node borrow from the item and are valid until [`hurray_item_destroy`].
199///
200/// Returns [`HURRAY_ERR_INDEX_OUT_OF_BOUNDS`] if `index` is past the end.
201///
202/// # Safety
203///
204/// `handle` MUST point to a live [`HurrayItem`], and `out_node` MUST be a valid, non-null
205/// pointer writable for one [`HurrayItemNode`].
206#[no_mangle]
207pub unsafe extern "C" fn hurray_item_node(
208    handle: *const HurrayItem,
209    index: u64,
210    out_node: *mut HurrayItemNode,
211) -> HurrayStatus {
212    catch(|| {
213        null_check!(handle, out_node);
214        // SAFETY: handle is non-null and points to a live HurrayItem.
215        let item = &*handle;
216        let Some(node) = item.nodes.get(index as usize) else {
217            return HURRAY_ERR_INDEX_OUT_OF_BOUNDS;
218        };
219        *out_node = HurrayItemNode {
220            kind: node.kind,
221            name: node.name.as_ref().map_or(std::ptr::null(), |n| n.as_ptr()),
222            name_len: node.name.as_ref().map_or(0, |n| n.len()),
223            tensor: node.tensor,
224            member_count: node.member_count,
225        };
226        HURRAY_OK
227    })
228}
229
230// ── Destructor ────────────────────────────────────────────────────────────────
231
232/// Destroys an item, every tensor and name it owns, then writes null through `handle`.
233///
234/// One call for a whole composite subtree: the nodes' `tensor` pointers were borrowed and
235/// are released here, so a caller that destroyed them individually would double-free.
236///
237/// # Safety
238///
239/// - `handle` MUST be a valid, non-null pointer to a `*mut HurrayItem` produced by
240///   [`hurray_reader_next_item`][crate::hurray_reader_next_item] or
241///   [`hurray_file_reader_read_composite`][crate::hurray_file_reader_read_composite].
242/// - That item MUST NOT have been destroyed already.
243#[no_mangle]
244pub unsafe extern "C" fn hurray_item_destroy(handle: *mut *mut HurrayItem) -> HurrayStatus {
245    catch(|| {
246        null_check!(handle);
247        let item = *handle;
248        if item.is_null() {
249            return HURRAY_ERR_INTERNAL;
250        }
251
252        #[cfg(debug_assertions)]
253        {
254            // SAFETY: item is non-null and points to a live HurrayItem. Read and mark
255            // through the raw pointer before Box::from_raw, so the dead marker lands in
256            // the heap slot a second destroy would inspect.
257            if (*item).sentinel == SENTINEL_DEAD {
258                return HURRAY_ERR_INTERNAL;
259            }
260            (*item).sentinel = SENTINEL_DEAD;
261        }
262
263        // SAFETY: produced by `item_handle` via Box::into_raw; the caller guarantees this
264        // is the first and only destroy call.
265        let mut boxed = Box::from_raw(item);
266        for node in boxed.nodes.iter_mut() {
267            let tensor = std::mem::replace(&mut node.tensor, std::ptr::null_mut());
268            if !tensor.is_null() {
269                // SAFETY: owned by this item and destroyed exactly once; nulled first so
270                // a re-entrant destroy cannot free it twice.
271                let mut tensor = tensor;
272                hurray_tensor_destroy(&mut tensor);
273            }
274        }
275
276        drop(boxed);
277        *handle = std::ptr::null_mut();
278        HURRAY_OK
279    })
280}
281
282// ── Handing a decoded item to C ───────────────────────────────────────────────
283
284/// Builds one node, taking ownership of a decoded tensor.
285fn owned_node(
286    kind: u8,
287    name: Option<String>,
288    descriptor: TensorDescriptor,
289    buffers: Vec<bytes::Bytes>,
290    member_count: u32,
291) -> OwnedNode {
292    OwnedNode {
293        kind,
294        name,
295        tensor: Box::into_raw(tensor_handle(descriptor, buffers)),
296        member_count,
297    }
298}
299
300/// Flattens a decoded stream item into the preorder array a C caller walks.
301///
302/// Nameless: a stream carries no tensor names.
303pub(crate) fn item_handle(item: StreamItem) -> Box<HurrayItem> {
304    let mut nodes = Vec::new();
305    // An explicit stack, not recursion (ADR-041 § 10) — a composite's depth is bounded by
306    // the reader's own limit, but this traversal needs no stack frame per level to say so.
307    // Members are pushed in reverse so the LIFO yields them in wire order.
308    let mut stack = vec![item];
309    while let Some(next) = stack.pop() {
310        match next {
311            StreamItem::Tensor(tensor) => nodes.push(owned_node(
312                HURRAY_ITEM_TENSOR,
313                None,
314                tensor.descriptor,
315                tensor.buffers,
316                0,
317            )),
318            StreamItem::Composite(composite) => {
319                // From the members present, not from the head's declared `member_count`:
320                // this number describes the array that follows, and the assembler has
321                // already refused any group where the two disagree.
322                nodes.push(owned_node(
323                    HURRAY_ITEM_COMPOSITE,
324                    None,
325                    composite.head,
326                    Vec::new(),
327                    composite.members.len() as u32,
328                ));
329                stack.extend(composite.members.into_iter().rev());
330            }
331        }
332    }
333    Box::new(HurrayItem {
334        #[cfg(debug_assertions)]
335        sentinel: SENTINEL_LIVE,
336        nodes,
337    })
338}
339
340/// Flattens a decoded file item into the preorder array a C caller walks.
341///
342/// Named: every tensor in a file has a footer-index entry, heads and members alike.
343///
344/// Deliberately not shared with [`item_handle`]: `StreamItem` and `FileItem` are distinct
345/// types with no common trait upstream, and the names are not an optional field on one of
346/// them but absent from the wire entirely. A local trait to join them would be more
347/// surface than the fifteen lines it saves — step 4b's conclusion for the same pair of
348/// write ops.
349pub(crate) fn file_item_handle(item: FileItem) -> Box<HurrayItem> {
350    let mut nodes = Vec::new();
351    let mut stack = vec![item];
352    while let Some(next) = stack.pop() {
353        match next {
354            FileItem::Tensor(tensor) => nodes.push(owned_node(
355                HURRAY_ITEM_TENSOR,
356                Some(tensor.name),
357                tensor.descriptor,
358                tensor.buffers,
359                0,
360            )),
361            FileItem::Composite(composite) => {
362                nodes.push(owned_node(
363                    HURRAY_ITEM_COMPOSITE,
364                    Some(composite.name),
365                    composite.head,
366                    Vec::new(),
367                    composite.members.len() as u32,
368                ));
369                stack.extend(composite.members.into_iter().rev());
370            }
371        }
372    }
373    Box::new(HurrayItem {
374        #[cfg(debug_assertions)]
375        sentinel: SENTINEL_LIVE,
376        nodes,
377    })
378}