1use std::any::Any;
14use std::cell::RefCell;
15use std::collections::BTreeMap;
16use std::ffi::{c_void, CString};
17use std::panic::{catch_unwind, AssertUnwindSafe};
18use std::sync::atomic::{AtomicBool, AtomicI64, Ordering};
19use std::sync::Arc;
20
21use crate::metadata::{Metadata, MetadataValue};
22use crate::status::{Error, Result};
23use crate::timestamp::Timestamp;
24
25#[derive(Clone, Copy, Default)]
26pub(crate) struct CowCopyCounters {
27 pub copies: u64,
28 pub bytes: u64,
29}
30
31thread_local! {
32 static COW_COPY_SCOPES: RefCell<Vec<CowCopyCounters>> = const { RefCell::new(Vec::new()) };
33}
34
35pub(crate) fn begin_cow_copy_scope() {
36 COW_COPY_SCOPES.with(|scopes| scopes.borrow_mut().push(CowCopyCounters::default()));
37}
38
39pub(crate) fn end_cow_copy_scope() -> CowCopyCounters {
40 COW_COPY_SCOPES.with(|scopes| scopes.borrow_mut().pop().unwrap_or_default())
41}
42
43fn record_cow_copy(bytes: u64) {
44 COW_COPY_SCOPES.with(|scopes| {
45 let mut scopes = scopes.borrow_mut();
46 let Some(counters) = scopes.last_mut() else {
47 return;
48 };
49 counters.copies = counters.copies.saturating_add(1);
50 counters.bytes = counters.bytes.saturating_add(bytes);
51 });
52}
53
54pub mod type_id {
56 pub const NONE: u64 = 0;
57 pub const BYTES: u64 = 1;
58 pub const I64: u64 = 2;
59 pub const F64: u64 = 3;
60 pub const BOOL: u64 = 4;
61 pub const STR: u64 = 5;
62 pub const BUFFER: u64 = 6;
63 pub const HOST_OBJECT: u64 = 7; }
65
66pub mod dtype {
68 pub const U8: i32 = 0;
69 pub const I8: i32 = 1;
70 pub const U16: i32 = 2;
71 pub const I16: i32 = 3;
72 pub const I32: i32 = 4;
73 pub const I64: i32 = 5;
74 pub const F16: i32 = 6;
75 pub const F32: i32 = 7;
76 pub const F64: i32 = 8;
77}
78
79pub const MAX_DIMS: usize = 8;
80
81pub unsafe trait InteropType: Any + Send + Sync {
97 const TYPE_NAME: &'static str;
98}
99
100#[derive(Clone, Debug, Eq, PartialEq)]
102pub struct TypeDescriptor {
103 pub name: String,
104 pub size: usize,
105 pub align: usize,
106}
107
108pub fn dtype_size(dt: i32) -> usize {
110 match dt {
111 dtype::U8 | dtype::I8 => 1,
112 dtype::U16 | dtype::I16 | dtype::F16 => 2,
113 dtype::I32 | dtype::F32 => 4,
114 dtype::I64 | dtype::F64 => 8,
115 _ => 0,
116 }
117}
118
119pub fn fnv1a_type_id(mangled_name: &str) -> u64 {
129 let mut h: u64 = 14695981039346656037;
130 for b in mangled_name.bytes() {
131 h ^= b as u64;
132 h = h.wrapping_mul(1099511628211);
133 }
134 if h < 16 {
136 h + 16
137 } else {
138 h
139 }
140}
141
142#[derive(Default)]
143struct TypeRegistry {
144 by_id: BTreeMap<u64, TypeDescriptor>,
145 by_name: BTreeMap<String, u64>,
146}
147
148static TYPE_REGISTRY: std::sync::LazyLock<std::sync::Mutex<TypeRegistry>> =
149 std::sync::LazyLock::new(|| std::sync::Mutex::new(TypeRegistry::default()));
150
151pub fn register_type_descriptor(id: u64, name: &str, size: usize, align: usize) -> Result<()> {
157 if id < 16 {
158 return Err(Error::InvalidArg(format!(
159 "custom type id {id} is reserved; custom ids must be >= 16"
160 )));
161 }
162 if name.is_empty() {
163 return Err(Error::InvalidArg(
164 "custom type name must not be empty".into(),
165 ));
166 }
167 if size == 0 || align == 0 || !align.is_power_of_two() {
168 return Err(Error::InvalidArg(format!(
169 "custom type `{name}` has invalid layout: size={size}, align={align}"
170 )));
171 }
172 let expected_id = fnv1a_type_id(name);
173 if id != expected_id {
174 return Err(Error::InvalidArg(format!(
175 "custom type `{name}` must use its stable-name id {expected_id}, got {id}"
176 )));
177 }
178
179 let descriptor = TypeDescriptor {
180 name: name.to_string(),
181 size,
182 align,
183 };
184 let mut registry = TYPE_REGISTRY.lock().unwrap_or_else(|e| e.into_inner());
185 if let Some(existing_id) = registry.by_name.get(name) {
186 if *existing_id != id {
187 return Err(Error::InvalidArg(format!(
188 "custom type name `{name}` is already registered with id {existing_id}, \
189 cannot reuse it with id {id}"
190 )));
191 }
192 }
193 if let Some(existing) = registry.by_id.get(&id) {
194 if existing == &descriptor {
195 return Ok(());
196 }
197 return Err(Error::InvalidArg(format!(
198 "custom type id {id} is already registered as `{}` (size={}, align={}), \
199 cannot register `{name}` (size={size}, align={align})",
200 existing.name, existing.size, existing.align
201 )));
202 }
203 registry.by_name.insert(name.to_string(), id);
204 registry.by_id.insert(id, descriptor);
205 Ok(())
206}
207
208pub fn type_descriptor(id: u64) -> Option<TypeDescriptor> {
209 TYPE_REGISTRY
210 .lock()
211 .unwrap_or_else(|e| e.into_inner())
212 .by_id
213 .get(&id)
214 .cloned()
215}
216
217pub(crate) fn validate_type_id(id: u64) -> Result<()> {
223 match id {
224 type_id::NONE..=type_id::BUFFER => Ok(()),
225 type_id::HOST_OBJECT => Err(Error::InvalidArg(
226 "LMFLOW_TYPE_HOST_OBJECT is reserved and not enabled (see ADR #26); use \
227 LMFLOW_TYPE_BUFFER for numeric collections, or LMFLOW_TYPE_STR carrying JSON for \
228 arbitrary metadata"
229 .into(),
230 )),
231 8..=15 => Err(Error::InvalidArg(format!(
232 "type id {id} is reserved; built-in ids currently end at LMFLOW_TYPE_HOST_OBJECT \
233 (7), and custom type ids must be >= 16"
234 ))),
235 _ if type_descriptor(id).is_none() => Err(Error::InvalidArg(format!(
236 "unregistered custom type id {id}; register its stable name, size, and alignment \
237 with lmflow_register_type_descriptor before using it in a graph"
238 ))),
239 _ => Ok(()),
240 }
241}
242
243pub fn type_name(id: u64) -> String {
245 match id {
246 type_id::NONE => "any/undeclared".to_string(),
247 type_id::BYTES => "Bytes".to_string(),
248 type_id::I64 => "I64".to_string(),
249 type_id::F64 => "F64".to_string(),
250 type_id::BOOL => "Bool".to_string(),
251 type_id::STR => "Str".to_string(),
252 type_id::BUFFER => "Buffer".to_string(),
253 type_id::HOST_OBJECT => "HostObject".to_string(),
254 other => TYPE_REGISTRY
255 .lock()
256 .unwrap_or_else(|e| e.into_inner())
257 .by_id
258 .get(&other)
259 .map(|descriptor| descriptor.name.clone())
260 .unwrap_or_else(|| format!("type#{other}")),
261 }
262}
263
264#[derive(Clone, Debug)]
266pub struct BufferData {
267 pub bytes: Vec<u8>,
268 pub shape: [i64; MAX_DIMS],
269 pub strides: [i64; MAX_DIMS],
270 pub ndim: i32,
271 pub dtype: i32,
272}
273
274impl BufferData {
275 pub(crate) fn checked_layout(
276 shape: &[i64],
277 dt: i32,
278 ) -> Result<([i64; MAX_DIMS], [i64; MAX_DIMS], usize)> {
279 if shape.is_empty() || shape.len() > MAX_DIMS {
280 return Err(Error::InvalidArg(format!(
281 "ndim must be in 1..={}, got {}",
282 MAX_DIMS,
283 shape.len()
284 )));
285 }
286 let esz = dtype_size(dt);
287 if esz == 0 {
288 return Err(Error::InvalidArg(format!("unknown dtype {dt}")));
289 }
290 let mut count: i64 = 1;
291 for &d in shape {
292 if d < 0 {
293 return Err(Error::InvalidArg("shape must not be negative".into()));
294 }
295 count = count
296 .checked_mul(d)
297 .ok_or_else(|| Error::InvalidArg("shape product overflow".into()))?;
298 }
299 let count = usize::try_from(count)
300 .map_err(|_| Error::InvalidArg("shape product exceeds addressable memory".into()))?;
301 let total = count
302 .checked_mul(esz)
303 .ok_or_else(|| Error::InvalidArg("buffer byte count overflow".into()))?;
304 let mut s = [0i64; MAX_DIMS];
305 let mut st = [0i64; MAX_DIMS];
306 s[..shape.len()].copy_from_slice(shape);
307 let mut acc = esz as i64;
308 for i in (0..shape.len()).rev() {
309 st[i] = acc;
310 acc = acc
311 .checked_mul(shape[i])
312 .ok_or_else(|| Error::InvalidArg("buffer stride overflow".into()))?;
313 }
314 Ok((s, st, total))
315 }
316
317 pub fn new(shape: &[i64], dt: i32) -> Result<Self> {
319 let (s, st, total) = Self::checked_layout(shape, dt)?;
320 let mut bytes = Vec::new();
321 bytes
322 .try_reserve_exact(total)
323 .map_err(|_| Error::InvalidArg(format!("cannot allocate {total} buffer bytes")))?;
324 bytes.resize(total, 0);
325 Ok(Self {
326 bytes,
327 shape: s,
328 strides: st,
329 ndim: shape.len() as i32,
330 dtype: dt,
331 })
332 }
333
334 pub fn shape_slice(&self) -> &[i64] {
335 &self.shape[..self.ndim as usize]
336 }
337}
338
339pub(crate) type BufferReleaseFn = unsafe extern "C" fn(*mut c_void);
340
341pub(crate) struct ExternalBuffer {
343 pub(crate) data: *mut c_void,
344 pub(crate) shape: [i64; MAX_DIMS],
345 pub(crate) strides: [i64; MAX_DIMS],
346 pub(crate) ndim: i32,
347 pub(crate) dtype: i32,
348 pub(crate) readonly: bool,
349 pub(crate) element_count: usize,
350 pub(crate) byte_size: u64,
351 pub(crate) release_fn: BufferReleaseFn,
352 pub(crate) user_data: *mut c_void,
353}
354
355impl ExternalBuffer {
356 pub fn byte_size(&self) -> u64 {
357 self.byte_size
358 }
359
360 fn packet_payload_ptr(&self) -> *mut c_void {
361 if self.data.is_null() {
362 std::ptr::NonNull::<u8>::dangling().as_ptr() as *mut c_void
363 } else {
364 self.data
365 }
366 }
367
368 fn copy_to_owned(&self) -> Result<BufferData> {
369 let dims = &self.shape[..self.ndim as usize];
370 let mut owned = BufferData::new(dims, self.dtype)?;
371 if self.element_count == 0 {
372 return Ok(owned);
373 }
374
375 let element_size = dtype_size(self.dtype);
376 let ndim = dims.len();
377 let mut index = vec![0i64; ndim];
378 for element in 0..self.element_count {
379 let offset = index
380 .iter()
381 .zip(self.strides)
382 .take(ndim)
383 .try_fold(0i128, |offset, (&coordinate, stride)| {
384 i128::from(coordinate)
385 .checked_mul(i128::from(stride))
386 .and_then(|term| offset.checked_add(term))
387 })
388 .ok_or_else(|| Error::InvalidArg("buffer element offset overflow".into()))?;
389 let offset = isize::try_from(offset)
390 .map_err(|_| Error::InvalidArg("buffer element offset overflow".into()))?;
391 let source = unsafe { (self.data as *const u8).offset(offset) };
392 let destination = unsafe { owned.bytes.as_mut_ptr().add(element * element_size) };
393 unsafe { std::ptr::copy_nonoverlapping(source, destination, element_size) };
394
395 for axis in (0..ndim).rev() {
396 index[axis] += 1;
397 if index[axis] < dims[axis] {
398 break;
399 }
400 index[axis] = 0;
401 }
402 }
403 Ok(owned)
404 }
405}
406
407impl Drop for ExternalBuffer {
408 fn drop(&mut self) {
409 unsafe { (self.release_fn)(self.user_data) };
410 self.user_data = std::ptr::null_mut();
411 self.data = std::ptr::null_mut();
412 }
413}
414
415impl std::fmt::Debug for ExternalBuffer {
416 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
417 formatter
418 .debug_struct("ExternalBuffer")
419 .field("data", &self.data)
420 .field("shape", &&self.shape[..self.ndim as usize])
421 .field("strides", &&self.strides[..self.ndim as usize])
422 .field("dtype", &self.dtype)
423 .field("readonly", &self.readonly)
424 .finish()
425 }
426}
427
428unsafe impl Send for ExternalBuffer {}
430unsafe impl Sync for ExternalBuffer {}
431
432#[derive(Clone, Copy)]
433pub(crate) struct BufferView {
434 pub data: *mut c_void,
435 pub shape: [i64; MAX_DIMS],
436 pub strides: [i64; MAX_DIMS],
437 pub ndim: i32,
438 pub dtype: i32,
439 pub readonly: bool,
440}
441
442#[derive(Clone, Debug)]
443pub enum Builtin {
444 Bytes(Vec<u8>),
445 I64(i64),
446 F64(f64),
447 Bool(bool),
448 Str(CString),
449 Buffer(BufferData),
450}
451
452impl Builtin {
453 pub fn type_id(&self) -> u64 {
454 match self {
455 Builtin::Bytes(_) => type_id::BYTES,
456 Builtin::I64(_) => type_id::I64,
457 Builtin::F64(_) => type_id::F64,
458 Builtin::Bool(_) => type_id::BOOL,
459 Builtin::Str(_) => type_id::STR,
460 Builtin::Buffer(_) => type_id::BUFFER,
461 }
462 }
463
464 pub fn byte_size(&self) -> u64 {
466 match self {
467 Builtin::Bytes(v) => v.len() as u64,
468 Builtin::I64(_) | Builtin::F64(_) => 8,
469 Builtin::Bool(_) => 1,
470 Builtin::Str(s) => s.as_bytes().len() as u64,
471 Builtin::Buffer(b) => b.bytes.len() as u64,
472 }
473 }
474
475 fn data_ptr(&self) -> *mut c_void {
477 match self {
478 Builtin::Bytes(v) => v.as_ptr() as *mut c_void,
479 Builtin::I64(v) => v as *const i64 as *mut c_void,
480 Builtin::F64(v) => v as *const f64 as *mut c_void,
481 Builtin::Bool(v) => v as *const bool as *mut c_void,
482 Builtin::Str(s) => s.as_ptr() as *mut c_void,
483 Builtin::Buffer(b) => b.bytes.as_ptr() as *mut c_void,
484 }
485 }
486}
487
488pub struct Foreign {
490 pub ptr: *mut c_void,
491 pub drop_fn: Option<unsafe extern "C" fn(*mut c_void)>,
492 pub type_id: u64,
493 pub byte_size: u64,
494}
495
496impl Drop for Foreign {
497 fn drop(&mut self) {
498 if let Some(f) = self.drop_fn {
499 if !self.ptr.is_null() {
500 unsafe { f(self.ptr) };
502 }
503 }
504 self.ptr = std::ptr::null_mut();
505 }
506}
507
508impl std::fmt::Debug for Foreign {
509 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
510 write!(f, "Foreign{{ptr:{:?}, type_id:{}}}", self.ptr, self.type_id)
511 }
512}
513
514unsafe impl Send for Foreign {}
518unsafe impl Sync for Foreign {}
519
520pub enum Payload {
521 Native(Box<dyn Any + Send + Sync>),
522 Builtin(Builtin),
523 Foreign(Foreign),
524}
525
526impl Payload {
527 pub fn type_id(&self) -> u64 {
528 match self {
529 Payload::Native(value) => {
531 if value.is::<ExternalBuffer>() {
532 type_id::BUFFER
533 } else {
534 type_id::NONE
535 }
536 }
537 Payload::Builtin(b) => b.type_id(),
538 Payload::Foreign(f) => f.type_id,
539 }
540 }
541
542 pub fn byte_size(&self) -> u64 {
543 match self {
544 Payload::Native(value) => value
545 .downcast_ref::<ExternalBuffer>()
546 .map_or(0, ExternalBuffer::byte_size),
547 Payload::Foreign(f) => f.byte_size,
548 Payload::Builtin(b) => b.byte_size(),
549 }
550 }
551
552 pub fn data_ptr(&self) -> *mut c_void {
554 match self {
555 Payload::Native(b) => {
556 if let Some(buffer) = b.downcast_ref::<ExternalBuffer>() {
557 return buffer.packet_payload_ptr();
558 }
559 let p: *const dyn Any = &**b;
561 p as *const () as *mut c_void
562 }
563 Payload::Builtin(b) => b.data_ptr(),
564 Payload::Foreign(f) => f.ptr,
565 }
566 }
567
568 pub(crate) fn buffer_view(&self) -> Option<BufferView> {
569 match self {
570 Payload::Builtin(Builtin::Buffer(buffer)) => Some(BufferView {
571 data: buffer.bytes.as_ptr() as *mut c_void,
572 shape: buffer.shape,
573 strides: buffer.strides,
574 ndim: buffer.ndim,
575 dtype: buffer.dtype,
576 readonly: true,
577 }),
578 Payload::Native(value) => {
579 value
580 .downcast_ref::<ExternalBuffer>()
581 .map(|buffer| BufferView {
582 data: buffer.data,
583 shape: buffer.shape,
584 strides: buffer.strides,
585 ndim: buffer.ndim,
586 dtype: buffer.dtype,
587 readonly: true,
588 })
589 }
590 _ => None,
591 }
592 }
593}
594
595impl std::fmt::Debug for Payload {
596 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
597 match self {
598 Payload::Native(value) => match value.downcast_ref::<ExternalBuffer>() {
599 Some(buffer) => write!(f, "{buffer:?}"),
600 None => f.write_str("Native"),
601 },
602 Payload::Builtin(b) => write!(f, "{b:?}"),
603 Payload::Foreign(x) => write!(f, "{x:?}"),
604 }
605 }
606}
607
608pub(crate) struct PacketBody {
615 pub payload: Arc<Payload>,
616 pub metadata: Arc<Metadata>,
617 pub(crate) ingress_us: AtomicI64,
618 pub(crate) e2e_observed: AtomicBool,
619}
620
621impl Clone for PacketBody {
622 fn clone(&self) -> Self {
623 Self {
624 payload: self.payload.clone(),
625 metadata: self.metadata.clone(),
626 ingress_us: AtomicI64::new(self.ingress_us.load(Ordering::Relaxed)),
627 e2e_observed: AtomicBool::new(self.e2e_observed.load(Ordering::Relaxed)),
628 }
629 }
630}
631
632impl PacketBody {
633 fn new(payload: Payload, metadata: Metadata) -> Self {
634 Self {
635 payload: Arc::new(payload),
636 metadata: Arc::new(metadata),
637 ingress_us: AtomicI64::new(0),
638 e2e_observed: AtomicBool::new(false),
639 }
640 }
641}
642
643static EMPTY_METADATA: std::sync::LazyLock<Metadata> = std::sync::LazyLock::new(Metadata::new);
644
645#[derive(Clone, Default)]
646pub struct Packet {
647 data: Option<Arc<PacketBody>>,
648 ts: Timestamp,
649}
650
651impl Packet {
652 pub(crate) fn mark_ingress_us(&self, now_us: i64) {
653 if let Some(body) = self.data.as_deref() {
654 let _ = body.ingress_us.compare_exchange(
655 0,
656 now_us.max(1),
657 Ordering::Relaxed,
658 Ordering::Relaxed,
659 );
660 }
661 }
662
663 pub(crate) fn rebase_ingress_us(&mut self, now_us: i64) {
664 let Some(body) = self.data.as_deref() else {
665 return;
666 };
667 self.data = Some(Arc::new(PacketBody {
668 payload: body.payload.clone(),
669 metadata: body.metadata.clone(),
670 ingress_us: AtomicI64::new(now_us.max(1)),
671 e2e_observed: AtomicBool::new(false),
672 }));
673 }
674
675 pub(crate) fn ingress_us(&self) -> i64 {
676 self.data
677 .as_deref()
678 .map_or(0, |body| body.ingress_us.load(Ordering::Relaxed))
679 }
680
681 pub(crate) fn mark_e2e_observed(&self) -> bool {
682 self.data.as_deref().is_some_and(|body| {
683 body.e2e_observed
684 .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
685 .is_ok()
686 })
687 }
688
689 pub fn empty() -> Self {
691 Self::default()
692 }
693
694 pub fn new<T: Any + Send + Sync>(value: T) -> Self {
716 Self {
717 data: Some(Arc::new(PacketBody::new(
718 Payload::Native(Box::new(value)),
719 Metadata::new(),
720 ))),
721 ts: Timestamp::unset(),
722 }
723 }
724
725 pub fn from_interop<T: InteropType>(value: T) -> Self {
727 let type_id = fnv1a_type_id(T::TYPE_NAME);
728 register_type_descriptor(
729 type_id,
730 T::TYPE_NAME,
731 std::mem::size_of::<T>(),
732 std::mem::align_of::<T>(),
733 )
734 .expect("conflicting InteropType declaration");
735 unsafe { Self::new_interop(value, type_id) }
737 }
738
739 pub unsafe fn new_interop<T: Any + Send + Sync>(value: T, type_id: u64) -> Self {
754 assert!(
755 type_id >= 16,
756 "custom interop type_id must be >= 16; built-in ids 0..=15 are reserved"
757 );
758 let descriptor = type_descriptor(type_id).unwrap_or_else(|| {
759 panic!(
760 "custom interop type_id {type_id} is not registered; register its stable name, \
761 size, and alignment before constructing a packet"
762 )
763 });
764 assert_eq!(
765 descriptor.size,
766 std::mem::size_of::<T>(),
767 "custom interop type `{}` size mismatch: registered {}, Rust value {}",
768 descriptor.name,
769 descriptor.size,
770 std::mem::size_of::<T>(),
771 );
772 assert_eq!(
773 descriptor.align,
774 std::mem::align_of::<T>(),
775 "custom interop type `{}` alignment mismatch: registered {}, Rust value {}",
776 descriptor.name,
777 descriptor.align,
778 std::mem::align_of::<T>(),
779 );
780 let boxed = Box::new(value);
781 let ptr = Box::into_raw(boxed) as *mut c_void;
782 unsafe extern "C" fn drop_boxed<T>(p: *mut c_void) {
783 let result = catch_unwind(AssertUnwindSafe(|| {
785 drop(unsafe { Box::from_raw(p as *mut T) })
786 }));
787 if result.is_err() {
788 crate::runtime::log_warn(&format!(
789 "Rust interop payload `{}` panicked during destruction (ignored)",
790 std::any::type_name::<T>()
791 ));
792 }
793 }
794 Self {
795 data: Some(Arc::new(PacketBody::new(
796 Payload::Foreign(Foreign {
797 ptr,
798 drop_fn: Some(drop_boxed::<T>),
799 type_id,
800 byte_size: descriptor.size as u64,
801 }),
802 Metadata::new(),
803 ))),
804 ts: Timestamp::unset(),
805 }
806 }
807
808 pub fn from_i64(v: i64) -> Self {
810 Self::from_builtin(Builtin::I64(v))
811 }
812 pub fn from_f64(v: f64) -> Self {
813 Self::from_builtin(Builtin::F64(v))
814 }
815 pub fn from_bool(v: bool) -> Self {
816 Self::from_builtin(Builtin::Bool(v))
817 }
818 pub fn from_bytes(v: Vec<u8>) -> Self {
819 Self::from_builtin(Builtin::Bytes(v))
820 }
821 pub fn from_string(v: &str) -> Self {
824 Self::from_builtin(Builtin::Str(std::ffi::CString::new(v).unwrap_or_default()))
825 }
826 pub fn as_i64(&self) -> Option<i64> {
828 match self.as_builtin()? {
829 Builtin::I64(v) => Some(*v),
830 _ => None,
831 }
832 }
833 pub fn as_f64(&self) -> Option<f64> {
834 match self.as_builtin()? {
835 Builtin::F64(v) => Some(*v),
836 _ => None,
837 }
838 }
839 pub fn as_str(&self) -> Option<&str> {
840 match self.as_builtin()? {
841 Builtin::Str(s) => s.to_str().ok(),
842 _ => None,
843 }
844 }
845
846 pub fn from_builtin(b: Builtin) -> Self {
847 Self {
848 data: Some(Arc::new(PacketBody::new(
849 Payload::Builtin(b),
850 Metadata::new(),
851 ))),
852 ts: Timestamp::unset(),
853 }
854 }
855
856 pub(crate) unsafe fn from_external_buffer(buffer: ExternalBuffer) -> Self {
861 Self {
862 data: Some(Arc::new(PacketBody::new(
863 Payload::Native(Box::new(buffer)),
864 Metadata::new(),
865 ))),
866 ts: Timestamp::unset(),
867 }
868 }
869
870 pub unsafe fn from_foreign(
875 ptr: *mut c_void,
876 type_id: u64,
877 drop_fn: Option<unsafe extern "C" fn(*mut c_void)>,
878 ) -> Self {
879 Self {
880 data: Some(Arc::new(PacketBody::new(
881 Payload::Foreign(Foreign {
882 ptr,
883 drop_fn,
884 type_id,
885 byte_size: type_descriptor(type_id)
886 .map_or(0, |descriptor| descriptor.size as u64),
887 }),
888 Metadata::new(),
889 ))),
890 ts: Timestamp::unset(),
891 }
892 }
893
894 pub fn at(mut self, ts: Timestamp) -> Self {
895 self.ts = ts;
896 self
897 }
898 pub fn set_timestamp(&mut self, ts: Timestamp) {
899 self.ts = ts;
900 }
901 pub fn timestamp(&self) -> Timestamp {
902 self.ts
903 }
904 pub fn is_empty(&self) -> bool {
905 self.data.is_none()
906 }
907 pub fn type_id(&self) -> u64 {
908 self.data.as_deref().map_or(type_id::NONE, |body| {
909 Payload::type_id(body.payload.as_ref())
910 })
911 }
912 pub fn byte_size(&self) -> u64 {
913 self.data
914 .as_deref()
915 .map_or(0, |body| body.payload.byte_size())
916 }
917 pub fn payload(&self) -> Option<&Payload> {
918 self.data.as_deref().map(|body| body.payload.as_ref())
919 }
920
921 pub fn ref_count(&self) -> usize {
923 self.data.as_ref().map_or(0, Arc::strong_count)
924 }
925
926 pub(crate) fn arc_ref(&self) -> Option<&Arc<PacketBody>> {
928 self.data.as_ref()
929 }
930
931 pub(crate) fn into_arc(self) -> Option<Arc<PacketBody>> {
933 self.data
934 }
935
936 pub(crate) fn from_arc(arc: Arc<PacketBody>, ts: Timestamp) -> Self {
938 Self {
939 data: Some(arc),
940 ts,
941 }
942 }
943
944 pub fn get<T: Any>(&self) -> Option<&T> {
946 match self.data.as_deref()?.payload.as_ref() {
947 Payload::Native(b) => b.downcast_ref::<T>(),
948 _ => None,
949 }
950 }
951
952 pub fn foreign_ptr(&self) -> Option<*const c_void> {
957 match self.data.as_deref()?.payload.as_ref() {
958 Payload::Foreign(f) => Some(f.ptr as *const c_void),
959 _ => None,
960 }
961 }
962
963 pub fn as_builtin(&self) -> Option<&Builtin> {
964 match self.data.as_deref()?.payload.as_ref() {
965 Payload::Builtin(b) => Some(b),
966 _ => None,
967 }
968 }
969
970 pub fn metadata(&self) -> &Metadata {
971 self.data
972 .as_deref()
973 .map_or(&EMPTY_METADATA, |body| body.metadata.as_ref())
974 }
975
976 pub fn metadata_value(&self, key: &str) -> Option<&MetadataValue> {
977 self.metadata().get(key)
978 }
979
980 pub fn has_metadata(&self, key: &str) -> bool {
981 self.metadata().contains_key(key)
982 }
983
984 pub fn metadata_keys(&self) -> impl Iterator<Item = &str> {
985 self.metadata().keys().map(String::as_str)
986 }
987
988 pub fn with_metadata<V: Into<MetadataValue>>(
989 mut self,
990 key: impl Into<String>,
991 value: V,
992 ) -> Self {
993 self.set_metadata(key, value);
994 self
995 }
996
997 pub fn set_metadata<V: Into<MetadataValue>>(&mut self, key: impl Into<String>, value: V) {
998 let Some(body) = self.data.as_mut() else {
999 return;
1000 };
1001 let body = Arc::make_mut(body);
1002 let metadata = Arc::make_mut(&mut body.metadata);
1003 metadata.insert(key.into(), value.into());
1004 }
1005
1006 pub fn remove_metadata(&mut self, key: &str) -> Option<MetadataValue> {
1007 let body = Arc::make_mut(self.data.as_mut()?);
1008 Arc::make_mut(&mut body.metadata).remove(key)
1009 }
1010
1011 pub(crate) fn make_mutable_buffer(&mut self) -> Result<BufferView> {
1012 let body = Arc::make_mut(self.data.as_mut().ok_or_else(|| {
1013 Error::InvalidArg("cannot get a writable view of an empty packet".into())
1014 })?);
1015 let arc = &mut body.payload;
1016 let external = match arc.as_ref() {
1017 Payload::Native(value) => value.downcast_ref::<ExternalBuffer>(),
1018 _ => None,
1019 };
1020 let must_copy =
1021 Arc::strong_count(arc) > 1 || external.is_some_and(|buffer| buffer.readonly);
1022 if must_copy {
1023 let owned = match arc.as_ref() {
1024 Payload::Builtin(Builtin::Buffer(buffer)) => {
1025 record_cow_copy(buffer.bytes.len() as u64);
1026 buffer.clone()
1027 }
1028 Payload::Native(value) if value.is::<ExternalBuffer>() => {
1029 let buffer = value
1030 .downcast_ref::<ExternalBuffer>()
1031 .expect("type checked above");
1032 record_cow_copy(buffer.byte_size());
1033 buffer.copy_to_owned()?
1034 }
1035 _ => {
1036 return Err(Error::InvalidArg(
1037 "this packet is not an LMFlowBuffer".into(),
1038 ))
1039 }
1040 };
1041 *arc = Arc::new(Payload::Builtin(Builtin::Buffer(owned)));
1042 }
1043
1044 match Arc::get_mut(arc).expect("buffer payload is exclusive after CoW") {
1045 Payload::Builtin(Builtin::Buffer(buffer)) => Ok(BufferView {
1046 data: buffer.bytes.as_mut_ptr() as *mut c_void,
1047 shape: buffer.shape,
1048 strides: buffer.strides,
1049 ndim: buffer.ndim,
1050 dtype: buffer.dtype,
1051 readonly: false,
1052 }),
1053 Payload::Native(value) if value.is::<ExternalBuffer>() => {
1054 let buffer = value
1055 .downcast_mut::<ExternalBuffer>()
1056 .expect("type checked above");
1057 Ok(BufferView {
1058 data: buffer.data,
1059 shape: buffer.shape,
1060 strides: buffer.strides,
1061 ndim: buffer.ndim,
1062 dtype: buffer.dtype,
1063 readonly: false,
1064 })
1065 }
1066 _ => Err(Error::InvalidArg(
1067 "this packet is not an LMFlowBuffer".into(),
1068 )),
1069 }
1070 }
1071
1072 pub fn make_mutable_builtin(&mut self) -> Result<&mut Builtin> {
1077 let body = Arc::make_mut(self.data.as_mut().ok_or_else(|| {
1078 Error::InvalidArg("cannot get a writable view of an empty packet".into())
1079 })?);
1080 let arc = &mut body.payload;
1081
1082 if Arc::get_mut(arc).is_none() {
1083 let copy = match arc.as_ref() {
1085 Payload::Builtin(b) => {
1086 record_cow_copy(b.byte_size());
1087 Payload::Builtin(b.clone())
1088 }
1089 _ => {
1090 return Err(Error::InvalidArg(
1091 "custom payload cannot be copied by the engine (only builtin types support CoW)".into(),
1092 ))
1093 }
1094 };
1095 *arc = Arc::new(copy);
1096 }
1097 match Arc::get_mut(arc).expect("just guaranteed exclusive") {
1099 Payload::Builtin(b) => Ok(b),
1100 _ => Err(Error::InvalidArg(
1101 "this packet is not a builtin payload; CoW not supported".into(),
1102 )),
1103 }
1104 }
1105
1106 pub fn debug_string(&self) -> String {
1108 let ty = match self.data.as_deref().map(|body| body.payload.as_ref()) {
1109 None => "Empty".to_string(),
1110 Some(Payload::Native(value)) => {
1111 if let Some(buffer) = value.downcast_ref::<ExternalBuffer>() {
1112 format!(
1113 "Buffer[{}] dtype={} external",
1114 buffer.shape[..buffer.ndim as usize]
1115 .iter()
1116 .map(|dimension| dimension.to_string())
1117 .collect::<Vec<_>>()
1118 .join("x"),
1119 buffer.dtype
1120 )
1121 } else {
1122 "Native".to_string()
1123 }
1124 }
1125 Some(Payload::Foreign(f)) => format!("Foreign(type#{})", f.type_id),
1126 Some(Payload::Builtin(b)) => match b {
1127 Builtin::Bytes(v) => format!("Bytes[{}]", v.len()),
1128 Builtin::I64(v) => format!("I64({v})"),
1129 Builtin::F64(v) => format!("F64({v})"),
1130 Builtin::Bool(v) => format!("Bool({v})"),
1131 Builtin::Str(s) => format!("Str({:?})", s.to_string_lossy()),
1132 Builtin::Buffer(b) => format!(
1133 "Buffer[{}] dtype={}",
1134 b.shape_slice()
1135 .iter()
1136 .map(|d| d.to_string())
1137 .collect::<Vec<_>>()
1138 .join("x"),
1139 b.dtype
1140 ),
1141 },
1142 };
1143 format!("Packet{{type={ty}, ts={}}}", self.ts)
1144 }
1145}
1146
1147impl std::fmt::Debug for Packet {
1148 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1149 f.write_str(&self.debug_string())
1150 }
1151}
1152
1153#[cfg(test)]
1154mod tests {
1155 use super::*;
1156
1157 #[test]
1158 fn native_roundtrip() {
1159 let p = Packet::new(42i32).at(Timestamp(7));
1160 assert_eq!(p.get::<i32>(), Some(&42));
1161 assert_eq!(
1162 p.get::<i64>(),
1163 None,
1164 "type mismatch must return None, not UB"
1165 );
1166 assert_eq!(p.timestamp(), Timestamp(7));
1167 assert!(!p.is_empty());
1168 }
1169
1170 #[test]
1171 fn empty_packet() {
1172 let p = Packet::empty();
1173 assert!(p.is_empty());
1174 assert_eq!(p.type_id(), type_id::NONE);
1175 assert_eq!(p.timestamp(), Timestamp::unset());
1176 }
1177
1178 #[test]
1179 fn clone_shares_and_does_not_copy() {
1180 let a = Packet::from_builtin(Builtin::I64(5));
1181 let b = a.clone();
1182 assert_eq!(a.ref_count(), 2);
1183 assert_eq!(b.ref_count(), 2);
1184 assert!(std::ptr::eq(
1186 a.payload().unwrap() as *const Payload,
1187 b.payload().unwrap() as *const Payload
1188 ));
1189 }
1190
1191 #[test]
1192 fn cow_is_zero_copy_when_exclusive() {
1193 let mut p = Packet::from_builtin(Builtin::Buffer(
1194 BufferData::new(&[2, 3], dtype::U8).unwrap(),
1195 ));
1196 let before = p.payload().unwrap() as *const Payload;
1197 p.make_mutable_builtin().unwrap();
1198 let after = p.payload().unwrap() as *const Payload;
1199 assert!(
1200 std::ptr::eq(before, after),
1201 "no copy should happen when exclusive"
1202 );
1203 }
1204
1205 #[test]
1206 fn cow_copies_when_shared_and_does_not_disturb_other() {
1207 let a = Packet::from_builtin(Builtin::Buffer(BufferData::new(&[4], dtype::U8).unwrap()));
1208 let mut b = a.clone();
1209 assert_eq!(a.ref_count(), 2);
1210
1211 match b.make_mutable_builtin().unwrap() {
1212 Builtin::Buffer(buf) => buf.bytes[0] = 0xAB,
1213 _ => unreachable!(),
1214 }
1215 assert_eq!(a.ref_count(), 1, "a should be exclusive again after copy");
1217 match a.as_builtin().unwrap() {
1218 Builtin::Buffer(buf) => assert_eq!(buf.bytes[0], 0),
1219 _ => unreachable!(),
1220 }
1221 match b.as_builtin().unwrap() {
1222 Builtin::Buffer(buf) => assert_eq!(buf.bytes[0], 0xAB),
1223 _ => unreachable!(),
1224 }
1225 }
1226
1227 #[test]
1228 fn cow_rejects_non_builtin() {
1229 let mut p = Packet::new(1u8);
1230 assert!(
1231 p.make_mutable_builtin().is_err(),
1232 "custom payload cannot be copied; must error rather than silently fail"
1233 );
1234 }
1235
1236 #[test]
1237 fn buffer_strides_are_row_major_contiguous() {
1238 let b = BufferData::new(&[2, 3, 4], dtype::F32).unwrap();
1239 assert_eq!(b.ndim, 3);
1240 assert_eq!(b.shape_slice(), &[2, 3, 4]);
1241 assert_eq!(&b.strides[..3], &[3 * 4 * 4, 4 * 4, 4]);
1243 assert_eq!(b.bytes.len(), 2 * 3 * 4 * 4);
1244 }
1245
1246 #[test]
1247 fn buffer_rejects_bad_args() {
1248 assert!(
1249 BufferData::new(&[], dtype::U8).is_err(),
1250 "ndim 0 is invalid"
1251 );
1252 assert!(
1253 BufferData::new(&[1, 2, 3, 4, 5, 6, 7, 8, 9], dtype::U8).is_err(),
1254 "must reject exceeding MAX_DIMS"
1255 );
1256 assert!(
1257 BufferData::new(&[2, 2], 999).is_err(),
1258 "must reject unknown dtype"
1259 );
1260 assert!(
1261 BufferData::new(&[-1], dtype::U8).is_err(),
1262 "must reject negative shape"
1263 );
1264 }
1265
1266 #[test]
1267 fn dtype_sizes() {
1268 assert_eq!(dtype_size(dtype::U8), 1);
1269 assert_eq!(dtype_size(dtype::F16), 2);
1270 assert_eq!(dtype_size(dtype::F32), 4);
1271 assert_eq!(dtype_size(dtype::F64), 8);
1272 assert_eq!(dtype_size(12345), 0, "unknown dtype returns 0");
1273 }
1274
1275 #[test]
1276 fn foreign_drop_fn_called_exactly_once() {
1277 use std::sync::atomic::{AtomicUsize, Ordering};
1278 static CALLS: AtomicUsize = AtomicUsize::new(0);
1279 unsafe extern "C" fn drop_it(p: *mut c_void) {
1280 CALLS.fetch_add(1, Ordering::SeqCst);
1281 drop(unsafe { Box::from_raw(p as *mut u32) });
1282 }
1283 let raw = Box::into_raw(Box::new(7u32)) as *mut c_void;
1284 {
1285 let p = unsafe { Packet::from_foreign(raw, 99, Some(drop_it)) };
1286 let q = p.clone();
1287 assert_eq!(q.type_id(), 99);
1288 assert_eq!(
1289 CALLS.load(Ordering::SeqCst),
1290 0,
1291 "must not free while references remain"
1292 );
1293 }
1294 assert_eq!(
1295 CALLS.load(Ordering::SeqCst),
1296 1,
1297 "must free exactly once when refcount hits zero"
1298 );
1299 }
1300
1301 #[test]
1302 fn fnv1a_matches_cpp_sugar_layer() {
1303 assert_eq!(fnv1a_type_id("i"), 12638195996648667684);
1306 assert_eq!(fnv1a_type_id("d"), 12638183902020757363);
1307 assert!(
1308 fnv1a_type_id("") >= 16,
1309 "empty name must also avoid the builtin range"
1310 );
1311 }
1312
1313 #[test]
1314 fn interop_packet_carries_explicit_type_id() {
1315 let name = "lmflow.test.ExplicitI32";
1316 let tid = fnv1a_type_id(name);
1317 register_type_descriptor(tid, name, 4, 4).unwrap();
1318 let p = unsafe { Packet::new_interop(7i32, tid) };
1319 assert_eq!(p.type_id(), tid);
1320 let ptr = p.foreign_ptr().expect("should get a data pointer");
1322 assert_eq!(unsafe { *(ptr as *const i32) }, 7);
1323 }
1324
1325 #[test]
1326 fn interop_drop_panic_does_not_abort_host() {
1327 const CHILD: &str = "LMFLOW_INTEROP_DROP_PANIC_CHILD";
1328 if std::env::var_os(CHILD).is_none() {
1329 let output = std::process::Command::new(std::env::current_exe().unwrap())
1330 .args([
1331 "--exact",
1332 "packet::tests::interop_drop_panic_does_not_abort_host",
1333 ])
1334 .env(CHILD, "1")
1335 .output()
1336 .unwrap();
1337 assert!(
1338 output.status.success(),
1339 "child failed: {}\n{}",
1340 output.status,
1341 String::from_utf8_lossy(&output.stderr)
1342 );
1343 return;
1344 }
1345
1346 struct PanicDrop {
1347 _value: u8,
1348 }
1349 impl Drop for PanicDrop {
1350 fn drop(&mut self) {
1351 panic!("interop payload destructor failed");
1352 }
1353 }
1354
1355 let name = "lmflow.test.PanicDropInterop";
1356 let id = fnv1a_type_id(name);
1357 register_type_descriptor(
1358 id,
1359 name,
1360 std::mem::size_of::<PanicDrop>(),
1361 std::mem::align_of::<PanicDrop>(),
1362 )
1363 .unwrap();
1364 let packet = unsafe { Packet::new_interop(PanicDrop { _value: 0 }, id) };
1365 drop(packet);
1366 }
1367
1368 #[test]
1369 #[should_panic(expected = "built-in ids 0..=15 are reserved")]
1370 fn interop_packet_rejects_builtin_type_ids() {
1371 let _ = unsafe { Packet::new_interop(7i32, type_id::I64) };
1372 }
1373
1374 #[test]
1375 #[should_panic(expected = "is not registered")]
1376 fn interop_packet_rejects_unregistered_custom_type() {
1377 let id = fnv1a_type_id("lmflow.test.UnregisteredInterop");
1378 let _ = unsafe { Packet::new_interop(7i32, id) };
1379 }
1380
1381 #[test]
1382 #[should_panic(expected = "size mismatch")]
1383 fn interop_packet_rejects_registered_layout_mismatch() {
1384 let name = "lmflow.test.LayoutMismatch";
1385 let id = fnv1a_type_id(name);
1386 register_type_descriptor(id, name, 16, 8).unwrap();
1387 let _ = unsafe { Packet::new_interop(7u64, id) };
1388 }
1389
1390 #[test]
1391 fn declared_interop_type_uses_stable_name() {
1392 #[repr(C)]
1393 struct Point {
1394 x: i32,
1395 y: i32,
1396 }
1397 unsafe impl InteropType for Point {
1398 const TYPE_NAME: &'static str = "lmflow.test.Point";
1399 }
1400
1401 let p = Packet::from_interop(Point { x: 3, y: 4 });
1402 assert_eq!(p.type_id(), fnv1a_type_id(Point::TYPE_NAME));
1403 assert_eq!(type_name(p.type_id()), Point::TYPE_NAME);
1404 assert_eq!(p.byte_size(), std::mem::size_of::<Point>() as u64);
1405 let point = unsafe { &*(p.foreign_ptr().unwrap() as *const Point) };
1406 assert_eq!((point.x, point.y), (3, 4));
1407 }
1408
1409 #[test]
1410 fn strict_type_registration_rejects_identity_and_layout_conflicts() {
1411 let id = fnv1a_type_id("lmflow.test.StrictRegistration");
1412 register_type_descriptor(id, "lmflow.test.StrictRegistration", 8, 4).unwrap();
1413 register_type_descriptor(id, "lmflow.test.StrictRegistration", 8, 4).unwrap();
1414
1415 let layout_error =
1416 register_type_descriptor(id, "lmflow.test.StrictRegistration", 16, 8).unwrap_err();
1417 assert!(layout_error.to_string().contains("already registered"));
1418
1419 let name_error =
1420 register_type_descriptor(id, "lmflow.test.OtherRegistration", 8, 4).unwrap_err();
1421 assert!(name_error.to_string().contains("stable-name id"));
1422 }
1423
1424 #[test]
1425 fn strict_type_registration_rejects_noncanonical_id() {
1426 let name = "lmflow.test.NonCanonical";
1427 let expected = fnv1a_type_id(name);
1428 let error = register_type_descriptor(expected + 1, name, 8, 8).unwrap_err();
1429 let message = error.to_string();
1430 assert!(message.contains(name), "{message}");
1431 assert!(message.contains(&expected.to_string()), "{message}");
1432 }
1433
1434 #[test]
1435 fn registered_foreign_payload_uses_descriptor_size() {
1436 unsafe extern "C" fn drop_u64(p: *mut c_void) {
1437 drop(unsafe { Box::from_raw(p as *mut u64) });
1438 }
1439 let name = "lmflow.test.ForeignSize";
1440 let id = fnv1a_type_id(name);
1441 register_type_descriptor(id, name, 8, 8).unwrap();
1442 let ptr = Box::into_raw(Box::new(42u64)) as *mut c_void;
1443 let packet = unsafe { Packet::from_foreign(ptr, id, Some(drop_u64)) };
1444 assert_eq!(packet.byte_size(), 8);
1445 }
1446
1447 #[test]
1448 fn byte_size_accounting() {
1449 assert_eq!(Packet::from_builtin(Builtin::I64(0)).byte_size(), 8);
1450 assert_eq!(
1451 Packet::from_builtin(Builtin::Bytes(vec![0; 100])).byte_size(),
1452 100
1453 );
1454 assert_eq!(Packet::new(0u64).byte_size(), 0);
1456 }
1457
1458 #[test]
1459 fn debug_string_is_informative() {
1460 let p = Packet::from_builtin(Builtin::Buffer(
1461 BufferData::new(&[3, 224, 224], dtype::F32).unwrap(),
1462 ))
1463 .at(Timestamp(42));
1464 let s = p.debug_string();
1465 assert!(s.contains("3x224x224"), "should contain shape: {s}");
1466 assert!(s.contains("42"), "should contain timestamp: {s}");
1467 }
1468
1469 #[test]
1470 fn metadata_can_be_listed_tested_and_removed_without_copying_payload() {
1471 let mut packet = Packet::from_i64(7)
1472 .with_metadata("confidence", 0.9)
1473 .with_metadata("category", "person");
1474 let shared = packet.clone();
1475 let payload = packet.payload().unwrap() as *const Payload;
1476
1477 assert!(packet.has_metadata("confidence"));
1478 assert_eq!(
1479 packet.metadata_keys().collect::<Vec<_>>(),
1480 vec!["category", "confidence"]
1481 );
1482 assert_eq!(
1483 packet.remove_metadata("category"),
1484 Some(MetadataValue::String("person".into()))
1485 );
1486 assert!(!packet.has_metadata("category"));
1487 assert!(shared.has_metadata("category"));
1488 assert!(std::ptr::eq(payload, packet.payload().unwrap()));
1489 assert!(std::ptr::eq(payload, shared.payload().unwrap()));
1490 }
1491}