vfxpreopenexr/deep/
deep_frame_buffer.rs

1use openexr_sys as sys;
2
3pub use crate::core::{
4    error::Error,
5    frame_buffer::{Frame, Slice, SliceRef},
6    refptr::{OpaquePtr, Ref, RefMut},
7    PixelType,
8};
9pub use imath_traits::{Bound2, Vec2};
10use std::marker::PhantomData;
11
12use std::ffi::{CStr, CString};
13
14type Result<T, E = Error> = std::result::Result<T, E>;
15
16pub struct DeepFrameBuffer {
17    pub(crate) ptr: *mut sys::Imf_DeepFrameBuffer_t,
18    pub(crate) sample_count_frame: Option<Frame>,
19    pub(crate) frames: Option<Vec<DeepFrame>>,
20}
21
22unsafe impl OpaquePtr for DeepFrameBuffer {
23    type SysPointee = sys::Imf_DeepFrameBuffer_t;
24    type Pointee = DeepFrameBuffer;
25}
26
27pub type DeepFrameBufferRef<'a> = Ref<'a, DeepFrameBuffer>;
28
29impl DeepFrameBuffer {
30    /// Create a new `DeepFrameBuffer`.
31    pub fn new() -> DeepFrameBuffer {
32        let mut ptr = std::ptr::null_mut();
33        unsafe {
34            sys::Imf_DeepFrameBuffer_ctor(&mut ptr);
35        }
36        DeepFrameBuffer {
37            ptr,
38            sample_count_frame: None,
39            frames: Some(Vec::new()),
40        }
41    }
42
43    /// Insert a [`DeepSlice`] into the `DeepFrameBuffer`.
44    ///
45    /// # Errors
46    /// * [`Error::InvalidArgument`] - if name is the empty string
47    ///
48    pub fn insert(&mut self, name: &str, slice: &DeepSlice) -> Result<()> {
49        let c_name =
50            CString::new(name).expect("Internal null bytes in filename");
51
52        unsafe {
53            sys::Imf_DeepFrameBuffer_insert(
54                self.ptr,
55                c_name.as_ptr(),
56                &slice.0,
57            )
58            .into_result()?;
59        }
60
61        Ok(())
62    }
63
64    /// Find the [`DeepSlice`] with the given `name` in the `DeepFrameBuffer`
65    ///
66    /// # Returns
67    /// * `Some([`DeepSliceRef`])` - if the [`DeepSlice`] is found
68    /// * `None` - otherwise
69    ///
70    pub fn get_slice<'a>(&'a self, name: &str) -> Option<DeepSliceRef<'a>> {
71        let c_name =
72            CString::new(name).expect("Internal null bytes in filename");
73
74        let mut ptr = std::ptr::null();
75        unsafe {
76            sys::Imf_DeepFrameBuffer_findSlice_const(
77                self.ptr,
78                &mut ptr,
79                c_name.as_ptr(),
80            );
81        }
82
83        if ptr.is_null() {
84            None
85        } else {
86            Some(DeepSliceRef::new(ptr))
87        }
88    }
89
90    /// Find the [`DeepSlice`] with the given `name` in the `DeepFrameBuffer` and get a
91    /// mutable Ref to it
92    ///
93    /// # Returns
94    /// * `Some([`DeepSliceRefMut`])` - if the [`DeepSlice`] is found
95    /// * `None` - otherwise
96    ///
97    pub fn get_slice_mut<'a>(
98        &'a mut self,
99        name: &str,
100    ) -> Option<DeepSliceRefMut<'a>> {
101        let c_name =
102            CString::new(name).expect("Internal null bytes in filename");
103
104        let mut ptr = std::ptr::null_mut();
105        unsafe {
106            sys::Imf_DeepFrameBuffer_findSlice(
107                self.ptr,
108                &mut ptr,
109                c_name.as_ptr(),
110            );
111        }
112
113        if ptr.is_null() {
114            None
115        } else {
116            Some(DeepSliceRefMut::new(ptr))
117        }
118    }
119
120    /// Get an iterator over the [`DeepSlice`]s in this `DeepFrameBuffer`
121    ///
122    pub fn iter(&self) -> DeepFrameBufferIter {
123        unsafe {
124            let mut ptr = sys::Imf_DeepFrameBuffer_ConstIterator_t::default();
125            sys::Imf_DeepFrameBuffer_begin_const(self.ptr, &mut ptr)
126                .into_result()
127                .unwrap();
128            let ptr = DeepFrameBufferConstIterator(ptr);
129
130            let mut end = sys::Imf_DeepFrameBuffer_ConstIterator_t::default();
131            sys::Imf_DeepFrameBuffer_end_const(self.ptr, &mut end)
132                .into_result()
133                .unwrap();
134            let end = DeepFrameBufferConstIterator(end);
135
136            DeepFrameBufferIter {
137                ptr,
138                end,
139                _p: PhantomData,
140            }
141        }
142    }
143
144    /// Insert a [`Slice`] to hold the sample count for each deep pixel.
145    ///
146    /// The slice must be of type [`PixelType::Uint`]
147    ///
148    /// # Errors
149    /// * [`Error::InvalidArgument`] - if the type of the slice is not [`PixelType::Uint`]
150    ///
151    pub fn set_sample_count_slice(
152        &mut self,
153        sample_count_slice: &Slice,
154    ) -> Result<()> {
155        unsafe {
156            sys::Imf_DeepFrameBuffer_insertSampleCountSlice(
157                self.ptr,
158                &sample_count_slice.0,
159            )
160            .into_result()?;
161        }
162
163        Ok(())
164    }
165
166    /// Set a [`Frame`](crate::core::frame_buffer::Frame) to hold the per-pixel sample
167    /// counts
168    ///
169    pub fn set_sample_count_frame(&mut self, frame: Frame) -> Result<()> {
170        let w = frame.data_window[2] - frame.data_window[0] + 1;
171        let ystride = w as usize * frame.stride;
172        self.set_sample_count_slice(
173            &Slice::with_data_window(
174                frame.channel_type,
175                frame.ptr,
176                frame.data_window,
177            )
178            .x_stride(frame.stride)
179            .y_stride(ystride)
180            .build()?,
181        )?;
182
183        self.sample_count_frame = Some(frame);
184
185        Ok(())
186    }
187
188    pub fn sample_count_frame(&self) -> Option<&Frame> {
189        self.sample_count_frame.as_ref()
190    }
191
192    /// Get the sample count slice
193    ///
194    pub fn sample_count_slice(&self) -> SliceRef {
195        let mut ptr = std::ptr::null();
196        unsafe {
197            sys::Imf_DeepFrameBuffer_getSampleCountSlice(self.ptr, &mut ptr);
198        }
199
200        SliceRef::new(ptr)
201    }
202
203    pub fn insert_deep_frame(&mut self, frame: DeepFrame) -> Result<()> {
204        let ptr = frame.ptr;
205        let data_window = frame.data_window;
206        let data_width = data_window[2] - data_window[0] + 1;
207        let offset = (-data_window[0] - data_window[1] * data_width) as isize;
208
209        unsafe {
210            self.insert(
211                &frame.channel_name,
212                &DeepSlice::builder(
213                    frame.channel_type,
214                    ptr.offset(offset) as *mut i8,
215                )
216                .x_stride(std::mem::size_of::<*const u8>())
217                .y_stride(
218                    std::mem::size_of::<*const u8>() * data_width as usize,
219                )
220                .sample_stride(frame.stride)
221                .build()?,
222            )?;
223        }
224
225        match &mut self.frames {
226            Some(v) => {
227                v.push(frame);
228            }
229            _ => unreachable!(),
230        }
231
232        Ok(())
233    }
234}
235
236impl Drop for DeepFrameBuffer {
237    fn drop(&mut self) {
238        unsafe {
239            sys::Imf_DeepFrameBuffer_dtor(self.ptr);
240        }
241    }
242}
243
244impl Default for DeepFrameBuffer {
245    fn default() -> Self {
246        DeepFrameBuffer::new()
247    }
248}
249
250#[repr(transparent)]
251#[derive(Clone)]
252pub(crate) struct DeepFrameBufferConstIterator(
253    pub(crate) sys::Imf_DeepFrameBuffer_ConstIterator_t,
254);
255
256// #[repr(transparent)]
257// pub(crate) struct DeepFrameBufferIterator(
258//     pub(crate) sys::Imf_DeepFrameBuffer_Iterator_t,
259// );
260
261pub struct DeepFrameBufferIter<'a> {
262    ptr: DeepFrameBufferConstIterator,
263    end: DeepFrameBufferConstIterator,
264    _p: PhantomData<&'a DeepFrameBuffer>,
265}
266
267impl<'a> Iterator for DeepFrameBufferIter<'a> {
268    type Item = (&'a str, DeepSliceRef<'a>);
269
270    fn next(&mut self) -> Option<(&'a str, DeepSliceRef<'a>)> {
271        let ptr_curr = self.ptr.clone();
272        let mut ptr_next = self.ptr.clone();
273        unsafe {
274            let mut dummy = std::ptr::null_mut();
275            sys::Imf_DeepFrameBuffer_ConstIterator__op_inc(
276                &mut ptr_next.0,
277                &mut dummy,
278            )
279            .into_result()
280            .unwrap();
281        }
282
283        if ptr_curr == self.end {
284            None
285        } else {
286            self.ptr = ptr_next;
287            unsafe {
288                let mut nameptr = std::ptr::null();
289                sys::Imf_DeepFrameBuffer_ConstIterator_name(
290                    &ptr_curr.0,
291                    &mut nameptr,
292                )
293                .into_result()
294                .unwrap();
295
296                if nameptr.is_null() {
297                    panic!(
298                        "DeepFrameBuffer::ConstIterator::name() returned NULL"
299                    );
300                }
301
302                let mut sliceptr = std::ptr::null();
303                sys::Imf_DeepFrameBuffer_ConstIterator_slice(
304                    &ptr_curr.0,
305                    &mut sliceptr,
306                )
307                .into_result()
308                .unwrap();
309
310                Some((
311                    CStr::from_ptr(nameptr)
312                        .to_str()
313                        .expect("NUL bytes in channel name"),
314                    DeepSliceRef::new(sliceptr),
315                ))
316            }
317        }
318    }
319}
320
321impl PartialEq for DeepFrameBufferConstIterator {
322    fn eq(&self, rhs: &DeepFrameBufferConstIterator) -> bool {
323        unsafe {
324            let mut result = false;
325            sys::Imf_deep_frame_buffer_const_iter_eq(
326                &mut result,
327                &self.0,
328                &rhs.0,
329            )
330            .into_result()
331            .unwrap();
332
333            result
334        }
335    }
336}
337
338#[repr(transparent)]
339pub struct DeepSlice(pub(crate) sys::Imf_DeepSlice_t);
340pub type DeepSliceRef<'a, P = DeepSlice> = Ref<'a, P>;
341pub type DeepSliceRefMut<'a, P = DeepSlice> = RefMut<'a, P>;
342
343unsafe impl OpaquePtr for DeepSlice {
344    type SysPointee = sys::Imf_DeepSlice_t;
345    type Pointee = DeepSlice;
346}
347
348pub struct DeepSliceBuilder {
349    pixel_type: PixelType,
350    data: *mut i8,
351    x_stride: usize,
352    y_stride: usize,
353    sample_stride: usize,
354    x_sampling: i32,
355    y_sampling: i32,
356    fill_value: f64,
357    x_tile_coords: bool,
358    y_tile_coords: bool,
359}
360
361impl DeepSliceBuilder {
362    pub fn build(self) -> Result<DeepSlice> {
363        let mut slice = sys::Imf_DeepSlice_t::default();
364        unsafe {
365            sys::Imf_DeepSlice_ctor(
366                &mut slice,
367                self.pixel_type.into(),
368                self.data,
369                self.x_stride as u64,
370                self.y_stride as u64,
371                self.sample_stride as u64,
372                self.x_sampling,
373                self.y_sampling,
374                self.fill_value,
375                self.x_tile_coords,
376                self.y_tile_coords,
377            )
378            .into_result()?;
379        }
380
381        Ok(DeepSlice(slice))
382    }
383
384    pub fn sample_stride(mut self, v: usize) -> Self {
385        self.sample_stride = v;
386        self
387    }
388
389    pub fn fill_value(mut self, v: f64) -> Self {
390        self.fill_value = v;
391        self
392    }
393
394    pub fn x_stride(mut self, x: usize) -> Self {
395        self.x_stride = x;
396        self
397    }
398
399    pub fn y_stride(mut self, y: usize) -> Self {
400        self.y_stride = y;
401        self
402    }
403
404    pub fn x_sampling(mut self, x: i32) -> Self {
405        self.x_sampling = x;
406        self
407    }
408
409    pub fn y_sampling(mut self, y: i32) -> Self {
410        self.y_sampling = y;
411        self
412    }
413
414    pub fn x_tile_coords(mut self, x: bool) -> Self {
415        self.x_tile_coords = x;
416        self
417    }
418
419    pub fn y_tile_coords(mut self, y: bool) -> Self {
420        self.y_tile_coords = y;
421        self
422    }
423}
424
425impl DeepSlice {
426    pub fn builder(pixel_type: PixelType, data: *mut i8) -> DeepSliceBuilder {
427        DeepSliceBuilder {
428            pixel_type,
429            data,
430            x_stride: std::mem::size_of::<*mut i8>(),
431            y_stride: 0,
432            sample_stride: 0,
433            x_sampling: 1,
434            y_sampling: 1,
435            fill_value: 0.0,
436            x_tile_coords: false,
437            y_tile_coords: false,
438        }
439    }
440
441    pub fn from_sample_ptr<S: DeepSample>(
442        data: *mut *mut S,
443        width: i32,
444    ) -> DeepSliceBuilder {
445        DeepSliceBuilder {
446            pixel_type: S::CHANNEL_TYPE,
447            data: data as *mut i8,
448            x_stride: std::mem::size_of::<*mut i8>(),
449            y_stride: std::mem::size_of::<*mut i8>() * width as usize,
450            sample_stride: std::mem::size_of::<S>(),
451            x_sampling: 1,
452            y_sampling: 1,
453            fill_value: 0.0,
454            x_tile_coords: false,
455            y_tile_coords: false,
456        }
457    }
458}
459
460impl Drop for DeepSlice {
461    fn drop(&mut self) {
462        unsafe {
463            sys::Imf_DeepSlice_dtor(&mut self.0);
464        }
465    }
466}
467
468pub trait DeepSample {
469    type Type;
470    const CHANNEL_TYPE: PixelType;
471    const STRIDE: usize = std::mem::size_of::<Self::Type>();
472}
473
474impl DeepSample for half::f16 {
475    type Type = Self;
476    const CHANNEL_TYPE: PixelType = PixelType::Half;
477}
478
479impl DeepSample for f32 {
480    type Type = Self;
481    const CHANNEL_TYPE: PixelType = PixelType::Float;
482}
483
484impl DeepSample for u32 {
485    type Type = Self;
486    const CHANNEL_TYPE: PixelType = PixelType::Uint;
487}
488
489pub struct DeepFrame {
490    pub(crate) channel_type: PixelType,
491    pub(crate) data_window: [i32; 4],
492    pub(crate) channel_name: String,
493    pub(crate) stride: usize,
494    pub(crate) ptr: *mut *mut u8,
495    pub(crate) len: usize,
496    pub(crate) byte_len: usize,
497    pub(crate) align: usize,
498    allocated: bool,
499}
500
501use std::alloc::{GlobalAlloc, Layout, System};
502impl DeepFrame {
503    pub fn new<T: DeepSample, B: Bound2<i32>>(
504        channel_name: &str,
505        data_window: B,
506    ) -> Result<DeepFrame> {
507        let data_window = *data_window.as_slice();
508        let w = data_window[2] - data_window[0] + 1;
509        let h = data_window[3] - data_window[1] + 1;
510        let len = (w * h) as usize;
511
512        let ptr = unsafe {
513            // FIXME: this needs to be MaybeUninit
514            System.alloc(Layout::array::<*mut u8>(len).unwrap()) as *mut *mut u8
515        };
516
517        Ok(DeepFrame {
518            channel_type: T::CHANNEL_TYPE,
519            data_window,
520            channel_name: channel_name.to_string(),
521            stride: T::STRIDE,
522            ptr,
523            len,
524            byte_len: len * std::mem::size_of::<*mut T>(),
525            align: std::mem::align_of::<*mut T>(),
526            allocated: false,
527        })
528    }
529
530    pub(crate) unsafe fn allocate_pixel_storage(
531        &mut self,
532        x: i32,
533        y: i32,
534        count: u32,
535    ) {
536        // if we're out of bounds, just ignore it. We expect to be called in a
537        // loop over the union of all windows
538        if x < self.data_window[0]
539            || x >= self.data_window[2]
540            || y < self.data_window[1]
541            || y >= self.data_window[3]
542        {
543            return;
544        }
545
546        let w = self.data_window[2] - self.data_window[0] + 1;
547
548        // offset index back to data window corner
549        let x = x - self.data_window[0];
550        let y = y - self.data_window[2];
551
552        // allocate the storage
553        let ptr = System
554            .alloc(Layout::array::<u8>(count as usize * self.stride).unwrap());
555
556        // set the sample pointer at the pixel location
557        let offset = (y * w + x) as isize;
558        *self.ptr.offset(offset) = ptr;
559    }
560}
561
562impl Drop for DeepFrame {
563    fn drop(&mut self) {
564        unsafe {
565            if self.allocated {
566                // TODO:
567            }
568
569            System.dealloc(
570                self.ptr as *mut u8,
571                Layout::array::<*mut u8>(self.len).unwrap(),
572            );
573        }
574    }
575}