Skip to main content

edgefirst_tensor/
pbo.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{
5    error::{Error, Result},
6    BufferIdentity, TensorMap, TensorMapTrait, TensorMemory, TensorTrait,
7};
8use log::trace;
9use num_traits::Num;
10use std::{
11    ffi::c_void,
12    fmt,
13    marker::PhantomData,
14    ops::{Deref, DerefMut},
15    ptr::NonNull,
16    sync::{Arc, Condvar, Mutex},
17};
18
19/// Raw mapped pointer from a PBO. CPU-accessible while the buffer is mapped.
20/// The pointer is only valid between map and unmap calls.
21pub struct PboMapping {
22    pub ptr: *mut u8,
23    pub size: usize,
24}
25
26// SAFETY: PboMapping is only created by PboOps::map_buffer which runs on the
27// GL thread, but the resulting pointer is used on the caller's thread. This is
28// safe because glMapBufferRange returns a CPU-visible pointer that can be
29// accessed from any thread while the buffer remains mapped.
30unsafe impl Send for PboMapping {}
31
32/// Trait for PBO GL operations, implemented by the image crate.
33///
34/// All methods are blocking — they send commands to the GL thread
35/// and wait for completion. Implementations must ensure GL context
36/// is current on the thread that executes the actual GL calls.
37///
38/// # Safety
39///
40/// Implementations must ensure:
41/// - `map_buffer` returns a valid, aligned pointer to `size` bytes of
42///   CPU-accessible memory that remains valid until `unmap_buffer` is called.
43/// - `unmap_buffer` invalidates the pointer and releases the mapping.
44/// - `delete_buffer` frees the GL buffer resources.
45pub unsafe trait PboOps: Send + Sync {
46    /// Map the PBO for CPU read/write access.
47    /// The returned PboMapping is valid until `unmap_buffer` is called.
48    fn map_buffer(&self, buffer_id: u32, size: usize) -> Result<PboMapping>;
49
50    /// Unmap a previously mapped PBO. Must be called before GL operations
51    /// on this buffer (GLES 3.0 requirement).
52    fn unmap_buffer(&self, buffer_id: u32) -> Result<()>;
53
54    /// Delete the PBO. Fire-and-forget — no reply needed.
55    /// Called from PboTensor's Drop impl.
56    fn delete_buffer(&self, buffer_id: u32);
57}
58
59/// Opaque handle to a PBO's GL resources.
60struct PboHandle {
61    ops: Arc<dyn PboOps>,
62    buffer_id: u32,
63    size: usize,
64    map_state: Mutex<MapState>,
65    /// Wakes threads waiting out an in-flight GL map/unmap. Needed because
66    /// `map_state` is deliberately NOT held across `PboOps` calls — see
67    /// [`PboHandle::acquire_map`].
68    map_cv: Condvar,
69}
70
71/// CPU-map state of a PBO's single GL buffer.
72///
73/// A GL buffer has exactly one mapping at a time, so these describe that
74/// one mapping's ownership rather than per-[`PboMap`] state:
75///
76/// * `Unmapped` — no CPU mapping exists.
77/// * `Exclusive` — one writable map holds it. Nothing else may map until
78///   that map drops.
79/// * `Shared` — `readers` read-only maps over ONE mapping, unmapped when
80///   the last of them drops.
81///
82/// Read sharing exists because tiled (SAHI) pre-processing has several
83/// worker threads convert different crops of the SAME source tensor
84/// concurrently, each taking `map_read()`. Read-only holders cannot observe
85/// one another's writes because there are none, so one mapping serves them
86/// all. Under a single-map rule every reader but one failed with
87/// [`Error::PboMapped`], which left callers no option but to serialize the
88/// workers that share a source.
89enum MapState {
90    Unmapped,
91    /// A thread is inside `PboOps::map_buffer` for this buffer right now.
92    /// Others must wait it out rather than start a second GL map.
93    Mapping,
94    Exclusive,
95    Shared {
96        readers: usize,
97        ptr: PboPtr,
98    },
99    /// A thread is inside `PboOps::unmap_buffer` for this buffer right now.
100    Unmapping,
101}
102
103impl PboHandle {
104    /// Acquire a CPU mapping of this buffer, returning the base pointer.
105    ///
106    /// Read-only acquisitions join an existing read-only mapping (bumping
107    /// its refcount) or create one. Writable acquisitions are exclusive.
108    /// Returns [`Error::PboMapped`] whenever the request cannot share what
109    /// is already held: anything against an exclusive map, or a writable
110    /// request against readers.
111    ///
112    /// # Locking
113    ///
114    /// `map_state` is NEVER held across a [`PboOps`] call. Those calls are
115    /// blocking round-trips to the GL thread, and the work the GL thread
116    /// runs (a convert, say) maps tensors itself — so a thread that waited
117    /// on GL while holding this mutex would deadlock against the very GL
118    /// thread it is waiting for. The in-flight GL call is published as
119    /// [`MapState::Mapping`] / [`MapState::Unmapping`] instead, and other
120    /// threads wait on `map_cv` until it resolves.
121    fn acquire_map(&self, writable: bool) -> Result<PboPtr> {
122        loop {
123            let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
124            // Wait out any GL map/unmap another thread is mid-way through.
125            while matches!(*guard, MapState::Mapping | MapState::Unmapping) {
126                guard = self
127                    .map_cv
128                    .wait(guard)
129                    .expect("PBO map state mutex poisoned");
130            }
131            match &mut *guard {
132                MapState::Shared { readers, ptr } if !writable => {
133                    *readers += 1;
134                    return Ok(PboPtr(ptr.0));
135                }
136                MapState::Shared { .. } | MapState::Exclusive => return Err(Error::PboMapped),
137                // Ruled out by the wait above; re-check rather than assume.
138                MapState::Mapping | MapState::Unmapping => continue,
139                MapState::Unmapped => *guard = MapState::Mapping,
140            }
141            drop(guard);
142
143            // No lock held here — see this function's Locking note.
144            let mapped = self.ops.map_buffer(self.buffer_id, self.size);
145            let base = match mapped {
146                Ok(mapping) => NonNull::new(mapping.ptr as *mut c_void),
147                Err(e) => {
148                    self.finish_map(None, writable);
149                    return Err(e);
150                }
151            };
152            let Some(base) = base else {
153                // GL reported success but handed back a null pointer. It
154                // considers the buffer mapped, so release it rather than
155                // strand it mapped forever.
156                //
157                // Order matters: the cleanup unmap runs while this thread
158                // still holds the `Mapping` claim, so waiters stay parked
159                // and cannot start a fresh `map_buffer` that would race
160                // this unmap on the GL queue. Only once it has completed is
161                // `Unmapped` published.
162                let _ = self.ops.unmap_buffer(self.buffer_id);
163                self.finish_map(None, writable);
164                return Err(Error::InvalidSize(self.size));
165            };
166            self.finish_map(Some(base), writable);
167            return Ok(PboPtr(base));
168        }
169    }
170
171    /// Publish the outcome of a GL map this thread had claimed via
172    /// [`MapState::Mapping`], waking anyone who waited on it. `None` means
173    /// the map failed and the buffer is unmapped again.
174    fn finish_map(&self, base: Option<NonNull<c_void>>, writable: bool) {
175        let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
176        *guard = match base {
177            None => MapState::Unmapped,
178            Some(_) if writable => MapState::Exclusive,
179            Some(base) => MapState::Shared {
180                readers: 1,
181                ptr: PboPtr(base),
182            },
183        };
184        self.map_cv.notify_all();
185    }
186
187    /// Release one acquisition, unmapping the GL buffer once the last
188    /// holder is gone. Same locking rule as [`Self::acquire_map`]: the GL
189    /// unmap runs with no lock held.
190    fn release_map(&self) {
191        let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
192        let last_holder = match &mut *guard {
193            MapState::Shared { readers, .. } => {
194                *readers = readers.saturating_sub(1);
195                *readers == 0
196            }
197            MapState::Exclusive => true,
198            // Nothing held — a release with no matching acquire. (A caller
199            // holding an acquisition cannot observe Mapping/Unmapping:
200            // only the last holder enters Unmapping, and it is this one.)
201            MapState::Unmapped | MapState::Mapping | MapState::Unmapping => false,
202        };
203        if !last_holder {
204            return;
205        }
206        *guard = MapState::Unmapping;
207        drop(guard);
208
209        trace!("Unmapping PBO buffer_id={}", self.buffer_id);
210        if let Err(e) = self.ops.unmap_buffer(self.buffer_id) {
211            log::warn!("Failed to unmap PBO buffer {}: {e}", self.buffer_id);
212        }
213
214        let mut guard = self.map_state.lock().expect("PBO map state mutex poisoned");
215        *guard = MapState::Unmapped;
216        self.map_cv.notify_all();
217    }
218
219    /// Whether any CPU mapping is currently held or being established.
220    fn is_mapped(&self) -> bool {
221        !matches!(
222            *self.map_state.lock().expect("PBO map state mutex poisoned"),
223            MapState::Unmapped
224        )
225    }
226}
227
228impl Drop for PboHandle {
229    fn drop(&mut self) {
230        self.ops.delete_buffer(self.buffer_id);
231    }
232}
233
234/// A tensor backed by an OpenGL Pixel Buffer Object.
235pub struct PboTensor<T>
236where
237    T: Num + Clone + fmt::Debug + Send + Sync,
238{
239    pub name: String,
240    pub shape: Vec<usize>,
241    handle: Arc<PboHandle>,
242    identity: BufferIdentity,
243    /// Byte offset of this tensor's window into the shared GL buffer. Non-zero
244    /// only for sub-views (`view`/`batch`), which share the `Arc<PboHandle>` and
245    /// `BufferIdentity` and address a sub-region by this offset — mirrors
246    /// `DmaTensor::mmap_offset` / `IoSurfaceTensor::view_offset`.
247    pub(crate) view_offset: usize,
248    _marker: PhantomData<T>,
249}
250
251impl<T> fmt::Debug for PboTensor<T>
252where
253    T: Num + Clone + fmt::Debug + Send + Sync,
254{
255    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
256        f.debug_struct("PboTensor")
257            .field("name", &self.name)
258            .field("shape", &self.shape)
259            .field("buffer_id", &self.handle.buffer_id)
260            .field("size", &self.handle.size)
261            .finish()
262    }
263}
264
265unsafe impl<T> Send for PboTensor<T> where T: Num + Clone + fmt::Debug + Send + Sync {}
266unsafe impl<T> Sync for PboTensor<T> where T: Num + Clone + fmt::Debug + Send + Sync {}
267
268impl<T> PboTensor<T>
269where
270    T: Num + Clone + fmt::Debug + Send + Sync,
271{
272    /// Create a new PBO tensor from an already-allocated GL buffer.
273    ///
274    /// Called by the image crate after creating the PBO on the GL thread.
275    /// Users should not call this directly — use `ImageProcessor::create_image()`.
276    ///
277    /// # Errors
278    ///
279    /// Returns `Error::ShapeMismatch` if `size` does not equal
280    /// `shape.iter().product::<usize>() * std::mem::size_of::<T>()`.
281    /// Returns `Error::InvalidSize` if `size` is zero.
282    pub fn from_pbo(
283        buffer_id: u32,
284        size: usize,
285        shape: &[usize],
286        name: Option<&str>,
287        ops: Arc<dyn PboOps>,
288    ) -> Result<Self> {
289        if size == 0 {
290            return Err(Error::InvalidSize(0));
291        }
292        let expected = shape.iter().product::<usize>() * std::mem::size_of::<T>();
293        // Allow `size >= expected`: PBOs allocated with a 64-byte-aligned row
294        // stride may be larger than the shape product.  Reject only if the
295        // allocation is strictly smaller than the logical content.
296        if size < expected {
297            return Err(Error::ShapeMismatch(format!(
298                "PBO size {size} is smaller than shape {shape:?} * sizeof({}) = {expected}",
299                std::any::type_name::<T>(),
300            )));
301        }
302        let name = name.unwrap_or("pbo_tensor").to_owned();
303        Ok(Self {
304            name,
305            shape: shape.to_vec(),
306            handle: Arc::new(PboHandle {
307                ops,
308                buffer_id,
309                size,
310                map_state: Mutex::new(MapState::Unmapped),
311                map_cv: Condvar::new(),
312            }),
313            identity: BufferIdentity::new(),
314            view_offset: 0,
315            _marker: PhantomData,
316        })
317    }
318
319    /// Returns the GL buffer ID for this PBO.
320    pub fn buffer_id(&self) -> u32 {
321        self.handle.buffer_id
322    }
323
324    /// Returns true unless the PBO is fully unmapped.
325    ///
326    /// That covers an exclusive map, one or more read-only holders, AND a
327    /// map or unmap another thread currently has in flight — during those
328    /// windows the GL buffer is (or is about to be) mapped, and no other
329    /// mapping may begin, so reporting `false` would be misleading. Use it
330    /// as "is this buffer free for GL operations?", not as "does a CPU
331    /// pointer exist right now".
332    pub fn is_mapped(&self) -> bool {
333        self.handle.is_mapped()
334    }
335}
336
337impl<T> TensorTrait<T> for PboTensor<T>
338where
339    T: Num + Clone + fmt::Debug + Send + Sync,
340{
341    fn new(_shape: &[usize], _name: Option<&str>) -> Result<Self> {
342        Err(Error::NotImplemented(
343            "PboTensor cannot be created directly — use ImageProcessor::create_image()".to_owned(),
344        ))
345    }
346
347    #[cfg(unix)]
348    fn from_fd(_fd: std::os::fd::OwnedFd, _shape: &[usize], _name: Option<&str>) -> Result<Self> {
349        Err(Error::NotImplemented(
350            "PboTensor does not support from_fd".to_owned(),
351        ))
352    }
353
354    #[cfg(unix)]
355    fn clone_fd(&self) -> Result<std::os::fd::OwnedFd> {
356        Err(Error::NotImplemented(
357            "PboTensor does not support clone_fd".to_owned(),
358        ))
359    }
360
361    fn memory(&self) -> TensorMemory {
362        TensorMemory::Pbo
363    }
364
365    fn name(&self) -> String {
366        self.name.clone()
367    }
368
369    fn shape(&self) -> &[usize] {
370        &self.shape
371    }
372
373    fn reshape(&mut self, shape: &[usize]) -> Result<()> {
374        if shape.is_empty() {
375            return Err(Error::InvalidSize(0));
376        }
377        let new_size = shape.iter().product::<usize>() * std::mem::size_of::<T>();
378        if new_size != self.handle.size {
379            return Err(Error::ShapeMismatch(format!(
380                "Cannot reshape incompatible shape: {:?} to {:?}",
381                self.shape, shape
382            )));
383        }
384        self.shape = shape.to_vec();
385        self.view_offset = 0;
386        Ok(())
387    }
388
389    fn capacity_bytes(&self) -> usize {
390        self.handle.size
391    }
392
393    /// Capacity-based reconfigure (mirrors Mem/Shm/DMA/IOSurface): allow any
394    /// shape whose byte size fits the PBO allocation, so an oversized reusable
395    /// pool can be `configure_image`d to a smaller image. Without this PBO fell
396    /// back to the strict-`reshape` default and rejected pool reuse.
397    fn set_logical_shape(&mut self, shape: &[usize]) -> Result<()> {
398        if shape.is_empty() {
399            return Err(Error::InvalidSize(0));
400        }
401        let needed = shape.iter().product::<usize>() * std::mem::size_of::<T>();
402        if needed > self.handle.size {
403            return Err(Error::InsufficientCapacity {
404                needed,
405                capacity: self.handle.size,
406            });
407        }
408        self.shape = shape.to_vec();
409        Ok(())
410    }
411
412    fn map_with(&self, access: crate::CpuAccess) -> Result<TensorMap<T>> {
413        self.map_internal(None, access)
414    }
415
416    fn buffer_identity(&self) -> &BufferIdentity {
417        &self.identity
418    }
419
420    /// Zero-copy sub-region view sharing this PBO's GL buffer (via the
421    /// `Arc<PboHandle>`) and [`BufferIdentity`], positioned at `offset_bytes`
422    /// from this tensor's own window with logical `shape`. The GL backend keys
423    /// the import on the shared identity and addresses the window via
424    /// `glViewport` / the staged copy; a CPU map adds `view_offset` to the
425    /// mapped base. Mirrors [`DmaTensor::view`](crate::TensorTrait::view).
426    ///
427    /// # Errors
428    ///
429    /// - [`Error::InvalidOperation`] if `offset_bytes` is mis-aligned for `T`.
430    /// - [`Error::InsufficientCapacity`] if the window exceeds the allocation.
431    fn view(&self, offset_bytes: usize, shape: &[usize]) -> Result<Self> {
432        if !offset_bytes.is_multiple_of(std::mem::align_of::<T>()) {
433            return Err(Error::InvalidOperation(format!(
434                "PboTensor::view: offset {offset_bytes} not aligned to align_of::<T>()={}",
435                std::mem::align_of::<T>()
436            )));
437        }
438        let abs_offset = self
439            .view_offset
440            .checked_add(offset_bytes)
441            .ok_or(Error::InvalidSize(offset_bytes))?;
442        let logical = shape.iter().product::<usize>() * std::mem::size_of::<T>();
443        let needed = abs_offset
444            .checked_add(logical)
445            .ok_or(Error::InvalidSize(logical))?;
446        if needed > self.handle.size {
447            return Err(Error::InsufficientCapacity {
448                needed,
449                capacity: self.handle.size,
450            });
451        }
452        Ok(Self {
453            name: self.name.clone(),
454            shape: shape.to_vec(),
455            handle: Arc::clone(&self.handle),
456            identity: self.identity.clone(),
457            view_offset: abs_offset,
458            _marker: PhantomData,
459        })
460    }
461}
462
463impl<T> PboTensor<T>
464where
465    T: Num + Clone + fmt::Debug + Send + Sync,
466{
467    /// Map the PBO so `as_slice()` exposes the full padded buffer (`byte_size`
468    /// bytes) rather than the shape-derived logical count. Mirrors
469    /// [`DmaTensor::map_with_byte_size`]: a CPU producer (e.g. the JPEG decoder)
470    /// or a strided convert source iterates rows via `effective_row_stride()`
471    /// without running past the slice. Crate-private; the only caller is
472    /// `Tensor::map()`, which already checks `byte_size <= capacity_bytes()`.
473    pub(crate) fn map_with_byte_size(
474        &self,
475        byte_size: usize,
476        access: crate::CpuAccess,
477    ) -> Result<TensorMap<T>> {
478        self.map_internal(Some(byte_size), access)
479    }
480
481    fn map_internal(
482        &self,
483        byte_size_override: Option<usize>,
484        access: crate::CpuAccess,
485    ) -> Result<TensorMap<T>> {
486        // Always map the full GL allocation (`handle.size`); the slice length is
487        // narrowed by `byte_size_override` (or the logical shape) at access time.
488        let writable = access.writes();
489        let ptr = self.handle.acquire_map(writable)?;
490        Ok(TensorMap::Pbo(PboMap {
491            ptr: Arc::new(Mutex::new(ptr)),
492            shape: self.shape.clone(),
493            handle: Arc::clone(&self.handle),
494            byte_size_override,
495            view_offset: self.view_offset,
496            writable,
497            released: false,
498            _marker: PhantomData,
499        }))
500    }
501}
502
503// -- PboMap --
504
505#[derive(Debug)]
506struct PboPtr(NonNull<c_void>);
507
508impl Deref for PboPtr {
509    type Target = NonNull<c_void>;
510    fn deref(&self) -> &Self::Target {
511        &self.0
512    }
513}
514
515unsafe impl Send for PboPtr {}
516
517pub struct PboMap<T>
518where
519    T: Num + Clone + fmt::Debug,
520{
521    ptr: Arc<Mutex<PboPtr>>,
522    shape: Vec<usize>,
523    handle: Arc<PboHandle>,
524    /// Optional override for `as_slice().len()`. `None` → `shape.product()`
525    /// elements (logical view). `Some(bytes)` → `bytes / sizeof(T)` elements,
526    /// exposing the full padded GL allocation so callers can iterate rows via
527    /// `row_stride` (set by `Tensor::map()` for strided PBO tensors). Mirrors
528    /// `DmaMap::byte_size_override`.
529    byte_size_override: Option<usize>,
530    /// Byte offset of the sub-view window into the mapped GL buffer. `as_slice`
531    /// advances the base pointer by this many bytes before exposing the slice.
532    view_offset: usize,
533    /// Whether mutable access is permitted (`map_read()` maps are not).
534    /// The GL mapping itself stays MAP_READ|MAP_WRITE (bit narrowing is a
535    /// follow-up); this enforces the API contract uniformly.
536    writable: bool,
537    /// Whether this map has already released its acquisition. `unmap()` is
538    /// public on the trait AND called from `Drop`, so without this an
539    /// explicit unmap followed by the drop would release twice — which,
540    /// now that read-only maps are refcounted, would drop a *sibling*
541    /// reader's mapping out from under it rather than being the harmless
542    /// double-store it used to be.
543    released: bool,
544    _marker: PhantomData<T>,
545}
546
547impl<T> fmt::Debug for PboMap<T>
548where
549    T: Num + Clone + fmt::Debug,
550{
551    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
552        f.debug_struct("PboMap")
553            .field("shape", &self.shape)
554            .field("buffer_id", &self.handle.buffer_id)
555            .finish()
556    }
557}
558
559impl<T> TensorMapTrait<T> for PboMap<T>
560where
561    T: Num + Clone + fmt::Debug,
562{
563    fn shape(&self) -> &[usize] {
564        &self.shape
565    }
566
567    fn unmap(&mut self) {
568        if self.released {
569            return;
570        }
571        self.released = true;
572        self.handle.release_map();
573    }
574
575    fn as_slice(&self) -> &[T] {
576        let ptr = self.ptr.lock().expect("Failed to lock PboMap pointer");
577        let base = unsafe { (ptr.as_ptr() as *const u8).add(self.view_offset) as *const T };
578        unsafe { std::slice::from_raw_parts(base, self.slice_len_elems()) }
579    }
580
581    fn as_mut_slice(&mut self) -> &mut [T] {
582        crate::assert_map_writable(self.writable, "Pbo");
583        let ptr = self.ptr.lock().expect("Failed to lock PboMap pointer");
584        let base = unsafe { (ptr.as_ptr() as *mut u8).add(self.view_offset) as *mut T };
585        unsafe { std::slice::from_raw_parts_mut(base, self.slice_len_elems()) }
586    }
587}
588
589impl<T> PboMap<T>
590where
591    T: Num + Clone + fmt::Debug,
592{
593    /// Number of `T` elements exposed by `as_slice()`. Honours
594    /// `byte_size_override` (the full padded GL allocation) when set; otherwise
595    /// the shape-derived logical count. Mirrors `DmaMap::slice_len_elems`.
596    fn slice_len_elems(&self) -> usize {
597        match self.byte_size_override {
598            Some(bytes) => bytes / std::mem::size_of::<T>(),
599            None => self.shape.iter().product(),
600        }
601    }
602}
603
604impl<T> Deref for PboMap<T>
605where
606    T: Num + Clone + fmt::Debug,
607{
608    type Target = [T];
609    fn deref(&self) -> &[T] {
610        self.as_slice()
611    }
612}
613
614impl<T> DerefMut for PboMap<T>
615where
616    T: Num + Clone + fmt::Debug,
617{
618    fn deref_mut(&mut self) -> &mut [T] {
619        self.as_mut_slice()
620    }
621}
622
623impl<T> Drop for PboMap<T>
624where
625    T: Num + Clone + fmt::Debug,
626{
627    fn drop(&mut self) {
628        self.unmap();
629    }
630}
631
632impl<T> Clone for PboTensor<T>
633where
634    T: Num + Clone + fmt::Debug + Send + Sync,
635{
636    fn clone(&self) -> Self {
637        Self {
638            name: self.name.clone(),
639            shape: self.shape.clone(),
640            handle: Arc::clone(&self.handle),
641            identity: self.identity.clone(),
642            view_offset: self.view_offset,
643            _marker: PhantomData,
644        }
645    }
646}
647
648#[cfg(test)]
649mod tests {
650    use super::*;
651
652    /// Mock PboOps that uses a Vec<u8> as backing storage instead of GL.
653    ///
654    /// Counts map/unmap calls so tests can assert that N read-only holders
655    /// share ONE GL mapping rather than each taking their own.
656    struct MockPboOps {
657        storage: Mutex<Vec<u8>>,
658        maps: std::sync::atomic::AtomicUsize,
659        unmaps: std::sync::atomic::AtomicUsize,
660    }
661
662    impl MockPboOps {
663        fn new(size: usize) -> Arc<Self> {
664            Arc::new(Self {
665                storage: Mutex::new(vec![0u8; size]),
666                maps: std::sync::atomic::AtomicUsize::new(0),
667                unmaps: std::sync::atomic::AtomicUsize::new(0),
668            })
669        }
670
671        fn map_count(&self) -> usize {
672            self.maps.load(std::sync::atomic::Ordering::Acquire)
673        }
674
675        fn unmap_count(&self) -> usize {
676            self.unmaps.load(std::sync::atomic::Ordering::Acquire)
677        }
678    }
679
680    // SAFETY: the returned pointer addresses a `Vec<u8>` that is allocated
681    // once in `new` and never resized, so it stays valid for the whole life
682    // of the `MockPboOps` — which outlives every map handed out, since the
683    // tensor holds an `Arc` to it. The storage mutex is NOT what keeps the
684    // pointer alive (`map_buffer` releases it before returning); it only
685    // guards the length assertion.
686    unsafe impl PboOps for MockPboOps {
687        fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
688            self.maps.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
689            let storage = self.storage.lock().expect("lock");
690            assert_eq!(storage.len(), size);
691            Ok(PboMapping {
692                ptr: storage.as_ptr() as *mut u8,
693                size,
694            })
695        }
696
697        fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
698            self.unmaps
699                .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
700            Ok(())
701        }
702
703        fn delete_buffer(&self, _buffer_id: u32) {}
704    }
705
706    #[test]
707    fn test_pbo_tensor_create_and_metadata() {
708        let ops = MockPboOps::new(24);
709        let tensor = PboTensor::<u8>::from_pbo(42, 24, &[2, 3, 4], Some("test_pbo"), ops).unwrap();
710        assert_eq!(tensor.memory(), TensorMemory::Pbo);
711        assert_eq!(tensor.name(), "test_pbo");
712        assert_eq!(tensor.shape(), &[2, 3, 4]);
713        assert_eq!(tensor.buffer_id(), 42);
714        assert!(!tensor.is_mapped());
715    }
716
717    #[test]
718    fn test_pbo_tensor_map_write_read() {
719        let ops = MockPboOps::new(12);
720        let tensor = PboTensor::<u8>::from_pbo(1, 12, &[3, 4], Some("rw_test"), ops).unwrap();
721        {
722            let mut map = tensor.map().expect("map should succeed");
723            assert_eq!(map.shape(), &[3, 4]);
724            assert!(tensor.is_mapped());
725            map.as_mut_slice().fill(0xAB);
726            assert!(map.as_slice().iter().all(|&b| b == 0xAB));
727        }
728        assert!(!tensor.is_mapped());
729    }
730
731    #[test]
732    fn test_pbo_tensor_double_map_fails() {
733        let ops = MockPboOps::new(8);
734        let tensor = PboTensor::<u8>::from_pbo(2, 8, &[8], None, ops).unwrap();
735        let _map1 = tensor.map().expect("first map should succeed");
736        assert!(tensor.is_mapped());
737        let result = tensor.map();
738        assert!(result.is_err(), "second map while mapped should fail");
739    }
740
741    /// Mock whose `map_buffer` reports success but hands back a null
742    /// pointer — the one path where a claimed map has to be released again.
743    #[derive(Default)]
744    struct NullMappingOps {
745        unmaps: std::sync::atomic::AtomicUsize,
746    }
747
748    impl NullMappingOps {
749        fn unmap_count(&self) -> usize {
750            self.unmaps.load(std::sync::atomic::Ordering::Acquire)
751        }
752    }
753
754    // SAFETY: never produces a usable pointer, so no caller can dereference
755    // one; the mapping it returns is rejected before reaching a slice.
756    unsafe impl PboOps for NullMappingOps {
757        fn map_buffer(&self, _buffer_id: u32, size: usize) -> Result<PboMapping> {
758            Ok(PboMapping {
759                ptr: std::ptr::null_mut(),
760                size,
761            })
762        }
763
764        fn unmap_buffer(&self, _buffer_id: u32) -> Result<()> {
765            self.unmaps
766                .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
767            Ok(())
768        }
769
770        fn delete_buffer(&self, _buffer_id: u32) {}
771    }
772
773    /// A map that "succeeds" with a null pointer leaves GL considering the
774    /// buffer mapped, so it must be released and the state returned to
775    /// unmapped — otherwise the buffer is stranded and every later map
776    /// fails. The release happens before `Unmapped` is published, so no
777    /// waiter can start a map that races the cleanup unmap.
778    #[test]
779    fn null_mapping_releases_the_buffer_and_resets_state() {
780        let ops = Arc::new(NullMappingOps::default());
781        let dyn_ops: Arc<dyn PboOps> = ops.clone();
782        let tensor = PboTensor::<u8>::from_pbo(12, 8, &[8], None, dyn_ops).unwrap();
783
784        let Err(err) = tensor.map_read() else {
785            panic!("a null mapping must not succeed");
786        };
787        assert!(
788            matches!(err, Error::InvalidSize(_)),
789            "unexpected error: {err:?}"
790        );
791        assert_eq!(
792            ops.unmap_count(),
793            1,
794            "GL still holds the buffer mapped — the failed map must release it"
795        );
796        assert!(
797            !tensor.is_mapped(),
798            "state must return to Unmapped so later maps can proceed"
799        );
800
801        // Not stranded: the buffer can be claimed again.
802        let Err(err) = tensor.map_read() else {
803            panic!("still null, so it must still fail");
804        };
805        assert!(matches!(err, Error::InvalidSize(_)));
806        assert_eq!(ops.unmap_count(), 2, "and released again");
807    }
808
809    /// Several read-only holders share ONE GL mapping. This is what lets
810    /// tiled pre-processing run several workers over one source tensor;
811    /// before read sharing, every reader after the first failed.
812    #[test]
813    fn read_maps_share_one_gl_mapping() {
814        let ops = MockPboOps::new(8);
815        let dyn_ops: Arc<dyn PboOps> = ops.clone();
816        let tensor = PboTensor::<u8>::from_pbo(9, 8, &[8], None, dyn_ops).unwrap();
817
818        let r1 = tensor.map_read().expect("first read map");
819        let r2 = tensor
820            .map_read()
821            .expect("second read map must share, not fail");
822        let r3 = tensor
823            .map_read()
824            .expect("third read map must share, not fail");
825
826        assert_eq!(ops.map_count(), 1, "one GL mapping serves every reader");
827        assert_eq!(ops.unmap_count(), 0, "nothing unmapped while readers live");
828        assert!(tensor.is_mapped());
829        // All readers address the same bytes.
830        assert_eq!(r1.as_slice(), r2.as_slice());
831        assert_eq!(r2.as_slice(), r3.as_slice());
832
833        drop(r1);
834        drop(r2);
835        assert_eq!(
836            ops.unmap_count(),
837            0,
838            "the mapping outlives every reader but the last"
839        );
840        assert!(tensor.is_mapped());
841
842        drop(r3);
843        assert_eq!(ops.unmap_count(), 1, "last reader out unmaps, exactly once");
844        assert!(!tensor.is_mapped());
845    }
846
847    /// Read sharing must not weaken write exclusion in either direction.
848    #[test]
849    fn writers_and_readers_still_exclude_each_other() {
850        let ops = MockPboOps::new(8);
851        let dyn_ops: Arc<dyn PboOps> = ops.clone();
852        let tensor = PboTensor::<u8>::from_pbo(10, 8, &[8], None, dyn_ops).unwrap();
853
854        let reader = tensor.map_read().expect("read map");
855        assert!(
856            tensor.map().is_err(),
857            "a writer must not join an existing reader set"
858        );
859        drop(reader);
860
861        let writer = tensor.map().expect("write map after readers drained");
862        assert!(
863            tensor.map_read().is_err(),
864            "a reader must not join an exclusive writer"
865        );
866        assert!(
867            tensor.map().is_err(),
868            "a second writer must not join an exclusive writer"
869        );
870        drop(writer);
871
872        // Fully released — a fresh map of either kind succeeds again.
873        assert!(tensor.map_read().is_ok());
874    }
875
876    /// `unmap()` is public on the trait and also runs from `Drop`. With
877    /// refcounted readers a double release would unmap a SIBLING reader's
878    /// mapping, so each map must release at most once.
879    #[test]
880    fn explicit_unmap_then_drop_releases_only_once() {
881        let ops = MockPboOps::new(8);
882        let dyn_ops: Arc<dyn PboOps> = ops.clone();
883        let tensor = PboTensor::<u8>::from_pbo(11, 8, &[8], None, dyn_ops).unwrap();
884
885        let keeper = tensor.map_read().expect("reader that must survive");
886        let mut early = tensor.map_read().expect("reader unmapped by hand");
887        early.unmap();
888        drop(early); // Drop must NOT release a second time.
889
890        assert_eq!(
891            ops.unmap_count(),
892            0,
893            "one reader still holds the mapping — it must not be unmapped"
894        );
895        assert!(tensor.is_mapped(), "the surviving reader still holds it");
896        // The survivor's pointer is still valid and readable.
897        assert_eq!(keeper.as_slice().len(), 8);
898
899        drop(keeper);
900        assert_eq!(ops.unmap_count(), 1);
901        assert!(!tensor.is_mapped());
902    }
903
904    #[test]
905    fn test_pbo_tensor_reshape() {
906        let ops = MockPboOps::new(24);
907        let mut tensor = PboTensor::<u8>::from_pbo(3, 24, &[2, 3, 4], None, ops).unwrap();
908        tensor
909            .reshape(&[4, 6])
910            .expect("compatible reshape should succeed");
911        assert_eq!(tensor.shape(), &[4, 6]);
912        let result = tensor.reshape(&[100]);
913        assert!(result.is_err(), "incompatible reshape should fail");
914    }
915
916    #[test]
917    fn test_pbo_tensor_set_logical_shape_capacity_based() {
918        // A 24-byte PBO can be reconfigured to any shape that fits (unlike the
919        // strict `reshape`), so an oversized reusable pool can be
920        // `configure_image`d to a smaller image (the native-chroma decode pool).
921        let ops = MockPboOps::new(24);
922        let mut tensor = PboTensor::<u8>::from_pbo(7, 24, &[24], None, ops).unwrap();
923        // Smaller-than-capacity logical shape is accepted (reshape would reject).
924        tensor
925            .set_logical_shape(&[4, 5])
926            .expect("shape within capacity should succeed");
927        assert_eq!(tensor.shape(), &[4, 5]);
928        // Exactly-capacity is fine.
929        tensor.set_logical_shape(&[24]).unwrap();
930        // Over-capacity is rejected.
931        assert!(
932            tensor.set_logical_shape(&[5, 5]).is_err(),
933            "shape exceeding PBO capacity must be rejected"
934        );
935    }
936
937    #[test]
938    fn test_pbo_tensor_buffer_identity() {
939        let ops1 = MockPboOps::new(8);
940        let ops2 = MockPboOps::new(8);
941        let t1 = PboTensor::<u8>::from_pbo(1, 8, &[8], None, ops1).unwrap();
942        let t2 = PboTensor::<u8>::from_pbo(2, 8, &[8], None, ops2).unwrap();
943        assert_ne!(t1.buffer_identity().id(), t2.buffer_identity().id());
944    }
945
946    #[test]
947    fn test_pbo_tensor_new_returns_error() {
948        let result = PboTensor::<u8>::new(&[8], None);
949        assert!(result.is_err(), "PboTensor::new() should fail");
950    }
951
952    #[cfg(unix)]
953    #[test]
954    fn test_pbo_tensor_fd_ops_return_error() {
955        let ops = MockPboOps::new(8);
956        let tensor = PboTensor::<u8>::from_pbo(1, 8, &[8], None, ops).unwrap();
957        assert!(tensor.clone_fd().is_err());
958    }
959
960    #[test]
961    fn test_pbo_tensor_from_pbo_size_mismatch() {
962        let ops = MockPboOps::new(24);
963        let result = PboTensor::<u8>::from_pbo(1, 24, &[2, 3, 5], None, ops);
964        assert!(result.is_err(), "mismatched size/shape should fail");
965    }
966
967    #[test]
968    fn test_pbo_tensor_from_pbo_zero_size() {
969        let ops = MockPboOps::new(0);
970        let result = PboTensor::<u8>::from_pbo(1, 0, &[0], None, ops);
971        assert!(result.is_err(), "zero size should fail");
972    }
973
974    #[test]
975    fn test_pbo_via_tensor_enum() {
976        let ops = MockPboOps::new(12);
977        let pbo = PboTensor::<u8>::from_pbo(10, 12, &[3, 4], Some("enum_test"), ops).unwrap();
978        let tensor = crate::Tensor::wrap(crate::TensorStorage::Pbo(pbo));
979        assert_eq!(tensor.memory(), TensorMemory::Pbo);
980        assert_eq!(tensor.name(), "enum_test");
981        assert_eq!(tensor.shape(), &[3, 4]);
982        let mut map = tensor.map().expect("map via enum");
983        map.as_mut_slice().fill(42);
984        assert!(map.as_slice().iter().all(|&b| b == 42));
985    }
986}