Skip to main content

edgefirst_tensor/
tensor_dyn.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{DType, PixelFormat, Tensor, TensorMemory, TensorTrait};
5use half::f16;
6use std::fmt;
7
8/// Type-erased tensor. Wraps a `Tensor<T>` with runtime element type.
9#[non_exhaustive]
10pub enum TensorDyn {
11    /// Unsigned 8-bit integer tensor.
12    U8(Tensor<u8>),
13    /// Signed 8-bit integer tensor.
14    I8(Tensor<i8>),
15    /// Unsigned 16-bit integer tensor.
16    U16(Tensor<u16>),
17    /// Signed 16-bit integer tensor.
18    I16(Tensor<i16>),
19    /// Unsigned 32-bit integer tensor.
20    U32(Tensor<u32>),
21    /// Signed 32-bit integer tensor.
22    I32(Tensor<i32>),
23    /// Unsigned 64-bit integer tensor.
24    U64(Tensor<u64>),
25    /// Signed 64-bit integer tensor.
26    I64(Tensor<i64>),
27    /// 16-bit floating-point tensor.
28    F16(Tensor<f16>),
29    /// 32-bit floating-point tensor.
30    F32(Tensor<f32>),
31    /// 64-bit floating-point tensor.
32    F64(Tensor<f64>),
33}
34
35/// Dispatch a method call across all TensorDyn variants.
36macro_rules! dispatch {
37    ($self:expr, $method:ident $(, $arg:expr)*) => {
38        match $self {
39            TensorDyn::U8(t) => t.$method($($arg),*),
40            TensorDyn::I8(t) => t.$method($($arg),*),
41            TensorDyn::U16(t) => t.$method($($arg),*),
42            TensorDyn::I16(t) => t.$method($($arg),*),
43            TensorDyn::U32(t) => t.$method($($arg),*),
44            TensorDyn::I32(t) => t.$method($($arg),*),
45            TensorDyn::U64(t) => t.$method($($arg),*),
46            TensorDyn::I64(t) => t.$method($($arg),*),
47            TensorDyn::F16(t) => t.$method($($arg),*),
48            TensorDyn::F32(t) => t.$method($($arg),*),
49            TensorDyn::F64(t) => t.$method($($arg),*),
50        }
51    };
52}
53
54/// Like [`dispatch!`], but for methods returning `Result<Tensor<T>>`: rewrap the
55/// typed result back into the matching `TensorDyn` variant. Keeps sub-region
56/// fan-out (`batch`, future `view`) to one line instead of an 11-arm match.
57macro_rules! dyn_fanout {
58    ($self:expr, $method:ident $(, $arg:expr)*) => {
59        match $self {
60            TensorDyn::U8(t) => t.$method($($arg),*).map(TensorDyn::U8),
61            TensorDyn::I8(t) => t.$method($($arg),*).map(TensorDyn::I8),
62            TensorDyn::U16(t) => t.$method($($arg),*).map(TensorDyn::U16),
63            TensorDyn::I16(t) => t.$method($($arg),*).map(TensorDyn::I16),
64            TensorDyn::U32(t) => t.$method($($arg),*).map(TensorDyn::U32),
65            TensorDyn::I32(t) => t.$method($($arg),*).map(TensorDyn::I32),
66            TensorDyn::U64(t) => t.$method($($arg),*).map(TensorDyn::U64),
67            TensorDyn::I64(t) => t.$method($($arg),*).map(TensorDyn::I64),
68            TensorDyn::F16(t) => t.$method($($arg),*).map(TensorDyn::F16),
69            TensorDyn::F32(t) => t.$method($($arg),*).map(TensorDyn::F32),
70            TensorDyn::F64(t) => t.$method($($arg),*).map(TensorDyn::F64),
71        }
72    };
73}
74
75/// Generate the three downcast methods (ref, mut ref, owned) for one variant.
76macro_rules! downcast_methods {
77    ($variant:ident, $ty:ty, $as_name:ident, $as_mut_name:ident, $into_name:ident) => {
78        /// Returns a shared reference to the inner tensor if the type matches.
79        pub fn $as_name(&self) -> Option<&Tensor<$ty>> {
80            match self {
81                Self::$variant(t) => Some(t),
82                _ => None,
83            }
84        }
85
86        /// Returns a mutable reference to the inner tensor if the type matches.
87        pub fn $as_mut_name(&mut self) -> Option<&mut Tensor<$ty>> {
88            match self {
89                Self::$variant(t) => Some(t),
90                _ => None,
91            }
92        }
93
94        /// Unwraps the inner tensor if the type matches, otherwise returns `self` as `Err`.
95        /// The Err variant is necessarily large (returns the unconsumed TensorDyn).
96        #[allow(clippy::result_large_err)]
97        pub fn $into_name(self) -> Result<Tensor<$ty>, Self> {
98            match self {
99                Self::$variant(t) => Ok(t),
100                other => Err(other),
101            }
102        }
103    };
104}
105
106impl TensorDyn {
107    /// Return the runtime element type discriminant.
108    pub fn dtype(&self) -> DType {
109        match self {
110            Self::U8(_) => DType::U8,
111            Self::I8(_) => DType::I8,
112            Self::U16(_) => DType::U16,
113            Self::I16(_) => DType::I16,
114            Self::U32(_) => DType::U32,
115            Self::I32(_) => DType::I32,
116            Self::U64(_) => DType::U64,
117            Self::I64(_) => DType::I64,
118            Self::F16(_) => DType::F16,
119            Self::F32(_) => DType::F32,
120            Self::F64(_) => DType::F64,
121        }
122    }
123
124    /// Return the tensor shape.
125    pub fn shape(&self) -> &[usize] {
126        dispatch!(self, shape)
127    }
128
129    /// Return the tensor name.
130    pub fn name(&self) -> String {
131        dispatch!(self, name)
132    }
133
134    /// Return the pixel format (None if not an image tensor).
135    pub fn format(&self) -> Option<PixelFormat> {
136        dispatch!(self, format)
137    }
138
139    /// Return the image width (None if not an image tensor).
140    pub fn width(&self) -> Option<usize> {
141        dispatch!(self, width)
142    }
143
144    /// Return the image height (None if not an image tensor).
145    pub fn height(&self) -> Option<usize> {
146        dispatch!(self, height)
147    }
148
149    /// Return the total size of this tensor in bytes.
150    pub fn size(&self) -> usize {
151        dispatch!(self, size)
152    }
153
154    /// Return the memory allocation type.
155    pub fn memory(&self) -> TensorMemory {
156        dispatch!(self, memory)
157    }
158
159    /// Reshape this tensor. Total element count must remain the same.
160    pub fn reshape(&mut self, shape: &[usize]) -> crate::Result<()> {
161        dispatch!(self, reshape, shape)
162    }
163
164    /// Attach pixel format metadata to this tensor.
165    ///
166    /// Validates that the tensor's shape is compatible with the format's
167    /// layout (packed, planar, or semi-planar).
168    ///
169    /// # Arguments
170    ///
171    /// * `format` - The pixel format to attach
172    ///
173    /// # Returns
174    ///
175    /// `Ok(())` on success, with the format stored as metadata on the tensor.
176    ///
177    /// # Errors
178    ///
179    /// Returns `Error::InvalidShape` if the tensor shape doesn't match
180    /// the expected layout for the given format.
181    pub fn set_format(&mut self, format: PixelFormat) -> crate::Result<()> {
182        dispatch!(self, set_format, format)
183    }
184
185    /// Attach pixel format metadata, consuming and returning self.
186    ///
187    /// Enables builder-style chaining.
188    ///
189    /// # Arguments
190    ///
191    /// * `format` - The pixel format to attach
192    ///
193    /// # Returns
194    ///
195    /// The tensor with format metadata attached.
196    ///
197    /// # Errors
198    ///
199    /// Returns `Error::InvalidShape` if the tensor shape doesn't match
200    /// the expected layout for the given format.
201    pub fn with_format(mut self, format: PixelFormat) -> crate::Result<Self> {
202        self.set_format(format)?;
203        Ok(self)
204    }
205
206    /// Colorimetry metadata (`None` = undefined; never auto-filled).
207    pub fn colorimetry(&self) -> Option<crate::Colorimetry> {
208        dispatch!(self, colorimetry)
209    }
210
211    /// Attach/clear colorimetry metadata.
212    pub fn set_colorimetry(&mut self, c: Option<crate::Colorimetry>) {
213        dispatch!(self, set_colorimetry, c)
214    }
215
216    /// Builder-style colorimetry attach (consumes and returns self).
217    pub fn with_colorimetry(mut self, c: crate::Colorimetry) -> Self {
218        self.set_colorimetry(Some(c));
219        self
220    }
221
222    /// Row stride in bytes (`None` = tightly packed).
223    pub fn row_stride(&self) -> Option<usize> {
224        dispatch!(self, row_stride)
225    }
226
227    /// Effective row stride: stored stride or computed from format and width.
228    pub fn effective_row_stride(&self) -> Option<usize> {
229        dispatch!(self, effective_row_stride)
230    }
231
232    /// Set logical dimensions + format to a decoded image, reusing the
233    /// allocation. See [`Tensor::configure_image`].
234    pub fn configure_image(
235        &mut self,
236        width: usize,
237        height: usize,
238        format: PixelFormat,
239    ) -> crate::Result<()> {
240        dispatch!(self, configure_image, width, height, format)
241    }
242
243    /// Set the row stride in bytes for externally allocated buffers with
244    /// row padding.
245    ///
246    /// Must be called before the tensor is first used for rendering. The
247    /// format must be set before calling this method.
248    pub fn set_row_stride(&mut self, stride: usize) -> crate::Result<()> {
249        dispatch!(self, set_row_stride, stride)
250    }
251
252    /// Builder-style: set row stride, consuming and returning self.
253    pub fn with_row_stride(mut self, stride: usize) -> crate::Result<Self> {
254        self.set_row_stride(stride)?;
255        Ok(self)
256    }
257
258    /// Byte offset within the DMA-BUF where image data starts (`None` = 0).
259    pub fn plane_offset(&self) -> Option<usize> {
260        dispatch!(self, plane_offset)
261    }
262
263    /// The parent-image snapshot if this tensor is a [`view`](Self::view)/
264    /// [`batch`](Self::batch) sub-region; `None` for a whole tensor. See
265    /// [`Tensor::view_origin`].
266    pub fn view_origin(&self) -> Option<crate::ViewOrigin> {
267        dispatch!(self, view_origin)
268    }
269
270    /// Set the byte offset within the DMA-BUF where image data starts.
271    pub fn set_plane_offset(&mut self, offset: usize) {
272        dispatch!(self, set_plane_offset, offset)
273    }
274
275    /// Borrow batch element `n` of a batched tensor (leading `N` dimension) as a
276    /// zero-copy view sharing this tensor's allocation. See [`Tensor::batch`].
277    pub fn batch(&self, n: usize) -> crate::Result<TensorDyn> {
278        dyn_fanout!(self, batch, n)
279    }
280
281    /// Borrow a rectangular spatial sub-region (the destination/source crop) as
282    /// a zero-copy view sharing this tensor's allocation. See [`Tensor::view`].
283    pub fn view(&self, region: crate::Region) -> crate::Result<TensorDyn> {
284        dyn_fanout!(self, view, region)
285    }
286
287    /// The CUDA registration for this tensor, if any.
288    ///
289    /// Returns `None` when no CUDA handle has been attached (the common non-CUDA case).
290    /// This check is a pure local field read — no thread routing occurs.
291    pub fn cuda(&self) -> Option<&crate::cuda::CudaHandle> {
292        dispatch!(self, cuda)
293    }
294
295    /// Fast-fail CUDA map: `None` when no handle is attached; else maps the
296    /// PBO through the GL worker and returns a scoped device-pointer guard.
297    ///
298    /// The same try-`cuda_map`-then-[`map`](crate::TensorTrait::map) fallback pattern that applies to
299    /// [`Tensor::cuda_map`](crate::Tensor::cuda_map) applies here: call `cuda_map()` first for a
300    /// zero-copy device pointer; when it returns `None` (no CUDA handle attached), fall back to the
301    /// typed host mapping via the inner tensor.
302    ///
303    /// # Example
304    ///
305    /// ```no_run
306    /// use edgefirst_tensor::TensorDyn;
307    /// # fn feed_tensorrt(_dptr: *mut std::ffi::c_void, _bytes: usize) {}
308    /// # fn demo(t: &TensorDyn) {
309    /// if let Some(cuda) = t.cuda_map() {
310    ///     feed_tensorrt(cuda.device_ptr(), cuda.len());
311    /// } else {
312    ///     // No CUDA handle — use the typed inner tensor for host access.
313    ///     // See Tensor::cuda_map for the full fallback example.
314    /// }
315    /// # }
316    /// ```
317    pub fn cuda_map(&self) -> Option<crate::cuda::CudaMap<'_>> {
318        dispatch!(self, cuda_map)
319    }
320
321    /// Quantization metadata. Returns `None` for float variants (F16, F32,
322    /// F64) — quantization does not apply to floating-point tensors.
323    /// Otherwise delegates to the typed `Tensor<T>::quantization()` accessor.
324    pub fn quantization(&self) -> Option<&crate::Quantization> {
325        match self {
326            Self::U8(t) => t.quantization(),
327            Self::I8(t) => t.quantization(),
328            Self::U16(t) => t.quantization(),
329            Self::I16(t) => t.quantization(),
330            Self::U32(t) => t.quantization(),
331            Self::I32(t) => t.quantization(),
332            Self::U64(t) => t.quantization(),
333            Self::I64(t) => t.quantization(),
334            Self::F16(_) | Self::F32(_) | Self::F64(_) => None,
335        }
336    }
337
338    /// Attach quantization metadata. Fails on float variants with
339    /// [`Error::QuantizationInvalid`]; delegates to the typed setter for
340    /// integer variants.
341    pub fn set_quantization(&mut self, q: crate::Quantization) -> crate::Result<()> {
342        match self {
343            Self::U8(t) => t.set_quantization(q),
344            Self::I8(t) => t.set_quantization(q),
345            Self::U16(t) => t.set_quantization(q),
346            Self::I16(t) => t.set_quantization(q),
347            Self::U32(t) => t.set_quantization(q),
348            Self::I32(t) => t.set_quantization(q),
349            Self::U64(t) => t.set_quantization(q),
350            Self::I64(t) => t.set_quantization(q),
351            Self::F16(_) | Self::F32(_) | Self::F64(_) => Err(crate::Error::QuantizationInvalid {
352                field: "dtype_is_integer",
353                expected: "integer tensor dtype (u8/i8/u16/i16/u32/i32/u64/i64)".to_string(),
354                got: format!("{:?}", self.dtype()),
355            }),
356        }
357    }
358
359    /// Builder-style variant of [`Self::set_quantization`]. Consumes self
360    /// and returns it with quantization applied (or the original error).
361    pub fn with_quantization(mut self, q: crate::Quantization) -> crate::Result<Self> {
362        self.set_quantization(q)?;
363        Ok(self)
364    }
365
366    /// Clear any quantization metadata. No-op on float variants.
367    pub fn clear_quantization(&mut self) {
368        match self {
369            Self::U8(t) => t.clear_quantization(),
370            Self::I8(t) => t.clear_quantization(),
371            Self::U16(t) => t.clear_quantization(),
372            Self::I16(t) => t.clear_quantization(),
373            Self::U32(t) => t.clear_quantization(),
374            Self::I32(t) => t.clear_quantization(),
375            Self::U64(t) => t.clear_quantization(),
376            Self::I64(t) => t.clear_quantization(),
377            Self::F16(_) | Self::F32(_) | Self::F64(_) => {}
378        }
379    }
380
381    /// Clone the file descriptor associated with this tensor.
382    #[cfg(unix)]
383    pub fn clone_fd(&self) -> crate::Result<std::os::fd::OwnedFd> {
384        dispatch!(self, clone_fd)
385    }
386
387    /// Clone the DMA-BUF file descriptor backing this tensor (Linux only).
388    ///
389    /// # Returns
390    ///
391    /// An owned duplicate of the DMA-BUF file descriptor.
392    ///
393    /// # Errors
394    ///
395    /// * `Error::NotImplemented` if the tensor is not DMA-backed (Mem/Shm/Pbo)
396    /// * `Error::IoError` if the fd clone syscall fails (e.g., fd limit reached)
397    #[cfg(target_os = "linux")]
398    pub fn dmabuf_clone(&self) -> crate::Result<std::os::fd::OwnedFd> {
399        if self.memory() != TensorMemory::Dma {
400            return Err(crate::Error::NotImplemented(format!(
401                "dmabuf_clone requires DMA-backed tensor, got {:?}",
402                self.memory()
403            )));
404        }
405        self.clone_fd()
406    }
407
408    /// Borrow the DMA-BUF file descriptor backing this tensor (Linux only).
409    ///
410    /// # Returns
411    ///
412    /// A borrowed reference to the DMA-BUF file descriptor, tied to `self`'s
413    /// lifetime.
414    ///
415    /// # Errors
416    ///
417    /// * `Error::NotImplemented` if the tensor is not DMA-backed
418    #[cfg(target_os = "linux")]
419    pub fn dmabuf(&self) -> crate::Result<std::os::fd::BorrowedFd<'_>> {
420        dispatch!(self, dmabuf)
421    }
422
423    /// Return `true` if this tensor uses separate plane allocations.
424    pub fn is_multiplane(&self) -> bool {
425        dispatch!(self, is_multiplane)
426    }
427
428    /// Return the [`BufferIdentity`](crate::BufferIdentity) of the underlying
429    /// allocation.
430    ///
431    /// Two `TensorDyn` values share a [`BufferIdentity::id`] iff they were
432    /// produced by cloning the same allocation (e.g. through
433    /// [`DmaTensor::try_clone`](crate::dma::DmaTensor::try_clone)). Separate
434    /// imports of the same physical buffer (e.g. two `from_fd` calls on the
435    /// same dmabuf fd) have **distinct** identities — use
436    /// [`aliases`](Self::aliases) if you need to detect that case.
437    pub fn buffer_identity(&self) -> &crate::BufferIdentity {
438        dispatch!(self, buffer_identity)
439    }
440
441    /// Return `true` if `self` and `other` reference the same underlying
442    /// buffer.
443    ///
444    /// This is the correct check for APIs that require distinct input and
445    /// output tensors (e.g. `ImageProcessor::draw_decoded_masks`, where
446    /// aliasing `dst` and `background` would cause the GL backend to read
447    /// and write the same texture — undefined behaviour on most drivers).
448    ///
449    /// Matching is conservative:
450    /// 1. Matching [`BufferIdentity::id`] → same buffer (always).
451    /// 2. Matching backing type + matching dmabuf fd number (Linux, DMA
452    ///    tensors only) → same buffer, even across separate `from_fd`
453    ///    imports in the same process.
454    ///
455    /// Two distinct `dup`'d fds pointing at the same kernel dma-buf are
456    /// **not** detected — there is no cheap way to resolve that without a
457    /// round-trip through the kernel.
458    pub fn aliases(&self, other: &Self) -> bool {
459        if self.buffer_identity().id() == other.buffer_identity().id() {
460            return true;
461        }
462        if self.memory() != other.memory() {
463            return false;
464        }
465        #[cfg(target_os = "linux")]
466        if self.memory() == TensorMemory::Dma {
467            use std::os::fd::AsRawFd;
468            if let (Ok(a), Ok(b)) = (self.dmabuf(), other.dmabuf()) {
469                return a.as_raw_fd() == b.as_raw_fd();
470            }
471        }
472        false
473    }
474
475    // --- Downcasting ---
476
477    downcast_methods!(U8, u8, as_u8, as_u8_mut, into_u8);
478    downcast_methods!(I8, i8, as_i8, as_i8_mut, into_i8);
479    downcast_methods!(U16, u16, as_u16, as_u16_mut, into_u16);
480    downcast_methods!(I16, i16, as_i16, as_i16_mut, into_i16);
481    downcast_methods!(U32, u32, as_u32, as_u32_mut, into_u32);
482    downcast_methods!(I32, i32, as_i32, as_i32_mut, into_i32);
483    downcast_methods!(U64, u64, as_u64, as_u64_mut, into_u64);
484    downcast_methods!(I64, i64, as_i64, as_i64_mut, into_i64);
485    downcast_methods!(F16, f16, as_f16, as_f16_mut, into_f16);
486    downcast_methods!(F32, f32, as_f32, as_f32_mut, into_f32);
487    downcast_methods!(F64, f64, as_f64, as_f64_mut, into_f64);
488
489    /// Create a type-erased tensor with the given shape and element type.
490    pub fn new(
491        shape: &[usize],
492        dtype: DType,
493        memory: Option<TensorMemory>,
494        name: Option<&str>,
495    ) -> crate::Result<Self> {
496        match dtype {
497            DType::U8 => Tensor::<u8>::new(shape, memory, name).map(Self::U8),
498            DType::I8 => Tensor::<i8>::new(shape, memory, name).map(Self::I8),
499            DType::U16 => Tensor::<u16>::new(shape, memory, name).map(Self::U16),
500            DType::I16 => Tensor::<i16>::new(shape, memory, name).map(Self::I16),
501            DType::U32 => Tensor::<u32>::new(shape, memory, name).map(Self::U32),
502            DType::I32 => Tensor::<i32>::new(shape, memory, name).map(Self::I32),
503            DType::U64 => Tensor::<u64>::new(shape, memory, name).map(Self::U64),
504            DType::I64 => Tensor::<i64>::new(shape, memory, name).map(Self::I64),
505            DType::F16 => Tensor::<f16>::new(shape, memory, name).map(Self::F16),
506            DType::F32 => Tensor::<f32>::new(shape, memory, name).map(Self::F32),
507            DType::F64 => Tensor::<f64>::new(shape, memory, name).map(Self::F64),
508        }
509    }
510
511    /// Create a type-erased tensor from a file descriptor.
512    #[cfg(unix)]
513    pub fn from_fd(
514        fd: std::os::fd::OwnedFd,
515        shape: &[usize],
516        dtype: DType,
517        name: Option<&str>,
518    ) -> crate::Result<Self> {
519        match dtype {
520            DType::U8 => Tensor::<u8>::from_fd(fd, shape, name).map(Self::U8),
521            DType::I8 => Tensor::<i8>::from_fd(fd, shape, name).map(Self::I8),
522            DType::U16 => Tensor::<u16>::from_fd(fd, shape, name).map(Self::U16),
523            DType::I16 => Tensor::<i16>::from_fd(fd, shape, name).map(Self::I16),
524            DType::U32 => Tensor::<u32>::from_fd(fd, shape, name).map(Self::U32),
525            DType::I32 => Tensor::<i32>::from_fd(fd, shape, name).map(Self::I32),
526            DType::U64 => Tensor::<u64>::from_fd(fd, shape, name).map(Self::U64),
527            DType::I64 => Tensor::<i64>::from_fd(fd, shape, name).map(Self::I64),
528            DType::F16 => Tensor::<f16>::from_fd(fd, shape, name).map(Self::F16),
529            DType::F32 => Tensor::<f32>::from_fd(fd, shape, name).map(Self::F32),
530            DType::F64 => Tensor::<f64>::from_fd(fd, shape, name).map(Self::F64),
531        }
532    }
533
534    /// Wrap externally-owned memory as a type-erased tensor without copying.
535    /// The tensor borrows `[ptr, ptr + shape.product() * dtype.size())` as
536    /// [`TensorMemory::Mem`]; `owner`, when `Some`, co-owns the source so it
537    /// outlives the tensor (and all derived views/maps). See
538    /// [`crate::ForeignOwner`] and [`Tensor::from_foreign`].
539    ///
540    /// # Safety
541    ///
542    /// `ptr` must be non-null, aligned to the element type, and valid for
543    /// `shape.product()` elements of `dtype` for as long as the returned
544    /// tensor — and every view/map sharing its backing — is alive. Pass an
545    /// `owner` that co-owns the source to uphold that contract.
546    pub unsafe fn from_foreign_ptr(
547        ptr: *mut u8,
548        shape: &[usize],
549        dtype: DType,
550        owner: Option<crate::ForeignOwner>,
551        name: Option<&str>,
552    ) -> crate::Result<Self> {
553        match dtype {
554            DType::U8 => Tensor::<u8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U8),
555            DType::I8 => Tensor::<i8>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I8),
556            DType::U16 => {
557                Tensor::<u16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U16)
558            }
559            DType::I16 => {
560                Tensor::<i16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I16)
561            }
562            DType::U32 => {
563                Tensor::<u32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U32)
564            }
565            DType::I32 => {
566                Tensor::<i32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I32)
567            }
568            DType::U64 => {
569                Tensor::<u64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::U64)
570            }
571            DType::I64 => {
572                Tensor::<i64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::I64)
573            }
574            DType::F16 => {
575                Tensor::<f16>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F16)
576            }
577            DType::F32 => {
578                Tensor::<f32>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F32)
579            }
580            DType::F64 => {
581                Tensor::<f64>::from_foreign(ptr.cast(), shape, owner, name).map(Self::F64)
582            }
583        }
584    }
585
586    /// Wrap an externally-allocated IOSurface as a type-erased tensor
587    /// (macOS/iOS only).
588    ///
589    /// # Safety
590    ///
591    /// `surface_ref` must be a valid live `IOSurfaceRef`. `shape` must
592    /// match the IOSurface's pixel dimensions and chosen element type.
593    #[cfg(any(target_os = "macos", target_os = "ios"))]
594    pub unsafe fn from_iosurface(
595        surface_ref: *mut std::ffi::c_void,
596        shape: &[usize],
597        dtype: DType,
598        name: Option<&str>,
599    ) -> crate::Result<Self> {
600        unsafe {
601            match dtype {
602                DType::U8 => Tensor::<u8>::from_iosurface(surface_ref, shape, name).map(Self::U8),
603                DType::I8 => Tensor::<i8>::from_iosurface(surface_ref, shape, name).map(Self::I8),
604                DType::U16 => {
605                    Tensor::<u16>::from_iosurface(surface_ref, shape, name).map(Self::U16)
606                }
607                DType::I16 => {
608                    Tensor::<i16>::from_iosurface(surface_ref, shape, name).map(Self::I16)
609                }
610                DType::U32 => {
611                    Tensor::<u32>::from_iosurface(surface_ref, shape, name).map(Self::U32)
612                }
613                DType::I32 => {
614                    Tensor::<i32>::from_iosurface(surface_ref, shape, name).map(Self::I32)
615                }
616                DType::U64 => {
617                    Tensor::<u64>::from_iosurface(surface_ref, shape, name).map(Self::U64)
618                }
619                DType::I64 => {
620                    Tensor::<i64>::from_iosurface(surface_ref, shape, name).map(Self::I64)
621                }
622                DType::F16 => {
623                    Tensor::<f16>::from_iosurface(surface_ref, shape, name).map(Self::F16)
624                }
625                DType::F32 => {
626                    Tensor::<f32>::from_iosurface(surface_ref, shape, name).map(Self::F32)
627                }
628                DType::F64 => {
629                    Tensor::<f64>::from_iosurface(surface_ref, shape, name).map(Self::F64)
630                }
631            }
632        }
633    }
634
635    /// IOSurfaceID for cross-process surface sharing (macOS/iOS only).
636    /// Returns `None` when the tensor is not IOSurface-backed.
637    #[cfg(any(target_os = "macos", target_os = "ios"))]
638    pub fn iosurface_id(&self) -> Option<u32> {
639        dispatch!(self, iosurface_id)
640    }
641
642    /// Borrow the raw `IOSurfaceRef` backing this tensor (macOS/iOS
643    /// only). Returns `None` when the tensor is not IOSurface-backed.
644    /// The pointer's lifetime is tied to `self`.
645    #[cfg(any(target_os = "macos", target_os = "ios"))]
646    pub fn iosurface_ref(&self) -> Option<*mut std::ffi::c_void> {
647        dispatch!(self, iosurface_ref)
648    }
649
650    /// Physical IOSurface dimensions in texels, independent of the logical
651    /// shape (macOS/iOS only). `None` when not IOSurface-backed. The GL
652    /// backend binds the EGL pbuffer at these dims so one cached pbuffer
653    /// serves every frame size a reused pool surface holds.
654    #[cfg(any(target_os = "macos", target_os = "ios"))]
655    pub fn iosurface_physical_dims(&self) -> Option<(usize, usize)> {
656        dispatch!(self, iosurface_physical_dims)
657    }
658
659    /// Wrap an externally-allocated AHardwareBuffer as a type-erased
660    /// tensor (Android only). Used to import buffers from
661    /// CameraX/ImageReader (via JNI), NNAPI, or cross-process binder
662    /// transfers.
663    ///
664    /// # Safety
665    ///
666    /// `buffer_ptr` must be a valid live AHardwareBuffer pointer. `shape`
667    /// must match the buffer's dimensions and chosen element type.
668    #[cfg(target_os = "android")]
669    pub unsafe fn from_hardware_buffer(
670        buffer_ptr: *mut std::ffi::c_void,
671        shape: &[usize],
672        dtype: DType,
673        name: Option<&str>,
674    ) -> crate::Result<Self> {
675        unsafe {
676            match dtype {
677                DType::U8 => {
678                    Tensor::<u8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U8)
679                }
680                DType::I8 => {
681                    Tensor::<i8>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I8)
682                }
683                DType::U16 => {
684                    Tensor::<u16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U16)
685                }
686                DType::I16 => {
687                    Tensor::<i16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I16)
688                }
689                DType::U32 => {
690                    Tensor::<u32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U32)
691                }
692                DType::I32 => {
693                    Tensor::<i32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I32)
694                }
695                DType::U64 => {
696                    Tensor::<u64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::U64)
697                }
698                DType::I64 => {
699                    Tensor::<i64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::I64)
700                }
701                DType::F16 => {
702                    Tensor::<f16>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F16)
703                }
704                DType::F32 => {
705                    Tensor::<f32>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F32)
706                }
707                DType::F64 => {
708                    Tensor::<f64>::from_hardware_buffer(buffer_ptr, shape, name).map(Self::F64)
709                }
710            }
711        }
712    }
713
714    /// Borrow the raw AHardwareBuffer pointer backing this tensor
715    /// (Android only). Returns `None` when the tensor is not
716    /// AHardwareBuffer-backed. The pointer's lifetime is tied to `self`.
717    #[cfg(target_os = "android")]
718    pub fn hardware_buffer_ptr(&self) -> Option<*mut std::ffi::c_void> {
719        dispatch!(self, hardware_buffer_ptr)
720    }
721
722    /// Physical AHardwareBuffer dimensions in texels, independent of the
723    /// logical shape (Android only). `None` when not
724    /// AHardwareBuffer-backed.
725    #[cfg(target_os = "android")]
726    pub fn hardware_buffer_physical_dims(&self) -> Option<(usize, usize)> {
727        dispatch!(self, hardware_buffer_physical_dims)
728    }
729
730    /// Copy the tensor's logical bytes into `dst`, compacting away any
731    /// recorded row-stride padding — see [`Tensor::copy_to_flat`] for the
732    /// full contract. `dst.len()` must equal the tight byte footprint
733    /// (`shape` product × element size).
734    pub fn copy_to_flat(&self, dst: &mut [u8]) -> crate::Result<()> {
735        dispatch!(self, copy_to_flat, dst)
736    }
737
738    /// Create a type-erased image tensor.
739    ///
740    /// # Arguments
741    ///
742    /// * `width` - Image width in pixels
743    /// * `height` - Image height in pixels
744    /// * `format` - Pixel format
745    /// * `dtype` - Element type discriminant
746    /// * `memory` - Optional memory backend (None selects the best available)
747    ///
748    /// # Returns
749    ///
750    /// A new `TensorDyn` wrapping an image tensor of the requested element type.
751    ///
752    /// # Errors
753    ///
754    /// Returns an error if the underlying `Tensor::image` call fails.
755    pub fn image(
756        width: usize,
757        height: usize,
758        format: PixelFormat,
759        dtype: DType,
760        memory: Option<TensorMemory>,
761        access: crate::CpuAccess,
762    ) -> crate::Result<Self> {
763        match dtype {
764            DType::U8 => Tensor::<u8>::image(width, height, format, memory, access).map(Self::U8),
765            DType::I8 => Tensor::<i8>::image(width, height, format, memory, access).map(Self::I8),
766            DType::U16 => {
767                Tensor::<u16>::image(width, height, format, memory, access).map(Self::U16)
768            }
769            DType::I16 => {
770                Tensor::<i16>::image(width, height, format, memory, access).map(Self::I16)
771            }
772            DType::U32 => {
773                Tensor::<u32>::image(width, height, format, memory, access).map(Self::U32)
774            }
775            DType::I32 => {
776                Tensor::<i32>::image(width, height, format, memory, access).map(Self::I32)
777            }
778            DType::U64 => {
779                Tensor::<u64>::image(width, height, format, memory, access).map(Self::U64)
780            }
781            DType::I64 => {
782                Tensor::<i64>::image(width, height, format, memory, access).map(Self::I64)
783            }
784            DType::F16 => {
785                Tensor::<f16>::image(width, height, format, memory, access).map(Self::F16)
786            }
787            DType::F32 => {
788                Tensor::<f32>::image(width, height, format, memory, access).map(Self::F32)
789            }
790            DType::F64 => {
791                Tensor::<f64>::image(width, height, format, memory, access).map(Self::F64)
792            }
793        }
794    }
795
796    /// Allocate an image tensor from a declarative [`ImageDesc`]
797    /// request — dispatching on `desc.dtype()`. See
798    /// [`Tensor::image_desc`] for the compression-request semantics.
799    pub fn image_desc(desc: &crate::ImageDesc) -> crate::Result<Self> {
800        match desc.dtype() {
801            DType::U8 => Tensor::<u8>::image_desc(desc).map(Self::U8),
802            DType::I8 => Tensor::<i8>::image_desc(desc).map(Self::I8),
803            DType::U16 => Tensor::<u16>::image_desc(desc).map(Self::U16),
804            DType::I16 => Tensor::<i16>::image_desc(desc).map(Self::I16),
805            DType::U32 => Tensor::<u32>::image_desc(desc).map(Self::U32),
806            DType::I32 => Tensor::<i32>::image_desc(desc).map(Self::I32),
807            DType::U64 => Tensor::<u64>::image_desc(desc).map(Self::U64),
808            DType::I64 => Tensor::<i64>::image_desc(desc).map(Self::I64),
809            DType::F16 => Tensor::<f16>::image_desc(desc).map(Self::F16),
810            DType::F32 => Tensor::<f32>::image_desc(desc).map(Self::F32),
811            DType::F64 => Tensor::<f64>::image_desc(desc).map(Self::F64),
812        }
813    }
814
815    /// The recorded vendor tile-compression scheme (see
816    /// [`Tensor::compression`]).
817    pub fn compression(&self) -> Option<crate::CompressionScheme> {
818        dispatch!(self, compression)
819    }
820
821    /// Create a DMA-backed image tensor with an explicit row stride that
822    /// may exceed the natural `width * channels * sizeof(T)` pitch.
823    ///
824    /// See [`Tensor::image_with_stride`] for the detailed contract and
825    /// constraints. The TensorDyn wrapper dispatches to the appropriate
826    /// monomorphised `Tensor<T>` based on `dtype`.
827    ///
828    /// # Example
829    ///
830    /// ```no_run
831    /// use edgefirst_tensor::{CpuAccess, TensorDyn, PixelFormat, DType, TensorMemory};
832    /// # fn main() -> edgefirst_tensor::Result<()> {
833    /// // Allocate a 3004×1688 RGBA8 canvas with 64-byte pitch alignment
834    /// // (12032 bytes per row instead of the natural 12016).
835    /// let img = TensorDyn::image_with_stride(
836    ///     3004, 1688,
837    ///     PixelFormat::Rgba, DType::U8,
838    ///     12032,
839    ///     Some(TensorMemory::Dma),
840    ///     CpuAccess::ReadWrite,
841    /// )?;
842    /// assert_eq!(img.width(), Some(3004));       // logical, unchanged
843    /// assert_eq!(img.effective_row_stride(), Some(12032)); // padded
844    /// # Ok(())
845    /// # }
846    /// ```
847    pub fn image_with_stride(
848        width: usize,
849        height: usize,
850        format: PixelFormat,
851        dtype: DType,
852        row_stride_bytes: usize,
853        memory: Option<TensorMemory>,
854        access: crate::CpuAccess,
855    ) -> crate::Result<Self> {
856        match dtype {
857            DType::U8 => Tensor::<u8>::image_with_stride(
858                width,
859                height,
860                format,
861                row_stride_bytes,
862                memory,
863                access,
864            )
865            .map(Self::U8),
866            DType::I8 => Tensor::<i8>::image_with_stride(
867                width,
868                height,
869                format,
870                row_stride_bytes,
871                memory,
872                access,
873            )
874            .map(Self::I8),
875            DType::U16 => Tensor::<u16>::image_with_stride(
876                width,
877                height,
878                format,
879                row_stride_bytes,
880                memory,
881                access,
882            )
883            .map(Self::U16),
884            DType::I16 => Tensor::<i16>::image_with_stride(
885                width,
886                height,
887                format,
888                row_stride_bytes,
889                memory,
890                access,
891            )
892            .map(Self::I16),
893            DType::U32 => Tensor::<u32>::image_with_stride(
894                width,
895                height,
896                format,
897                row_stride_bytes,
898                memory,
899                access,
900            )
901            .map(Self::U32),
902            DType::I32 => Tensor::<i32>::image_with_stride(
903                width,
904                height,
905                format,
906                row_stride_bytes,
907                memory,
908                access,
909            )
910            .map(Self::I32),
911            DType::U64 => Tensor::<u64>::image_with_stride(
912                width,
913                height,
914                format,
915                row_stride_bytes,
916                memory,
917                access,
918            )
919            .map(Self::U64),
920            DType::I64 => Tensor::<i64>::image_with_stride(
921                width,
922                height,
923                format,
924                row_stride_bytes,
925                memory,
926                access,
927            )
928            .map(Self::I64),
929            DType::F16 => Tensor::<f16>::image_with_stride(
930                width,
931                height,
932                format,
933                row_stride_bytes,
934                memory,
935                access,
936            )
937            .map(Self::F16),
938            DType::F32 => Tensor::<f32>::image_with_stride(
939                width,
940                height,
941                format,
942                row_stride_bytes,
943                memory,
944                access,
945            )
946            .map(Self::F32),
947            DType::F64 => Tensor::<f64>::image_with_stride(
948                width,
949                height,
950                format,
951                row_stride_bytes,
952                memory,
953                access,
954            )
955            .map(Self::F64),
956        }
957    }
958}
959
960// --- From impls ---
961
962impl From<Tensor<u8>> for TensorDyn {
963    fn from(t: Tensor<u8>) -> Self {
964        Self::U8(t)
965    }
966}
967
968impl From<Tensor<i8>> for TensorDyn {
969    fn from(t: Tensor<i8>) -> Self {
970        Self::I8(t)
971    }
972}
973
974impl From<Tensor<u16>> for TensorDyn {
975    fn from(t: Tensor<u16>) -> Self {
976        Self::U16(t)
977    }
978}
979
980impl From<Tensor<i16>> for TensorDyn {
981    fn from(t: Tensor<i16>) -> Self {
982        Self::I16(t)
983    }
984}
985
986impl From<Tensor<u32>> for TensorDyn {
987    fn from(t: Tensor<u32>) -> Self {
988        Self::U32(t)
989    }
990}
991
992impl From<Tensor<i32>> for TensorDyn {
993    fn from(t: Tensor<i32>) -> Self {
994        Self::I32(t)
995    }
996}
997
998impl From<Tensor<u64>> for TensorDyn {
999    fn from(t: Tensor<u64>) -> Self {
1000        Self::U64(t)
1001    }
1002}
1003
1004impl From<Tensor<i64>> for TensorDyn {
1005    fn from(t: Tensor<i64>) -> Self {
1006        Self::I64(t)
1007    }
1008}
1009
1010impl From<Tensor<f16>> for TensorDyn {
1011    fn from(t: Tensor<f16>) -> Self {
1012        Self::F16(t)
1013    }
1014}
1015
1016impl From<Tensor<f32>> for TensorDyn {
1017    fn from(t: Tensor<f32>) -> Self {
1018        Self::F32(t)
1019    }
1020}
1021
1022impl From<Tensor<f64>> for TensorDyn {
1023    fn from(t: Tensor<f64>) -> Self {
1024        Self::F64(t)
1025    }
1026}
1027
1028impl fmt::Debug for TensorDyn {
1029    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1030        dispatch!(self, fmt, f)
1031    }
1032}
1033
1034#[cfg(test)]
1035mod tests {
1036    use super::*;
1037
1038    #[test]
1039    fn from_typed_tensor() {
1040        let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1041        let dyn_t: TensorDyn = t.into();
1042        assert_eq!(dyn_t.dtype(), DType::U8);
1043        assert_eq!(dyn_t.shape(), &[10]);
1044    }
1045
1046    #[test]
1047    fn from_foreign_ptr_wraps_borrowed_memory() {
1048        use crate::TensorMapTrait;
1049        // The CUDA zero-copy export shape: wrap an externally-allocated buffer as
1050        // a type-erased Mem tensor, with an owner that frees it on last drop.
1051        let mut vec: Vec<f32> = vec![0.0; 4];
1052        let ptr = vec.as_mut_ptr() as *mut u8;
1053        let owner: crate::ForeignOwner = Box::new(vec);
1054        let t = unsafe {
1055            TensorDyn::from_foreign_ptr(ptr, &[2, 2], DType::F32, Some(owner), Some("trt_output"))
1056        }
1057        .unwrap();
1058        assert_eq!(t.dtype(), DType::F32);
1059        assert_eq!(t.memory(), TensorMemory::Mem);
1060        assert_eq!(t.shape(), &[2, 2]);
1061        {
1062            let mut m = t.as_f32().unwrap().map().unwrap();
1063            m.as_mut_slice().copy_from_slice(&[1.0, 2.0, 3.0, 4.0]);
1064        }
1065        let m = t.as_f32().unwrap().map().unwrap();
1066        assert_eq!(m.as_slice(), &[1.0, 2.0, 3.0, 4.0]);
1067    }
1068
1069    // -------------------------------------------------------------------------
1070    // TensorDyn::from_foreign_ptr guard paths.
1071    //
1072    // The happy path (F32) is covered by `from_foreign_ptr_wraps_borrowed_memory`
1073    // above. These cells add the null-ptr, empty-shape, and overflow guards, plus
1074    // a U8 dtype to confirm the match-arm dispatch is exercised for integer types.
1075    // -------------------------------------------------------------------------
1076
1077    #[test]
1078    fn from_foreign_ptr_rejects_null_ptr() {
1079        let err = unsafe {
1080            TensorDyn::from_foreign_ptr(std::ptr::null_mut(), &[4], DType::U8, None, None)
1081        }
1082        .unwrap_err();
1083        // The null guard fires inside Tensor<u8>::from_foreign.
1084        assert!(
1085            matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("non-null")),
1086            "expected InvalidArgument(non-null), got {err:?}"
1087        );
1088    }
1089
1090    #[test]
1091    fn from_foreign_ptr_rejects_empty_shape() {
1092        let mut dummy: u8 = 0;
1093        let err = unsafe {
1094            TensorDyn::from_foreign_ptr(&mut dummy as *mut u8, &[], DType::U8, None, None)
1095        }
1096        .unwrap_err();
1097        assert!(
1098            matches!(err, crate::error::Error::InvalidSize(0)),
1099            "expected InvalidSize(0) for empty shape, got {err:?}"
1100        );
1101    }
1102
1103    #[test]
1104    fn from_foreign_ptr_rejects_overflow_shape() {
1105        let mut dummy: u8 = 0;
1106        let huge = [usize::MAX / 2 + 1, 2];
1107        let err = unsafe { TensorDyn::from_foreign_ptr(&mut dummy, &huge, DType::U8, None, None) }
1108            .unwrap_err();
1109        assert!(
1110            matches!(err, crate::error::Error::InvalidArgument(ref m) if m.contains("overflow")),
1111            "expected InvalidArgument(overflow), got {err:?}"
1112        );
1113    }
1114
1115    #[test]
1116    fn from_foreign_ptr_u8_dtype_dispatch() {
1117        // Exercises the U8 arm of from_foreign_ptr's match, which wraps
1118        // the raw pointer as Tensor<u8> and downcasts correctly.
1119        let mut buf: Vec<u8> = vec![1, 2, 3, 4];
1120        let ptr = buf.as_mut_ptr();
1121        let owner: crate::ForeignOwner = Box::new(buf);
1122        let t = unsafe {
1123            TensorDyn::from_foreign_ptr(ptr, &[4], DType::U8, Some(owner), Some("u8_foreign"))
1124        }
1125        .unwrap();
1126        assert_eq!(t.dtype(), DType::U8);
1127        assert_eq!(t.shape(), &[4]);
1128        let m = t.as_u8().unwrap().map().unwrap();
1129        use crate::TensorMapTrait;
1130        assert_eq!(m.as_slice(), &[1u8, 2, 3, 4]);
1131    }
1132
1133    #[test]
1134    fn downcast_ref() {
1135        let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1136        let dyn_t: TensorDyn = t.into();
1137        assert!(dyn_t.as_u8().is_some());
1138        assert!(dyn_t.as_i8().is_none());
1139    }
1140
1141    #[test]
1142    fn downcast_into() {
1143        let t = Tensor::<u8>::new(&[10], None, None).unwrap();
1144        let dyn_t: TensorDyn = t.into();
1145        let back = dyn_t.into_u8().unwrap();
1146        assert_eq!(back.shape(), &[10]);
1147    }
1148
1149    #[test]
1150    fn image_accessors() {
1151        let t = Tensor::<u8>::image(
1152            640,
1153            480,
1154            PixelFormat::Rgba,
1155            None,
1156            crate::CpuAccess::ReadWrite,
1157        )
1158        .unwrap();
1159        let dyn_t: TensorDyn = t.into();
1160        assert_eq!(dyn_t.format(), Some(PixelFormat::Rgba));
1161        assert_eq!(dyn_t.width(), Some(640));
1162        assert_eq!(dyn_t.height(), Some(480));
1163        assert!(!dyn_t.is_multiplane());
1164    }
1165
1166    #[test]
1167    fn image_constructor() {
1168        let dyn_t = TensorDyn::image(
1169            640,
1170            480,
1171            PixelFormat::Rgb,
1172            DType::U8,
1173            None,
1174            crate::CpuAccess::ReadWrite,
1175        )
1176        .unwrap();
1177        assert_eq!(dyn_t.dtype(), DType::U8);
1178        assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
1179        assert_eq!(dyn_t.width(), Some(640));
1180    }
1181
1182    #[test]
1183    fn image_constructor_i8() {
1184        let dyn_t = TensorDyn::image(
1185            640,
1186            480,
1187            PixelFormat::Rgb,
1188            DType::I8,
1189            None,
1190            crate::CpuAccess::ReadWrite,
1191        )
1192        .unwrap();
1193        assert_eq!(dyn_t.dtype(), DType::I8);
1194        assert_eq!(dyn_t.format(), Some(PixelFormat::Rgb));
1195    }
1196
1197    #[test]
1198    fn set_format_packed() {
1199        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1200        assert_eq!(t.format(), None);
1201        t.set_format(PixelFormat::Rgb).unwrap();
1202        assert_eq!(t.format(), Some(PixelFormat::Rgb));
1203        assert_eq!(t.width(), Some(640));
1204        assert_eq!(t.height(), Some(480));
1205    }
1206
1207    #[test]
1208    fn set_format_planar() {
1209        let mut t = TensorDyn::new(&[3, 480, 640], DType::U8, None, None).unwrap();
1210        t.set_format(PixelFormat::PlanarRgb).unwrap();
1211        assert_eq!(t.format(), Some(PixelFormat::PlanarRgb));
1212        assert_eq!(t.width(), Some(640));
1213        assert_eq!(t.height(), Some(480));
1214    }
1215
1216    #[test]
1217    fn set_format_rejects_wrong_shape() {
1218        let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1219        assert!(t.set_format(PixelFormat::Rgb).is_err());
1220    }
1221
1222    #[test]
1223    fn with_format_builder() {
1224        let t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1225            .unwrap()
1226            .with_format(PixelFormat::Rgba)
1227            .unwrap();
1228        assert_eq!(t.format(), Some(PixelFormat::Rgba));
1229        assert_eq!(t.width(), Some(640));
1230        assert_eq!(t.height(), Some(480));
1231    }
1232
1233    #[cfg(target_os = "linux")]
1234    #[test]
1235    fn dmabuf_clone_mem_tensor_fails() {
1236        let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1237        assert_eq!(t.memory(), TensorMemory::Mem);
1238        assert!(t.dmabuf_clone().is_err());
1239    }
1240
1241    #[cfg(target_os = "linux")]
1242    #[test]
1243    fn dmabuf_mem_tensor_fails() {
1244        let t = TensorDyn::new(&[480, 640, 3], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1245        assert!(t.dmabuf().is_err());
1246    }
1247
1248    #[test]
1249    fn set_format_semi_planar_nv12() {
1250        // 720 rows = 480 * 3/2 (NV12: height + height/2 for chroma)
1251        let mut t = TensorDyn::new(&[720, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1252        t.set_format(PixelFormat::Nv12).unwrap();
1253        assert_eq!(t.format(), Some(PixelFormat::Nv12));
1254        assert_eq!(t.width(), Some(640));
1255        assert_eq!(t.height(), Some(480));
1256    }
1257
1258    #[test]
1259    fn set_format_semi_planar_nv16() {
1260        // 960 rows = 480 * 2 (NV16: height + height for chroma)
1261        let mut t = TensorDyn::new(&[960, 640], DType::U8, Some(TensorMemory::Mem), None).unwrap();
1262        t.set_format(PixelFormat::Nv16).unwrap();
1263        assert_eq!(t.format(), Some(PixelFormat::Nv16));
1264        assert_eq!(t.width(), Some(640));
1265        assert_eq!(t.height(), Some(480));
1266    }
1267
1268    #[test]
1269    fn with_format_rejects_wrong_shape() {
1270        let result = TensorDyn::new(&[480, 640, 4], DType::U8, None, None)
1271            .unwrap()
1272            .with_format(PixelFormat::Rgb);
1273        assert!(result.is_err());
1274    }
1275
1276    #[test]
1277    fn set_format_preserved_after_rejection() {
1278        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1279        t.set_format(PixelFormat::Rgb).unwrap();
1280        assert_eq!(t.format(), Some(PixelFormat::Rgb));
1281
1282        // Rgba requires 4 channels, should fail on a 3-channel tensor
1283        assert!(t.set_format(PixelFormat::Rgba).is_err());
1284
1285        // Original format should be preserved
1286        assert_eq!(t.format(), Some(PixelFormat::Rgb));
1287    }
1288
1289    #[test]
1290    fn set_format_idempotent() {
1291        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1292        t.set_format(PixelFormat::Rgb).unwrap();
1293        t.set_format(PixelFormat::Rgb).unwrap();
1294        assert_eq!(t.format(), Some(PixelFormat::Rgb));
1295        assert_eq!(t.width(), Some(640));
1296        assert_eq!(t.height(), Some(480));
1297    }
1298
1299    // --- Row stride tests ---
1300
1301    #[test]
1302    fn set_row_stride_valid() {
1303        // RGBA 100px wide: min stride = 400, set 512
1304        let mut t = TensorDyn::image(
1305            100,
1306            100,
1307            PixelFormat::Rgba,
1308            DType::U8,
1309            None,
1310            crate::CpuAccess::ReadWrite,
1311        )
1312        .unwrap();
1313        t.set_row_stride(512).unwrap();
1314        assert_eq!(t.row_stride(), Some(512));
1315        assert_eq!(t.effective_row_stride(), Some(512));
1316    }
1317
1318    #[test]
1319    fn set_row_stride_equals_min() {
1320        // RGB 100px: min stride = 300, set exactly 300
1321        let mut t = TensorDyn::image(
1322            100,
1323            100,
1324            PixelFormat::Rgb,
1325            DType::U8,
1326            None,
1327            crate::CpuAccess::ReadWrite,
1328        )
1329        .unwrap();
1330        t.set_row_stride(300).unwrap();
1331        assert_eq!(t.row_stride(), Some(300));
1332    }
1333
1334    #[test]
1335    fn set_row_stride_too_small() {
1336        // RGBA 64px (a 64-aligned width: 64*4 = 256, already a multiple of 64)
1337        // carries no implicit stride. min stride = 256; setting 200 must error
1338        // and leave row_stride unset. (Non-64-aligned widths now record the
1339        // padded stride at allocation — see `Tensor::image`.)
1340        let mut t = TensorDyn::image(
1341            64,
1342            100,
1343            PixelFormat::Rgba,
1344            DType::U8,
1345            None,
1346            crate::CpuAccess::ReadWrite,
1347        )
1348        .unwrap();
1349        assert!(t.set_row_stride(200).is_err());
1350        assert_eq!(t.row_stride(), None);
1351    }
1352
1353    #[test]
1354    fn set_row_stride_zero() {
1355        let mut t = TensorDyn::image(
1356            100,
1357            100,
1358            PixelFormat::Rgb,
1359            DType::U8,
1360            None,
1361            crate::CpuAccess::ReadWrite,
1362        )
1363        .unwrap();
1364        assert!(t.set_row_stride(0).is_err());
1365    }
1366
1367    #[test]
1368    fn set_row_stride_requires_format() {
1369        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1370        assert!(t.set_row_stride(2048).is_err());
1371    }
1372
1373    #[test]
1374    fn effective_row_stride_without_stride() {
1375        // A 64-aligned-width packed image carries no explicit stride; the
1376        // effective stride falls back to the computed tight pitch. (Width 64
1377        // RGB → 64*3 = 192, already a multiple of 64, so no padding is added.
1378        // Non-aligned widths now record the padded stride — see `Tensor::image`.)
1379        let t = TensorDyn::image(
1380            64,
1381            100,
1382            PixelFormat::Rgb,
1383            DType::U8,
1384            None,
1385            crate::CpuAccess::ReadWrite,
1386        )
1387        .unwrap();
1388        assert_eq!(t.row_stride(), None);
1389        assert_eq!(t.effective_row_stride(), Some(192)); // 64 * 3
1390    }
1391
1392    #[test]
1393    fn effective_row_stride_padded_packed_dma() {
1394        // A non-64-aligned packed width on a DMA buffer records the 64-aligned
1395        // stride so the EGLImage import is accepted by Mali/Vivante (RGB 100px:
1396        // 100*3 = 300 → padded to 320). This padding is DMA-specific — host-only
1397        // memory keeps the tight pitch — so skip when DMA is unavailable (e.g. CI
1398        // without dma_heap); the behaviour is also validated on-target.
1399        let t = match TensorDyn::image(
1400            100,
1401            100,
1402            PixelFormat::Rgb,
1403            DType::U8,
1404            Some(TensorMemory::Dma),
1405            crate::CpuAccess::ReadWrite,
1406        ) {
1407            Ok(t) if t.memory() == TensorMemory::Dma => t,
1408            _ => return,
1409        };
1410        assert_eq!(t.row_stride(), Some(320));
1411        assert_eq!(t.effective_row_stride(), Some(320));
1412    }
1413
1414    #[test]
1415    fn effective_row_stride_no_format() {
1416        let t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1417        assert_eq!(t.effective_row_stride(), None);
1418    }
1419
1420    #[test]
1421    fn with_row_stride_builder() {
1422        let t = TensorDyn::image(
1423            100,
1424            100,
1425            PixelFormat::Rgba,
1426            DType::U8,
1427            None,
1428            crate::CpuAccess::ReadWrite,
1429        )
1430        .unwrap()
1431        .with_row_stride(512)
1432        .unwrap();
1433        assert_eq!(t.row_stride(), Some(512));
1434        assert_eq!(t.effective_row_stride(), Some(512));
1435    }
1436
1437    #[test]
1438    fn with_row_stride_rejects_small() {
1439        let result = TensorDyn::image(
1440            100,
1441            100,
1442            PixelFormat::Rgba,
1443            DType::U8,
1444            None,
1445            crate::CpuAccess::ReadWrite,
1446        )
1447        .unwrap()
1448        .with_row_stride(200);
1449        assert!(result.is_err());
1450    }
1451
1452    #[test]
1453    fn set_format_clears_row_stride() {
1454        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1455        t.set_format(PixelFormat::Rgb).unwrap();
1456        t.set_row_stride(2048).unwrap();
1457        assert_eq!(t.row_stride(), Some(2048));
1458
1459        // Incompatible format change (4-chan on 3-chan shape) fails — stride preserved
1460        let _ = t.set_format(PixelFormat::Bgra);
1461        assert_eq!(t.row_stride(), Some(2048));
1462
1463        // Re-set to same format — stride preserved
1464        t.set_format(PixelFormat::Rgb).unwrap();
1465        assert_eq!(t.row_stride(), Some(2048));
1466
1467        // Reshape clears format and stride
1468        t.reshape(&[480 * 640 * 3]).unwrap();
1469        assert_eq!(t.row_stride(), None);
1470        assert_eq!(t.format(), None);
1471    }
1472
1473    #[test]
1474    fn set_format_different_compatible_clears_stride() {
1475        // RGBA and BGRA are both 4-channel packed — switching between them
1476        // succeeds and must clear the stored stride.
1477        let mut t = TensorDyn::new(&[480, 640, 4], DType::U8, None, None).unwrap();
1478        t.set_format(PixelFormat::Rgba).unwrap();
1479        t.set_row_stride(4096).unwrap();
1480        assert_eq!(t.row_stride(), Some(4096));
1481
1482        // Successful format change to a different compatible format clears stride
1483        t.set_format(PixelFormat::Bgra).unwrap();
1484        assert_eq!(t.format(), Some(PixelFormat::Bgra));
1485        assert_eq!(t.row_stride(), None);
1486    }
1487
1488    #[test]
1489    fn set_format_same_preserves_stride() {
1490        let mut t = TensorDyn::image(
1491            100,
1492            100,
1493            PixelFormat::Rgb,
1494            DType::U8,
1495            None,
1496            crate::CpuAccess::ReadWrite,
1497        )
1498        .unwrap();
1499        t.set_row_stride(512).unwrap();
1500        // Re-setting the same format should not clear stride
1501        t.set_format(PixelFormat::Rgb).unwrap();
1502        assert_eq!(t.row_stride(), Some(512));
1503    }
1504
1505    #[test]
1506    fn effective_row_stride_planar() {
1507        let t = TensorDyn::image(
1508            640,
1509            480,
1510            PixelFormat::PlanarRgb,
1511            DType::U8,
1512            None,
1513            crate::CpuAccess::ReadWrite,
1514        )
1515        .unwrap();
1516        assert_eq!(t.effective_row_stride(), Some(640)); // planar: width only
1517    }
1518
1519    #[test]
1520    fn effective_row_stride_nv12() {
1521        let t = TensorDyn::image(
1522            640,
1523            480,
1524            PixelFormat::Nv12,
1525            DType::U8,
1526            None,
1527            crate::CpuAccess::ReadWrite,
1528        )
1529        .unwrap();
1530        assert_eq!(t.effective_row_stride(), Some(640)); // semi-planar: width only
1531    }
1532
1533    #[test]
1534    fn map_rejects_strided_tensor() {
1535        let mut t = Tensor::<u8>::image(
1536            100,
1537            100,
1538            PixelFormat::Rgba,
1539            Some(TensorMemory::Mem),
1540            crate::CpuAccess::ReadWrite,
1541        )
1542        .unwrap();
1543        // Map works before stride is set
1544        assert!(t.map().is_ok());
1545        // After setting stride, map should be rejected
1546        t.set_row_stride(512).unwrap();
1547        let err = t.map();
1548        assert!(err.is_err());
1549    }
1550
1551    // ── plane_offset tests ──────────────────────────────────────────
1552
1553    #[test]
1554    fn plane_offset_default_none() {
1555        let t = TensorDyn::image(
1556            100,
1557            100,
1558            PixelFormat::Rgba,
1559            DType::U8,
1560            None,
1561            crate::CpuAccess::ReadWrite,
1562        )
1563        .unwrap();
1564        assert_eq!(t.plane_offset(), None);
1565    }
1566
1567    #[test]
1568    fn set_plane_offset_basic() {
1569        let mut t = TensorDyn::image(
1570            100,
1571            100,
1572            PixelFormat::Rgba,
1573            DType::U8,
1574            None,
1575            crate::CpuAccess::ReadWrite,
1576        )
1577        .unwrap();
1578        t.set_plane_offset(4096);
1579        assert_eq!(t.plane_offset(), Some(4096));
1580    }
1581
1582    #[test]
1583    fn set_plane_offset_zero() {
1584        let mut t = TensorDyn::image(
1585            100,
1586            100,
1587            PixelFormat::Rgb,
1588            DType::U8,
1589            None,
1590            crate::CpuAccess::ReadWrite,
1591        )
1592        .unwrap();
1593        t.set_plane_offset(0);
1594        assert_eq!(t.plane_offset(), Some(0));
1595    }
1596
1597    #[test]
1598    fn set_plane_offset_no_format() {
1599        // plane_offset does not require format (it is format-independent)
1600        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1601        t.set_plane_offset(4096);
1602        assert_eq!(t.plane_offset(), Some(4096));
1603    }
1604
1605    #[test]
1606    fn set_format_clears_plane_offset() {
1607        let mut t = TensorDyn::new(&[480, 640, 3], DType::U8, None, None).unwrap();
1608        t.set_format(PixelFormat::Rgb).unwrap();
1609        t.set_plane_offset(4096);
1610        assert_eq!(t.plane_offset(), Some(4096));
1611
1612        // Re-set same format — offset preserved
1613        t.set_format(PixelFormat::Rgb).unwrap();
1614        assert_eq!(t.plane_offset(), Some(4096));
1615
1616        // Reshape clears everything
1617        t.reshape(&[480 * 640 * 3]).unwrap();
1618        assert_eq!(t.plane_offset(), None);
1619        assert_eq!(t.format(), None);
1620    }
1621
1622    #[test]
1623    fn map_rejects_out_of_bounds_offset() {
1624        let mut t = Tensor::<u8>::image(
1625            100,
1626            100,
1627            PixelFormat::Rgba,
1628            Some(TensorMemory::Mem),
1629            crate::CpuAccess::ReadWrite,
1630        )
1631        .unwrap();
1632        // Map works before offset is set.
1633        assert!(t.map().is_ok());
1634        // Heap offsets are now honored, but an offset that pushes the full
1635        // logical window (40000 bytes) past the allocation must be rejected.
1636        t.set_plane_offset(4096);
1637        assert!(t.map().is_err());
1638    }
1639
1640    #[test]
1641    fn mem_subview_in_bounds_maps_at_offset() {
1642        // An in-bounds heap sub-view now maps at its offset (previously every
1643        // non-zero heap offset was rejected outright).
1644        let parent = Tensor::<u8>::image(
1645            100,
1646            100,
1647            PixelFormat::Rgba,
1648            Some(TensorMemory::Mem),
1649            crate::CpuAccess::ReadWrite,
1650        )
1651        .unwrap();
1652        // A 10x10 RGBA window (400 bytes) at byte offset 4096 fits in 40000.
1653        let view = parent.subview(4096, &[10, 10, 4]).unwrap();
1654        assert_eq!(view.plane_offset(), Some(4096));
1655        assert!(view.map().is_ok());
1656    }
1657
1658    #[test]
1659    fn dyn_batch_dispatches_every_dtype() {
1660        // `TensorDyn::batch` fans out across all 11 dtype arms via `dyn_fanout!`;
1661        // exercise each so element `n` preserves the element type and shape.
1662        // A `[N=2, 4]` raw parent: element 1 is the contiguous 4-element window.
1663        use DType::*;
1664        for dt in [U8, I8, U16, I16, U32, I32, U64, I64, F16, F32, F64] {
1665            let parent = TensorDyn::new(&[2, 4], dt, Some(TensorMemory::Mem), None).unwrap();
1666            let view = parent.batch(1).unwrap();
1667            assert_eq!(view.dtype(), dt, "batch must preserve dtype {dt:?}");
1668            assert_eq!(view.shape(), &[4], "{dt:?}");
1669        }
1670    }
1671
1672    #[test]
1673    fn map_accepts_zero_offset_tensor() {
1674        let mut t = Tensor::<u8>::image(
1675            100,
1676            100,
1677            PixelFormat::Rgba,
1678            Some(TensorMemory::Mem),
1679            crate::CpuAccess::ReadWrite,
1680        )
1681        .unwrap();
1682        t.set_plane_offset(0);
1683        // Zero offset is fine for CPU mapping
1684        assert!(t.map().is_ok());
1685    }
1686
1687    #[test]
1688    fn dyn_configure_image_nv12() {
1689        let mut t = TensorDyn::image(
1690            640,
1691            480,
1692            PixelFormat::Rgb,
1693            DType::U8,
1694            None,
1695            crate::CpuAccess::ReadWrite,
1696        )
1697        .unwrap();
1698        t.configure_image(320, 240, PixelFormat::Nv12).unwrap();
1699        assert_eq!(t.format(), Some(PixelFormat::Nv12));
1700        assert_eq!((t.width(), t.height()), (Some(320), Some(240)));
1701    }
1702
1703    #[test]
1704    fn tensordyn_colorimetry_roundtrip() {
1705        use crate::{ColorEncoding, Colorimetry, DType, PixelFormat};
1706        let mut t = TensorDyn::image(
1707            1280,
1708            720,
1709            PixelFormat::Nv12,
1710            DType::U8,
1711            None,
1712            crate::CpuAccess::ReadWrite,
1713        )
1714        .unwrap();
1715        assert_eq!(t.colorimetry(), None);
1716        let c = Colorimetry::default().with_encoding(ColorEncoding::Bt709);
1717        t.set_colorimetry(Some(c));
1718        assert_eq!(t.colorimetry(), Some(c));
1719    }
1720
1721    #[test]
1722    fn from_planes_propagates_plane_offset() {
1723        let mut luma =
1724            Tensor::<u8>::new(&[480, 640], Some(TensorMemory::Mem), Some("luma")).unwrap();
1725        luma.set_plane_offset(4096);
1726        let chroma =
1727            Tensor::<u8>::new(&[240, 640], Some(TensorMemory::Mem), Some("chroma")).unwrap();
1728        let combined = Tensor::<u8>::from_planes(luma, chroma, PixelFormat::Nv12).unwrap();
1729        assert_eq!(combined.plane_offset(), Some(4096));
1730    }
1731
1732    #[test]
1733    fn cuda_passthrough_none_for_mem_tensor() {
1734        // Build a Mem-backed dynamic tensor the same way the other tests here do,
1735        // then confirm the CUDA accessors pass through to None (no handle).
1736        let t: TensorDyn = Tensor::<f32>::new(&[10], Some(TensorMemory::Mem), None)
1737            .unwrap()
1738            .into();
1739        assert!(t.cuda().is_none());
1740        assert!(t.cuda_map().is_none());
1741    }
1742}