switch to refcounted buffers

6 files changed, 173 insertions, 40 deletions

diff --git a/nihav-core/src/codecs/blockdsp.rs b/nihav-core/src/codecs/blockdsp.rs
index 33d682a..ccd44cb 100644
--- a/nihav-core/src/codecs/blockdsp.rs
+++ b/nihav-core/src/codecs/blockdsp.rs
@@ -8,7 +8,7 @@ pub fn put_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[
     let mut idxu = buf.get_offset(1) + xpos *  8 + ypos *  8 * strideu;
     let mut idxv = buf.get_offset(2) + xpos *  8 + ypos *  8 * stridev;
 
-    let mut data = buf.get_data_mut();
+    let data = buf.get_data_mut().unwrap();
     let framebuf: &mut [u8] = data.as_mut_slice();
 
     for j in 0..8 {
@@ -62,7 +62,7 @@ pub fn add_blocks(buf: &mut NAVideoBuffer<u8>, xpos: usize, ypos: usize, blk: &[
     let mut idxu = buf.get_offset(1) + xpos *  8 + ypos *  8 * strideu;
     let mut idxv = buf.get_offset(2) + xpos *  8 + ypos *  8 * stridev;
 
-    let mut data = buf.get_data_mut();
+    let data = buf.get_data_mut().unwrap();
     let framebuf: &mut [u8] = data.as_mut_slice();
 
     for j in 0..8 {
@@ -149,7 +149,7 @@ pub fn copy_blocks(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>,
         for comp in 0..3 {
             let dstride = dst.get_stride(comp);
             let doff    = dst.get_offset(comp);
-            let mut ddta    = dst.get_data_mut();
+            let ddta    = dst.get_data_mut().unwrap();
             let dbuf: &mut [u8] = ddta.as_mut_slice();
             let x   = if comp > 0 { dx/2 } else { dx };
             let y   = if comp > 0 { dy/2 } else { dy };
@@ -171,7 +171,7 @@ pub fn copy_blocks(dst: &mut NAVideoBuffer<u8>, src: &NAVideoBuffer<u8>,
             let sbuf: &[u8] = sdta.as_slice();
             let dstride = dst.get_stride(comp);
             let doff    = dst.get_offset(comp);
-            let mut ddta    = dst.get_data_mut();
+            let ddta    = dst.get_data_mut().unwrap();
             let dbuf: &mut [u8] = ddta.as_mut_slice();
             let x   = if comp > 0 { dx/2 } else { dx };
             let y   = if comp > 0 { dy/2 } else { dy };
diff --git a/nihav-core/src/codecs/h263/code.rs b/nihav-core/src/codecs/h263/code.rs
index dd1279e..c18a9ea 100644
--- a/nihav-core/src/codecs/h263/code.rs
+++ b/nihav-core/src/codecs/h263/code.rs
@@ -347,7 +347,7 @@ impl H263BlockDSP {
 
 fn deblock_hor(buf: &mut NAVideoBuffer<u8>, comp: usize, q: u8, off: usize) {
     let stride = buf.get_stride(comp);
-    let mut dptr = buf.get_data_mut();
+    let dptr = buf.get_data_mut().unwrap();
     let buf = dptr.as_mut_slice();
     for x in 0..8 {
         let a = buf[off - 2 * stride + x] as i16;
@@ -377,7 +377,7 @@ fn deblock_hor(buf: &mut NAVideoBuffer<u8>, comp: usize, q: u8, off: usize) {
 
 fn deblock_ver(buf: &mut NAVideoBuffer<u8>, comp: usize, q: u8, off: usize) {
     let stride = buf.get_stride(comp);
-    let mut dptr = buf.get_data_mut();
+    let dptr = buf.get_data_mut().unwrap();
     let buf = dptr.as_mut_slice();
     for y in 0..8 {
         let a = buf[off - 2 + y * stride] as i16;
diff --git a/nihav-core/src/demuxers/mod.rs b/nihav-core/src/demuxers/mod.rs
index 201e163..48529d4 100644
--- a/nihav-core/src/demuxers/mod.rs
+++ b/nihav-core/src/demuxers/mod.rs
@@ -39,7 +39,7 @@ impl<'a> NAPacketReader for ByteReader<'a> {
     }
     fn fill_packet(&mut self, pkt: &mut NAPacket) -> DemuxerResult<()> {
         let mut refbuf = pkt.get_buffer();
-        let buf = Rc::make_mut(&mut refbuf);
+        let buf = refbuf.as_mut().unwrap();
         let res = self.read_buf(buf.as_mut_slice());
         if let Err(_) = res { return Err(DemuxerError::IOError); }
         Ok(())
@@ -226,4 +226,4 @@ impl RegisteredDemuxers {
         }
         None
     }
-}
\ No newline at end of file
+}
diff --git a/nihav-core/src/frame.rs b/nihav-core/src/frame.rs
index 253842a..ea9f203 100644
--- a/nihav-core/src/frame.rs
+++ b/nihav-core/src/frame.rs
@@ -4,6 +4,7 @@ use std::fmt;
 pub use std::rc::Rc;
 pub use std::cell::*;
 pub use crate::formats::*;
+pub use crate::refs::*;
 
 #[allow(dead_code)]
 #[derive(Clone,Copy,PartialEq)]
@@ -102,12 +103,10 @@ impl fmt::Display for NACodecTypeInfo {
     }
 }
 
-pub type NABufferRefT<T> = Rc<RefCell<Vec<T>>>;
-
 #[derive(Clone)]
 pub struct NAVideoBuffer<T> {
     info:    NAVideoInfo,
-    data:    NABufferRefT<T>,
+    data:    NABufferRef<Vec<T>>,
     offs:    Vec<usize>,
     strides: Vec<usize>,
 }
@@ -118,16 +117,16 @@ impl<T: Clone> NAVideoBuffer<T> {
         else { self.offs[idx] }
     }
     pub fn get_info(&self) -> NAVideoInfo { self.info }
-    pub fn get_data(&self) -> Ref<Vec<T>> { self.data.borrow() }
-    pub fn get_data_mut(&mut self) -> RefMut<Vec<T>> { self.data.borrow_mut() }
+    pub fn get_data(&self) -> &Vec<T> { self.data.as_ref() }
+    pub fn get_data_mut(&mut self) -> Option<&mut Vec<T>> { self.data.as_mut() }
     pub fn copy_buffer(&mut self) -> Self {
-        let mut data: Vec<T> = Vec::with_capacity(self.data.borrow().len());
-        data.clone_from(self.data.borrow().as_ref());
+        let mut data: Vec<T> = Vec::with_capacity(self.data.len());
+        data.clone_from(self.data.as_ref());
         let mut offs: Vec<usize> = Vec::with_capacity(self.offs.len());
         offs.clone_from(&self.offs);
         let mut strides: Vec<usize> = Vec::with_capacity(self.strides.len());
         strides.clone_from(&self.strides);
-        NAVideoBuffer { info: self.info, data: Rc::new(RefCell::new(data)), offs: offs, strides: strides }
+        NAVideoBuffer { info: self.info, data: NABufferRef::new(data), offs: offs, strides: strides }
     }
     pub fn get_stride(&self, idx: usize) -> usize {
         if idx >= self.strides.len() { return 0; }
@@ -141,7 +140,7 @@ impl<T: Clone> NAVideoBuffer<T> {
 #[derive(Clone)]
 pub struct NAAudioBuffer<T> {
     info:   NAAudioInfo,
-    data:   NABufferRefT<T>,
+    data:   NABufferRef<Vec<T>>,
     offs:   Vec<usize>,
     chmap:  NAChannelMap,
     len:    usize,
@@ -154,21 +153,21 @@ impl<T: Clone> NAAudioBuffer<T> {
     }
     pub fn get_info(&self) -> NAAudioInfo { self.info }
     pub fn get_chmap(&self) -> NAChannelMap { self.chmap.clone() }
-    pub fn get_data(&self) -> Ref<Vec<T>> { self.data.borrow() }
-    pub fn get_data_mut(&mut self) -> RefMut<Vec<T>> { self.data.borrow_mut() }
+    pub fn get_data(&self) -> &Vec<T> { self.data.as_ref() }
+    pub fn get_data_mut(&mut self) -> Option<&mut Vec<T>> { self.data.as_mut() }
     pub fn copy_buffer(&mut self) -> Self {
-        let mut data: Vec<T> = Vec::with_capacity(self.data.borrow().len());
-        data.clone_from(self.data.borrow().as_ref());
+        let mut data: Vec<T> = Vec::with_capacity(self.data.len());
+        data.clone_from(self.data.as_ref());
         let mut offs: Vec<usize> = Vec::with_capacity(self.offs.len());
         offs.clone_from(&self.offs);
-        NAAudioBuffer { info: self.info, data: Rc::new(RefCell::new(data)), offs: offs, chmap: self.get_chmap(), len: self.len }
+        NAAudioBuffer { info: self.info, data: NABufferRef::new(data), offs: offs, chmap: self.get_chmap(), len: self.len }
     }
     pub fn get_length(&self) -> usize { self.len }
 }
 
 impl NAAudioBuffer<u8> {
-    pub fn new_from_buf(info: NAAudioInfo, data: NABufferRefT<u8>, chmap: NAChannelMap) -> Self {
-        let len = data.borrow().len();
+    pub fn new_from_buf(info: NAAudioInfo, data: NABufferRef<Vec<u8>>, chmap: NAChannelMap) -> Self {
+        let len = data.len();
         NAAudioBuffer { info: info, data: data, chmap: chmap, offs: Vec::new(), len: len }
     }
 }
@@ -184,7 +183,7 @@ pub enum NABufferType {
     AudioI32   (NAAudioBuffer<i32>),
     AudioF32   (NAAudioBuffer<f32>),
     AudioPacked(NAAudioBuffer<u8>),
-    Data       (NABufferRefT<u8>),
+    Data       (NABufferRef<Vec<u8>>),
     None,
 }
 
@@ -309,7 +308,7 @@ pub fn alloc_video_buffer(vinfo: NAVideoInfo, align: u8) -> Result<NABufferType,
         strides.push(stride);
         let mut data: Vec<u8> = Vec::with_capacity(new_size.unwrap());
         data.resize(new_size.unwrap(), 0);
-        let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+        let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
         Ok(NABufferType::Video(buf))
     } else if !all_packed {
         for i in 0..fmt.get_num_comp() {
@@ -334,17 +333,17 @@ pub fn alloc_video_buffer(vinfo: NAVideoInfo, align: u8) -> Result<NABufferType,
         if max_depth <= 8 {
             let mut data: Vec<u8> = Vec::with_capacity(new_size);
             data.resize(new_size, 0);
-            let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+            let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
             Ok(NABufferType::Video(buf))
         } else if max_depth <= 16 {
             let mut data: Vec<u16> = Vec::with_capacity(new_size);
             data.resize(new_size, 0);
-            let buf: NAVideoBuffer<u16> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+            let buf: NAVideoBuffer<u16> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
             Ok(NABufferType::Video16(buf))
         } else {
             let mut data: Vec<u32> = Vec::with_capacity(new_size);
             data.resize(new_size, 0);
-            let buf: NAVideoBuffer<u32> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+            let buf: NAVideoBuffer<u32> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
             Ok(NABufferType::Video32(buf))
         }
     } else if all_bytealigned || unfit_elem_size {
@@ -357,7 +356,7 @@ pub fn alloc_video_buffer(vinfo: NAVideoInfo, align: u8) -> Result<NABufferType,
         let mut data: Vec<u8> = Vec::with_capacity(new_size);
         data.resize(new_size, 0);
         strides.push(line_sz.unwrap());
-        let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+        let buf: NAVideoBuffer<u8> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
         Ok(NABufferType::VideoPacked(buf))
     } else {
         let elem_sz = fmt.get_elem_size();
@@ -369,14 +368,14 @@ pub fn alloc_video_buffer(vinfo: NAVideoInfo, align: u8) -> Result<NABufferType,
                     let mut data: Vec<u16> = Vec::with_capacity(new_size);
                     data.resize(new_size, 0);
                     strides.push(width);
-                    let buf: NAVideoBuffer<u16> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+                    let buf: NAVideoBuffer<u16> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
                     Ok(NABufferType::Video16(buf))
                 },
             4 => {
                     let mut data: Vec<u32> = Vec::with_capacity(new_size);
                     data.resize(new_size, 0);
                     strides.push(width);
-                    let buf: NAVideoBuffer<u32> = NAVideoBuffer { data: Rc::new(RefCell::new(data)), info: vinfo, offs: offs, strides: strides };
+                    let buf: NAVideoBuffer<u32> = NAVideoBuffer { data: NABufferRef::new(data), info: vinfo, offs: offs, strides: strides };
                     Ok(NABufferType::Video32(buf))
                 },
             _ => unreachable!(),
@@ -397,7 +396,7 @@ pub fn alloc_audio_buffer(ainfo: NAAudioInfo, nsamples: usize, chmap: NAChannelM
             if ainfo.format.get_bits() == 32 {
                 let mut data: Vec<f32> = Vec::with_capacity(length);
                 data.resize(length, 0.0);
-                let buf: NAAudioBuffer<f32> = NAAudioBuffer { data: Rc::new(RefCell::new(data)), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
+                let buf: NAAudioBuffer<f32> = NAAudioBuffer { data: NABufferRef::new(data), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
                 Ok(NABufferType::AudioF32(buf))
             } else {
                 Err(AllocatorError::TooLargeDimensions)
@@ -406,12 +405,12 @@ pub fn alloc_audio_buffer(ainfo: NAAudioInfo, nsamples: usize, chmap: NAChannelM
             if ainfo.format.get_bits() == 8 && !ainfo.format.is_signed() {
                 let mut data: Vec<u8> = Vec::with_capacity(length);
                 data.resize(length, 0);
-                let buf: NAAudioBuffer<u8> = NAAudioBuffer { data: Rc::new(RefCell::new(data)), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
+                let buf: NAAudioBuffer<u8> = NAAudioBuffer { data: NABufferRef::new(data), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
                 Ok(NABufferType::AudioU8(buf))
             } else if ainfo.format.get_bits() == 16 && ainfo.format.is_signed() {
                 let mut data: Vec<i16> = Vec::with_capacity(length);
                 data.resize(length, 0);
-                let buf: NAAudioBuffer<i16> = NAAudioBuffer { data: Rc::new(RefCell::new(data)), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
+                let buf: NAAudioBuffer<i16> = NAAudioBuffer { data: NABufferRef::new(data), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
                 Ok(NABufferType::AudioI16(buf))
             } else {
                 Err(AllocatorError::TooLargeDimensions)
@@ -423,7 +422,7 @@ pub fn alloc_audio_buffer(ainfo: NAAudioInfo, nsamples: usize, chmap: NAChannelM
         let length = ainfo.format.get_audio_size(len.unwrap() as u64);
         let mut data: Vec<u8> = Vec::with_capacity(length);
         data.resize(length, 0);
-        let buf: NAAudioBuffer<u8> = NAAudioBuffer { data: Rc::new(RefCell::new(data)), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
+        let buf: NAAudioBuffer<u8> = NAAudioBuffer { data: NABufferRef::new(data), info: ainfo, offs: offs, chmap: chmap, len: nsamples };
         Ok(NABufferType::AudioPacked(buf))
     }
 }
@@ -431,7 +430,7 @@ pub fn alloc_audio_buffer(ainfo: NAAudioInfo, nsamples: usize, chmap: NAChannelM
 pub fn alloc_data_buffer(size: usize) -> Result<NABufferType, AllocatorError> {
     let mut data: Vec<u8> = Vec::with_capacity(size);
     data.resize(size, 0);
-    let buf: NABufferRefT<u8> = Rc::new(RefCell::new(data));
+    let buf: NABufferRef<Vec<u8>> = NABufferRef::new(data);
     Ok(NABufferType::Data(buf))
 }
 
@@ -439,6 +438,52 @@ pub fn copy_buffer(buf: NABufferType) -> NABufferType {
     buf.clone()
 }
 
+pub struct NABufferPool {
+    pool:       Vec<NABufferRef<NABufferType>>,
+    max_len:    usize,
+}
+
+impl NABufferPool {
+    pub fn new(max_len: usize) -> Self {
+        Self {
+            pool:       Vec::with_capacity(max_len),
+            max_len,
+        }
+    }
+    pub fn prealloc_video(&mut self, vinfo: NAVideoInfo, align: u8) -> Result<(), AllocatorError> {
+        let nbufs = self.max_len - self.pool.len();
+        for _ in 0..nbufs {
+            let buf = alloc_video_buffer(vinfo.clone(), align)?;
+            self.pool.push(NABufferRef::new(buf));
+        }
+        Ok(())
+    }
+    pub fn prealloc_audio(&mut self, ainfo: NAAudioInfo, nsamples: usize, chmap: NAChannelMap) -> Result<(), AllocatorError> {
+        let nbufs = self.max_len - self.pool.len();
+        for _ in 0..nbufs {
+            let buf = alloc_audio_buffer(ainfo.clone(), nsamples, chmap.clone())?;
+            self.pool.push(NABufferRef::new(buf));
+        }
+        Ok(())
+    }
+    pub fn add(&mut self, buf: NABufferType) -> bool {
+        if self.pool.len() < self.max_len {
+            self.pool.push(NABufferRef::new(buf));
+            true
+        } else {
+            false
+        }
+    }
+    pub fn get_free(&mut self) -> Option<NABufferRef<NABufferType>> {
+        for e in self.pool.iter() {
+            if e.get_num_refs() == 1 {
+                return Some(e.clone());
+            }
+        }
+        None
+    }
+}
+
 #[allow(dead_code)]
 #[derive(Clone)]
 pub struct NACodecInfo {
@@ -690,7 +735,7 @@ impl fmt::Display for NAStream {
 pub struct NAPacket {
     stream:         Rc<NAStream>,
     ts:             NATimeInfo,
-    buffer:         Rc<Vec<u8>>,
+    buffer:         NABufferRef<Vec<u8>>,
     keyframe:       bool,
 //    options:        HashMap<String, NAValue<'a>>,
 }
@@ -699,7 +744,7 @@ impl NAPacket {
     pub fn new(str: Rc<NAStream>, ts: NATimeInfo, kf: bool, vec: Vec<u8>) -> Self {
 //        let mut vec: Vec<u8> = Vec::new();
 //        vec.resize(size, 0);
-        NAPacket { stream: str, ts: ts, keyframe: kf, buffer: Rc::new(vec) }
+        NAPacket { stream: str, ts: ts, keyframe: kf, buffer: NABufferRef::new(vec) }
     }
     pub fn get_stream(&self) -> Rc<NAStream> { self.stream.clone() }
     pub fn get_time_information(&self) -> NATimeInfo { self.ts }
@@ -707,7 +752,7 @@ impl NAPacket {
     pub fn get_dts(&self) -> Option<u64> { self.ts.get_dts() }
     pub fn get_duration(&self) -> Option<u64> { self.ts.get_duration() }
     pub fn is_keyframe(&self) -> bool { self.keyframe }
-    pub fn get_buffer(&self) -> Rc<Vec<u8>> { self.buffer.clone() }
+    pub fn get_buffer(&self) -> NABufferRef<Vec<u8>> { self.buffer.clone() }
 }
 
 impl Drop for NAPacket {
diff --git a/nihav-core/src/lib.rs b/nihav-core/src/lib.rs
index 05b7ad9..615a0cf 100644
--- a/nihav-core/src/lib.rs
+++ b/nihav-core/src/lib.rs
@@ -7,6 +7,7 @@ pub mod demuxers;
 pub mod formats;
 pub mod frame;
 pub mod io;
+pub mod refs;
 pub mod register;
 pub mod detect;
 
diff --git a/nihav-core/src/refs.rs b/nihav-core/src/refs.rs
new file mode 100644
index 0000000..f2b0577
--- /dev/null
+++ b/nihav-core/src/refs.rs
@@ -0,0 +1,87 @@
+use std::ops::{Deref, DerefMut};
+use std::sync::atomic::*;
+
+struct NABufferData<T> {
+    data:       T,
+    refs:       AtomicUsize,
+}
+
+impl<T> NABufferData<T> {
+    fn new(data: T) -> Self {
+        Self {
+            data:       data,
+            refs:       AtomicUsize::new(1),
+        }
+    }
+    fn inc_refs(obj: &mut Self) {
+        obj.refs.fetch_add(1, Ordering::SeqCst);
+    }
+    fn dec_refs(obj: &mut Self) {
+        if obj.refs.fetch_sub(1, Ordering::SeqCst) == 0 {
+            std::mem::forget(obj);
+        }
+    }
+    fn get_num_refs(obj: &Self) -> usize {
+        obj.refs.load(Ordering::Relaxed)
+    }
+    fn get_read_ptr(obj: &Self) -> &T {
+        &obj.data
+    }
+    fn get_write_ptr(obj: &mut Self) -> Option<&mut T> {
+        Some(&mut obj.data)
+    }
+}
+
+pub struct NABufferRef<T> {
+    ptr: *mut NABufferData<T>,
+}
+
+impl<T> NABufferRef<T> {
+    pub fn new(val: T) -> Self {
+        let bdata = NABufferData::new(val);
+        let nbox: Box<_> = Box::new(bdata);
+        Self { ptr: Box::into_raw(nbox) }
+    }
+    pub fn get_num_refs(&self) -> usize {
+        unsafe {
+            NABufferData::get_num_refs(self.ptr.as_mut().unwrap())
+        }
+    }
+    pub fn as_ref(&self) -> &T {
+        unsafe {
+            NABufferData::get_read_ptr(self.ptr.as_mut().unwrap())
+        }
+    }
+    pub fn as_mut(&mut self) -> Option<&mut T> {
+        unsafe {
+            NABufferData::get_write_ptr(self.ptr.as_mut().unwrap())
+        }
+    }
+}
+
+impl<T> Deref for NABufferRef<T> {
+    type Target = T;
+    fn deref(&self) -> &T { self.as_ref() }
+}
+
+impl<T> DerefMut for NABufferRef<T> {
+    fn deref_mut(&mut self) -> &mut T { self.as_mut().unwrap() }
+}
+
+impl<T> Clone for NABufferRef<T> {
+    fn clone(&self) -> Self {
+        unsafe {
+            NABufferData::inc_refs(self.ptr.as_mut().unwrap());
+        }
+        Self { ptr: self.ptr }
+    }
+}
+
+impl<T> Drop for NABufferRef<T> {
+    fn drop(&mut self) {
+        unsafe {
+            NABufferData::dec_refs(self.ptr.as_mut().unwrap());
+        }
+    }
+}
+