allow encoder to work with multiple inputs
[nihav-encoder.git] / src / demux.rs
CommitLineData
91a15e39
KS
1use std::io::SeekFrom;
2use nihav_core::codecs::*;
3use nihav_core::demuxers::*;
4use nihav_core::muxers::*;
5use nihav_registry::detect;
6use nihav_core::io::byteio::ByteReader;
7use nihav_allstuff::*;
8use crate::null::*;
90683bc3 9use crate::imgseq::*;
4a835a2a 10use crate::{SBBox, SelfBorrow};
91a15e39
KS
11
12pub struct FullRegister {
13 pub dmx_reg: RegisteredDemuxers,
14 pub rdmx_reg: RegisteredRawDemuxers,
15 pub pkt_reg: RegisteredPacketisers,
16 pub dec_reg: RegisteredDecoders,
17 pub enc_reg: RegisteredEncoders,
18 pub mux_reg: RegisteredMuxers,
19}
20
21impl FullRegister {
22 pub fn new() -> Self {
23 let mut dmx_reg = RegisteredDemuxers::new();
24 nihav_register_all_demuxers(&mut dmx_reg);
25 let mut rdmx_reg = RegisteredRawDemuxers::new();
26 nihav_register_all_raw_demuxers(&mut rdmx_reg);
27 let mut dec_reg = RegisteredDecoders::new();
28 nihav_register_all_decoders(&mut dec_reg);
29 let mut pkt_reg = RegisteredPacketisers::new();
30 nihav_register_all_packetisers(&mut pkt_reg);
31 let mut enc_reg = RegisteredEncoders::new();
32 nihav_register_all_encoders(&mut enc_reg);
33 enc_reg.add_encoder(NULL_ENCODER);
34 let mut mux_reg = RegisteredMuxers::new();
35 nihav_register_all_muxers(&mut mux_reg);
36 mux_reg.add_muxer(NULL_MUXER);
37 Self { dmx_reg, rdmx_reg, pkt_reg, dec_reg, enc_reg, mux_reg }
38 }
39}
40
41pub struct RawStreamCtx<'a> {
42 stream: NAStreamRef,
43 sm: StreamManager,
44 packetiser: Box<dyn NAPacketiser + Send>,
45 br: &'a mut ByteReader<'a>,
206dd66d 46 pts: u64,
91a15e39
KS
47}
48
49impl<'a> RawStreamCtx<'a> {
50 fn new(stream: NAStreamRef, packetiser: Box<dyn NAPacketiser + Send>, br: &'a mut ByteReader<'a>) -> Self {
51 let mut sm = StreamManager::new();
52 sm.add_stream_ref(stream.clone());
206dd66d 53 Self { stream, sm, packetiser, br, pts: 0 }
91a15e39
KS
54 }
55}
56
4a835a2a
KS
57type ReaderBox<'a> = SBBox<Box<dyn ByteIO>, ByteReader<'a>>;
58
91a15e39
KS
59pub enum DemuxerObject<'a> {
60 None,
4a835a2a
KS
61 Normal(SBBox<ReaderBox<'a>, Demuxer<'a>>),
62 Raw(SBBox<ReaderBox<'a>, RawDemuxer<'a>>, Vec<Option<Box<dyn NAPacketiser + Send>>>, bool),
63 RawStream(SBBox<ReaderBox<'a>, RawStreamCtx<'a>>),
90683bc3 64 ImageSequence(ImgSeqDemuxer),
91a15e39
KS
65}
66
67impl<'a> DemuxerObject<'a> {
4a835a2a
KS
68 pub fn create(mut brb: ReaderBox<'a>, reg: &FullRegister, name: &str, ifmt: &Option<String>, is_raw: bool) -> DemuxerObject<'a> {
69 let br = brb.get_object_mut();
91a15e39
KS
70 if !is_raw {
71 let dmx_name = if let Some(ref dname) = ifmt {
72 dname.as_str()
73 } else if let Some((dname, score)) = detect::detect_format(name, br) {
74 println!("detected {} with score {:?}", dname, score);
75 dname
76 } else {
77 ""
78 };
76ea2e41 79 if !dmx_name.is_empty() {
91a15e39
KS
80 println!("trying demuxer {} on {}", dmx_name, name);
81 if let Some(dmx_fact) = reg.dmx_reg.find_demuxer(dmx_name) {
82 br.seek(SeekFrom::Start(0)).unwrap();
4a835a2a
KS
83
84 let dmx = SelfBorrow::new(brb, |br_| {
85 unsafe {
86 create_demuxer(dmx_fact, br_.as_mut().unwrap().get_object_mut()).unwrap()
87 }
88 });
91a15e39
KS
89 return DemuxerObject::Normal(dmx);
90 }
91 }
92 if ifmt.is_some() {
93 return DemuxerObject::None;
94 }
76ea2e41 95 if !dmx_name.is_empty() {
91a15e39
KS
96 println!("trying raw demuxer {} on {}", dmx_name, name);
97 if let Some(rdmx_fact) = reg.rdmx_reg.find_demuxer(dmx_name) {
98 br.seek(SeekFrom::Start(0)).unwrap();
4a835a2a
KS
99 let dmx = SelfBorrow::new(brb, |br_| {
100 unsafe {
101 create_raw_demuxer(rdmx_fact, br_.as_mut().unwrap().get_object_mut()).unwrap()
102 }
103 });
91a15e39 104 let mut pkts = Vec::new();
4a835a2a 105 for stream in dmx.get_object().get_streams() {
91a15e39
KS
106 if let Some(pcreate) = reg.pkt_reg.find_packetiser(stream.get_info().get_name()) {
107 let packetiser = (pcreate)();
108 pkts.push(Some(packetiser));
109 } else {
110 pkts.push(None);
111 }
112 }
113 return DemuxerObject::Raw(dmx, pkts, false);
114 }
115 }
116 for rdmx in reg.rdmx_reg.iter() {
117 if rdmx.check_format(br) {
118 println!("detected {} as {}", name, rdmx.get_name());
119 br.seek(SeekFrom::Start(0)).unwrap();
4a835a2a
KS
120 let dmx = SelfBorrow::new(brb, |br_| {
121 unsafe {
122 create_raw_demuxer(*rdmx, br_.as_mut().unwrap().get_object_mut()).unwrap()
123 }
124 });
91a15e39 125 let mut pkts = Vec::new();
4a835a2a 126 for stream in dmx.get_object().get_streams() {
91a15e39
KS
127 if let Some(pcreate) = reg.pkt_reg.find_packetiser(stream.get_info().get_name()) {
128 let packetiser = (pcreate)();
129 pkts.push(Some(packetiser));
130 } else {
131 pkts.push(None);
132 }
133 }
134 return DemuxerObject::Raw(dmx, pkts, false);
135 }
136 }
137 }
138 br.seek(SeekFrom::Start(0)).unwrap();
139 let mut buf = vec![0; 1048576];
140 let size = br.peek_buf(&mut buf).unwrap();
141 let mut pname = "";
142
143 for pinfo in reg.pkt_reg.iter() {
144 let mut packetiser = (pinfo.get_packetiser)();
145 packetiser.add_data(&buf[..size]);
146 if packetiser.parse_stream(0).is_ok() {
147 pname = pinfo.name;
148 break;
149 }
150 }
76ea2e41 151 if !pname.is_empty() {
91a15e39
KS
152 println!("found raw stream of type {} for {}", pname, name);
153 let pcreate = reg.pkt_reg.find_packetiser(pname).unwrap();
4a835a2a
KS
154 let rctx = SelfBorrow::new(brb, |br_| {
155 unsafe {
156 let mut packetiser = (pcreate)();
157 packetiser.add_data(&buf[..size]);
158 let stream = packetiser.parse_stream(0).unwrap();
159 packetiser.reset();
160 RawStreamCtx::new(stream, packetiser, br_.as_mut().unwrap().get_object_mut())
161 }
162 });
163 DemuxerObject::RawStream(rctx)
91a15e39
KS
164 } else {
165 DemuxerObject::None
166 }
167 }
90683bc3
KS
168 pub fn create_imgseq(isd: ImgSeqDemuxer) -> Self {
169 DemuxerObject::ImageSequence(isd)
170 }
91a15e39 171 pub fn is_none(&self) -> bool {
76ea2e41 172 matches!(*self, DemuxerObject::None)
91a15e39 173 }
5ec8115f 174 pub fn get_duration(&self) -> u64 {
91a15e39 175 match *self {
4a835a2a
KS
176 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_duration(),
177 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_duration(),
4dd9047c 178 DemuxerObject::RawStream(ref ctx) => {
4a835a2a 179 NATimeInfo::ts_to_time(ctx.get_object().stream.duration, 1000, ctx.get_object().stream.tb_num, ctx.get_object().stream.tb_den)
4dd9047c 180 },
91a15e39
KS
181 _ => 0,
182 }
5ec8115f 183 }
91a15e39
KS
184 pub fn get_num_streams(&self) -> usize {
185 match *self {
186 DemuxerObject::None => 0,
4a835a2a
KS
187 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_num_streams(),
188 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_num_streams(),
91a15e39 189 DemuxerObject::RawStream(_) => 1,
90683bc3 190 DemuxerObject::ImageSequence(_) => 1,
91a15e39
KS
191 }
192 }
193 pub fn get_stream(&self, idx: usize) -> Option<NAStreamRef> {
194 match *self {
4a835a2a
KS
195 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_stream(idx),
196 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_stream(idx),
197 DemuxerObject::RawStream(ref ctx) if idx == 0 => Some(ctx.get_object().stream.clone()),
90683bc3 198 DemuxerObject::ImageSequence(ref ctx) if idx == 0 => Some(ctx.stream.clone()),
91a15e39
KS
199 _ => None,
200 }
201 }
4a835a2a 202 /*pub fn get_streams(&self) -> StreamIter {
91a15e39 203 match *self {
4a835a2a
KS
204 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_streams(),
205 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_streams(),
206 DemuxerObject::RawStream(ref ctx) => ctx.get_object().sm.iter(),
90683bc3 207 DemuxerObject::ImageSequence(ref ctx) => ctx.sm.iter(),
91a15e39
KS
208 _ => unreachable!(),
209 }
4a835a2a 210 }*/
91a15e39
KS
211 pub fn get_stream_manager(&self) -> &StreamManager {
212 match *self {
4a835a2a
KS
213 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_stream_manager(),
214 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_stream_manager(),
215 DemuxerObject::RawStream(ref ctx) => &ctx.get_object().sm,
90683bc3 216 DemuxerObject::ImageSequence(ref ctx) => &ctx.sm,
91a15e39
KS
217 _ => unreachable!(),
218 }
219 }
220 pub fn get_frame(&mut self) -> DemuxerResult<NAPacket> {
221 match *self {
4a835a2a 222 DemuxerObject::Normal(ref mut dmx) => dmx.get_object_mut().get_frame(),
91a15e39 223 DemuxerObject::Raw(ref mut dmx, ref mut packetisers, ref mut eof) => {
4a835a2a 224 let dmx = dmx.get_object_mut();
91a15e39
KS
225 loop {
226 let mut has_some = false;
227 for (stream, p) in dmx.get_streams().zip(packetisers.iter_mut()) {
228 if let Some(ref mut pkts) = p {
229 match pkts.get_packet(stream.clone()) {
230 Ok(Some(pkt)) => return Ok(pkt),
231 Ok(None) | Err(DecoderError::ShortData) => {
232 if *eof {
233 *p = None;
234 }
235 },
236 Err(err) => {
237 println!("packetisation error {:?}", err);
238 return Err(DemuxerError::InvalidData);
239 }
240 };
241 has_some |= p.is_some();
242 }
243 }
244 if !has_some {
245 return Err(DemuxerError::EOF);
246 }
247 if let Ok(data) = dmx.get_data() {
248 let id = data.get_stream().get_id();
249 for (i, stream) in dmx.get_streams().enumerate() {
250 if stream.get_id() == id {
251 if let Some(ref mut pkts) = packetisers[i] {
252 pkts.add_data(&data.get_buffer());
253 }
254 break;
255 }
256 }
257 } else {
258 *eof = true;
259 }
260 }
261 },
262 DemuxerObject::RawStream(ref mut ctx) => {
4a835a2a 263 let ctx = ctx.get_object_mut();
91a15e39
KS
264 let mut buf = [0; 65536];
265 loop {
266 match ctx.packetiser.get_packet(ctx.stream.clone()) {
206dd66d
KS
267 Ok(Some(mut packet)) => {
268 if packet.get_pts().is_none() && packet.get_duration().is_some() {
269 packet.ts.pts = Some(ctx.pts);
270 }
271 ctx.pts += packet.get_duration().unwrap_or(0);
272 return Ok(packet);
273 },
91a15e39
KS
274 Ok(None) => {},
275 Err(DecoderError::ShortData) => {},
276 _ => return Err(DemuxerError::InvalidData),
277 };
278 match ctx.br.read_buf_some(&mut buf) {
279 Ok(size) => {
280 ctx.packetiser.add_data(&buf[..size]);
281 },
282 Err(_) => {
283 match ctx.packetiser.get_packet(ctx.stream.clone()) {
206dd66d
KS
284 Ok(Some(mut packet)) => {
285 if packet.get_pts().is_none() && packet.get_duration().is_some() {
286 packet.ts.pts = Some(ctx.pts);
287 }
288 ctx.pts += packet.get_duration().unwrap_or(0);
289 return Ok(packet);
290 },
91a15e39
KS
291 Ok(None) | Err(DecoderError::ShortData) => return Err(DemuxerError::EOF),
292 _ => return Err(DemuxerError::InvalidData),
293 };
294 },
295 };
296 }
297 },
90683bc3 298 DemuxerObject::ImageSequence(ref mut ctx) => ctx.get_frame(),
91a15e39
KS
299 _ => unreachable!(),
300 }
301 }
302 pub fn seek(&mut self, seek_time: NATimePoint) -> DemuxerResult<()> {
303 match *self {
4a835a2a
KS
304 DemuxerObject::Normal(ref mut dmx) => dmx.get_object_mut().seek(seek_time),
305 DemuxerObject::Raw(ref mut dmx, _, _) => dmx.get_object_mut().seek(seek_time),
90683bc3 306 DemuxerObject::ImageSequence(ref mut ctx) => ctx.seek(seek_time),
91a15e39
KS
307 _ => Err(DemuxerError::NotImplemented),
308 }
309 }
310}
311
312impl<'a> NAOptionHandler for DemuxerObject<'a> {
313 fn get_supported_options(&self) -> &[NAOptionDefinition] {
314 match *self {
4a835a2a
KS
315 DemuxerObject::Normal(ref dmx) => dmx.get_object().get_supported_options(),
316 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().get_supported_options(),
90683bc3 317 DemuxerObject::ImageSequence(ref ctx) => ctx.get_supported_options(),
91a15e39
KS
318 _ => &[],
319 }
320 }
321 fn set_options(&mut self, options: &[NAOption]) {
322 match *self {
4a835a2a
KS
323 DemuxerObject::Normal(ref mut dmx) => dmx.get_object_mut().set_options(options),
324 DemuxerObject::Raw(ref mut dmx, _, _) => dmx.get_object_mut().set_options(options),
90683bc3 325 DemuxerObject::ImageSequence(ref mut ctx) => ctx.set_options(options),
91a15e39
KS
326 _ => {},
327 }
328 }
329 fn query_option_value(&self, name: &str) -> Option<NAValue> {
330 match *self {
4a835a2a
KS
331 DemuxerObject::Normal(ref dmx) => dmx.get_object().query_option_value(name),
332 DemuxerObject::Raw(ref dmx, _, _) => dmx.get_object().query_option_value(name),
90683bc3 333 DemuxerObject::ImageSequence(ref ctx) => ctx.query_option_value(name),
91a15e39
KS
334 _ => None,
335 }
336 }
337}
338
339pub fn detect_tags(br: &mut ByteReader) -> (bool, u64, Option<u64>) {
340 let mut is_raw = false;
341 let mut start = 0;
342 let mut end = None;
343
344 // check for ID3v{2-4}
345 let mut buf = [0; 5];
346 br.peek_buf(&mut buf).unwrap();
347 if &buf[0..3] == b"ID3" && buf[3] > 0 && buf[3] < 5 && buf[4] == 0 { //ID3 tag found, must be a raw stream
348 br.read_skip(6).unwrap();
349 let mut size = 0;
350 for _ in 0..4 {
351 let b = br.read_byte().unwrap();
352 if (b & 0x80) != 0 {
353 println!("Invalid ID3 size");
354 break;
355 }
356 size = (size << 7) | u64::from(b);
357 }
358 start = size + 10;
359 is_raw = true;
360 }
361 // check for ID3v1
362 br.seek(SeekFrom::End(-128)).unwrap();
363 let off = br.tell();
364 br.peek_buf(&mut buf[..3]).unwrap();
365 if &buf[0..3] == b"TAG" {
366 end = Some(off);
367 }
368 // check for APETAG
369 let mut buf = [0; 8];
370 if let Some(off) = end {
371 br.seek(SeekFrom::Start(off - 32)).unwrap();
372 } else {
373 br.seek(SeekFrom::End(-32)).unwrap();
374 }
375 let off = br.tell();
376 br.read_buf(&mut buf).unwrap();
377 if &buf == b"APETAGEX" {
378 let ver = br.read_u32le().unwrap();
379 let size = u64::from(br.read_u32le().unwrap());
380 let _items = br.read_u32le().unwrap();
381 let flags = br.read_u32le().unwrap();
382 if ver == 1000 || (flags & 0x80000000) == 0 {
383 end = Some(off - size + 32);
384 } else {
385 end = Some(off - size);
386 }
387 }
388
389 (is_raw, start, end)
390}