add frame reordering

1 files changed, 74 insertions, 39 deletions

diff --git a/src/main.rs b/src/main.rs
index a6d464a..ac5a92c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,6 +10,7 @@ use nihav_core::options::*;
 use nihav_core::codecs::*;
 use nihav_core::demuxers::*;
 use nihav_core::muxers::*;
+use nihav_core::reorder::*;
 use nihav_core::scale::*;
 use nihav_core::soundcvt::*;
 use nihav_registry::detect;
@@ -71,7 +72,7 @@ struct Transcoder {
     mux_opts:       Vec<OptionArgs>,
     istr_opts:      Vec<InputStreamOptions>,
     ostr_opts:      Vec<OutputStreamOptions>,
-    decoders:       Vec<Option<(Box<NADecoderSupport>, Box<dyn NADecoder>)>>,
+    decoders:       Vec<Option<(Box<NADecoderSupport>, Box<dyn NADecoder>, Box<dyn FrameReorderer>)>>,
     encoders:       Vec<OutputMode>,
     no_video:       bool,
     no_audio:       bool,
@@ -585,6 +586,43 @@ println!("stream {} ({}) can't be handled", istr.id, istr.get_info().get_name())
     }
 }
 
+fn encode_frame(dst_id: u32, encoder: &mut Box<NAEncoder>, cvt: &mut OutputConvert, frm: NAFrameRef) -> bool {
+    let buf = frm.get_buffer();
+    let cbuf = match cvt {
+            OutputConvert::None => buf,
+            OutputConvert::Video(ref mut scaler, ref mut dbuf) => {
+                let cur_ifmt = get_scale_fmt_from_pic(&buf);
+                let last_ifmt = scaler.get_in_fmt();
+                if cur_ifmt != last_ifmt {
+                    let ofmt = scaler.get_out_fmt();
+                    let ret = NAScale::new(cur_ifmt, ofmt);
+                    if ret.is_err() {
+                        println!("error re-initialising scaler for {} -> {}", cur_ifmt, ofmt);
+                        return false;
+                    }
+                    *scaler = ret.unwrap();
+                }
+                let ret = scaler.convert(&buf, dbuf);
+                if ret.is_err() {
+                    println!("error converting frame for encoding");
+                    return false;
+                }
+                dbuf.clone()
+            },
+            OutputConvert::Audio(ref dinfo, ref dchmap) => {
+                let ret = convert_audio_frame(&buf, dinfo, dchmap);
+                if ret.is_err() {
+                    println!("error converting audio for stream {}", dst_id);
+                    return false;
+                }
+                ret.unwrap()
+            },
+        };
+    let cfrm = NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cbuf);
+    encoder.encode(&cfrm).unwrap();
+    true
+}
+
 macro_rules! next_arg {
     ($args: expr, $arg_idx: expr) => {
         if $arg_idx + 1 >= $args.len() {
@@ -853,7 +891,15 @@ println!("stream {} - {} {}", i, s, info.get_name());
                 return;
             }
             transcoder.apply_decoder_options(dec.as_mut(), str_id);
-            transcoder.decoders.push(Some((dsupp, dec)));
+            let desc = register::get_codec_description(info.get_name());
+            let has_b = if let Some(desc) = desc {
+                    desc.has_reorder()
+                } else {
+                    println!("No codec description found, using B-frame reorderer.");
+                    true
+                };
+            let reord: Box<dyn FrameReorderer> = if has_b { Box::new(IPBReorderer::new()) } else { Box::new(NoReorderer::new()) };
+            transcoder.decoders.push(Some((dsupp, dec, reord)));
         } else {
             println!("No decoder for stream {} ({}) is found", str_id, info.get_name());
             transcoder.decoders.push(None);
@@ -933,7 +979,7 @@ println!("stream {} - {} {}", i, s, info.get_name());
                 }
             },
             OutputMode::Encode(dst_id, ref mut encoder, ref mut cvt) => {
-                if let Some((ref mut dsupp, ref mut decoder)) = transcoder.decoders[src_id] {
+                if let Some((ref mut dsupp, ref mut decoder, ref mut reorderer)) = transcoder.decoders[src_id] {
                     let ret = decoder.decode(dsupp, &pkt);
                     if let (true, Err(DecoderError::MissingReference)) = (transcoder.start != NATimePoint::None, &ret) {
                         continue;
@@ -943,42 +989,15 @@ println!("stream {} - {} {}", i, s, info.get_name());
                         break;
                     }
                     let frm = ret.unwrap();
-                    let buf = frm.get_buffer();
-                    let cbuf = match cvt {
-                            OutputConvert::None => buf,
-                            OutputConvert::Video(ref mut scaler, ref mut dbuf) => {
-                                let cur_ifmt = get_scale_fmt_from_pic(&buf);
-                                let last_ifmt = scaler.get_in_fmt();
-                                if cur_ifmt != last_ifmt {
-                                    let ofmt = scaler.get_out_fmt();
-                                    let ret = NAScale::new(cur_ifmt, ofmt);
-                                    if ret.is_err() {
-                                        println!("error re-initialising scaler for {} -> {}", cur_ifmt, ofmt);
-                                        break;
-                                    }
-                                    *scaler = ret.unwrap();
-                                }
-                                let ret = scaler.convert(&buf, dbuf);
-                                if ret.is_err() {
-                                    println!("error converting frame for encoding");
-                                    break;
-                                }
-                                dbuf.clone()
-                            },
-                            OutputConvert::Audio(ref dinfo, ref dchmap) => {
-                                let ret = convert_audio_frame(&buf, dinfo, dchmap);
-                                if ret.is_err() {
-                                    println!("error converting audio for stream {}", dst_id);
-                                    break;
-                                }
-                                ret.unwrap()
-                            },
-                        };
-                    let cfrm = NAFrame::new(frm.get_time_information(), frm.frame_type, frm.key, frm.get_info(), cbuf);
-                    encoder.encode(&cfrm).unwrap();
-                    while let Ok(Some(pkt)) = encoder.get_packet() {
-                        if transcoder.end != NATimePoint::None && !pkt.ts.less_than(transcoder.end) { break 'main_loop; }
-                        mux.mux_frame(pkt).unwrap();
+                    reorderer.add_frame(frm);
+                    while let Some(frm) = reorderer.get_frame() {
+                        if !encode_frame(dst_id, encoder, cvt, frm) {
+                            break;
+                        }
+                        while let Ok(Some(pkt)) = encoder.get_packet() {
+                            if transcoder.end != NATimePoint::None && !pkt.ts.less_than(transcoder.end) { break 'main_loop; }
+                            mux.mux_frame(pkt).unwrap();
+                        }
                     }
                 } else {
                     println!("no decoder for stream {}", src_id);
@@ -987,6 +1006,22 @@ println!("stream {} - {} {}", i, s, info.get_name());
             },
         };
     }
+    'reord_flush_loop: for str in dmx.get_streams() {
+        let src_id = str.get_num();
+        if let OutputMode::Encode(dst_id, ref mut encoder, ref mut cvt) = transcoder.encoders[src_id] {
+            if let Some((_, _, ref mut reorderer)) = transcoder.decoders[src_id] {
+                while let Some(frm) = reorderer.get_last_frames() {
+                    if !encode_frame(dst_id, encoder, cvt, frm) {
+                        break;
+                    }
+                    while let Ok(Some(pkt)) = encoder.get_packet() {
+                        if transcoder.end != NATimePoint::None && !pkt.ts.less_than(transcoder.end) { break 'reord_flush_loop; }
+                        mux.mux_frame(pkt).unwrap();
+                    }
+                }
+            }
+        }
+    }
     'flush_loop: for enc in transcoder.encoders.iter_mut() {
         match enc {
             OutputMode::Encode(str_id, ref mut encoder, _) => {