wmv_decoder/
api.rs

1//! Public library API.
2
3use std::collections::HashMap;
4use std::io::{Read, Seek, SeekFrom};
5
6use crate::asf::{AsfFile, AsfPayload, VideoStreamInfo};
7use crate::decoder::{MacroblockDecoder, YuvFrame};
8use crate::error::{DecoderError, Result};
9use crate::wma::{PcmFrameF32, WmaDecoder};
10use crate::wmv2::{Wmv2FrameHeader, Wmv2FrameType, Wmv2Params};
11
12/// A decoded video frame with timing metadata.
13#[derive(Clone)]
14pub struct DecodedFrame {
15    pub pts_ms: u32,
16    pub is_key_frame: bool,
17    pub frame: YuvFrame,
18}
19
20/// A decoded audio frame with timing metadata.
21#[derive(Clone)]
22pub struct DecodedAudioFrame {
23    pub pts_ms: u32,
24    pub frame: PcmFrameF32,
25}
26
27/// WMV2 (Windows Media Video 8) decoder.
28///
29/// The picture header parsing and macroblock decode paths are aligned with upstream.
30pub struct Wmv2Decoder {
31    params: Wmv2Params,
32    mb_dec: MacroblockDecoder,
33    cur: YuvFrame,
34    locked_hdr_off: Option<usize>,
35}
36
37impl Wmv2Decoder {
38    /// Create a decoder for a fixed resolution.
39    ///
40    /// `extradata` is the 4-byte WMV2 ext header typically carried in ASF stream properties.
41    pub fn new(width: u32, height: u32, extradata: &[u8]) -> Self {
42        let params = Wmv2Params::new(width, height);
43        let mut mb_dec = MacroblockDecoder::new(width, height);
44        mb_dec.wmv2_set_extradata(extradata);
45        let cur = YuvFrame::new(width, height);
46        Self {
47            params,
48            mb_dec,
49            cur,
50            locked_hdr_off: None,
51        }
52    }
53
54    pub fn width(&self) -> u32 {
55        self.params.width
56    }
57
58    pub fn height(&self) -> u32 {
59        self.params.height
60    }
61
62    /// Borrow the internal YUV420p frame buffer.
63    ///
64    /// The returned reference stays valid until the next successful decode.
65    pub fn current_frame(&self) -> &YuvFrame {
66        &self.cur
67    }
68
69
70    /// Decode one assembled WMV2 frame payload.
71    ///
72    /// Returns `Ok(None)` if no plausible picture header can be found.
73    pub fn decode_frame(&mut self, payload: &[u8], is_key_frame: bool) -> Result<Option<&YuvFrame>> {
74        if payload.is_empty() {
75            return Ok(None);
76        }
77
78        let mut best_score: i64 = -1;
79        let mut best_off: usize = 0;
80        let mut best_hdr: Option<Wmv2FrameHeader> = None;
81
82        // Try the previously locked offset first, then fall back to a small scan.
83        let mut offs: Vec<usize> = Vec::with_capacity(18);
84        if let Some(o) = self.locked_hdr_off {
85            offs.push(o);
86        }
87        for o in 0..=16 {
88            if Some(o) != self.locked_hdr_off {
89                offs.push(o);
90            }
91        }
92
93        for off in offs {
94            if off > payload.len() {
95                continue;
96            }
97            let cands = Wmv2FrameHeader::parse_candidates(&payload[off..], self.mb_dec.width_mb, self.mb_dec.height_mb);
98            if cands.is_empty() {
99                continue;
100            }
101            for h in cands {
102                // ASF keyframe marking should correspond to WMV2 I pictures.
103                if is_key_frame && h.frame_type != Wmv2FrameType::I {
104                    continue;
105                }
106
107                // upstream-aligned scoring strategy.
108                let mut sc: i64 = if h.frame_skipped {
109                    1
110                } else if is_key_frame {
111                    2
112                } else {
113                    self.mb_dec.probe_wmv2_payload(&payload[off..], &h) as i64
114                };
115
116                if Some(off) == self.locked_hdr_off {
117                    sc += 64;
118                }
119
120                if sc > best_score {
121                    best_score = sc;
122                    best_off = off;
123                    best_hdr = Some(h);
124                }
125            }
126        }
127
128        let Some(hdr) = best_hdr else {
129            return Ok(None);
130        };
131
132        if self.locked_hdr_off.is_none() {
133            self.locked_hdr_off = Some(best_off);
134        }
135
136        let frame_data = &payload[best_off..];
137        self.mb_dec.decode_wmv2_frame(frame_data, &hdr, &self.params, &mut self.cur)?;
138        Ok(Some(&self.cur))
139    }
140
141    /// Decode and return an owned frame buffer (clone).
142    pub fn decode_frame_owned(&mut self, payload: &[u8], is_key_frame: bool) -> Result<Option<YuvFrame>> {
143        let Some(f) = self.decode_frame(payload, is_key_frame)? else {
144            return Ok(None);
145        };
146        Ok(Some(f.clone()))
147    }
148}
149
150// ─────────────────────────────────────────────────────────────────────────────
151// ASF media-object reassembly (frame reassembly)
152// ─────────────────────────────────────────────────────────────────────────────
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
155struct FrameKey {
156    stream_number: u8,
157    object_id: u32,
158}
159
160#[derive(Debug, Clone)]
161struct FrameAssembly {
162    total: usize,
163    pts_ms: u32,
164    is_key: bool,
165    data: Vec<u8>,
166    ranges: Vec<(usize, usize)>,
167}
168
169impl FrameAssembly {
170    fn new(total: usize, pts_ms: u32, is_key: bool) -> Self {
171        Self {
172            total,
173            pts_ms,
174            is_key,
175            data: vec![0u8; total],
176            ranges: Vec::new(),
177        }
178    }
179
180    fn insert(&mut self, offset: usize, frag: &[u8]) {
181        if self.total == 0 || offset >= self.total || frag.is_empty() {
182            return;
183        }
184        let end = (offset + frag.len()).min(self.total);
185        let n = end - offset;
186        self.data[offset..end].copy_from_slice(&frag[..n]);
187        self.add_range(offset, end);
188    }
189
190    fn add_range(&mut self, start: usize, end: usize) {
191        if start >= end {
192            return;
193        }
194        self.ranges.push((start, end));
195        self.ranges.sort_by_key(|r| r.0);
196
197        let mut merged: Vec<(usize, usize)> = Vec::with_capacity(self.ranges.len());
198        for (s, e) in self.ranges.drain(..) {
199            if let Some(last) = merged.last_mut() {
200                if s <= last.1 {
201                    last.1 = last.1.max(e);
202                    continue;
203                }
204            }
205            merged.push((s, e));
206        }
207        self.ranges = merged;
208    }
209
210    fn covered_len(&self) -> usize {
211        self.ranges.iter().map(|(s, e)| e - s).sum()
212    }
213
214    fn is_complete(&self) -> bool {
215        self.total > 0
216            && self.covered_len() >= self.total
217            && self.ranges.len() == 1
218            && self.ranges[0] == (0, self.total)
219    }
220}
221
222#[derive(Default)]
223struct FrameAssembler {
224    in_flight: HashMap<FrameKey, FrameAssembly>,
225}
226
227impl FrameAssembler {
228    fn push(&mut self, payload: AsfPayload) -> Option<(u32, bool, Vec<u8>)> {
229        if payload.data.is_empty() {
230            return None;
231        }
232
233        let key = FrameKey {
234            stream_number: payload.stream_number,
235            object_id: payload.object_id,
236        };
237
238        // Fast path: complete media object in one payload (or size unknown).
239        if payload.obj_offset == 0 {
240            let osz = payload.obj_size as usize;
241            if osz == 0 || osz == payload.data.len() {
242                return Some((payload.pts_ms, payload.is_key_frame, payload.data));
243            }
244        }
245
246        // If the total object size is unknown, we cannot reliably reassemble.
247        if payload.obj_size == 0 {
248            return Some((payload.pts_ms, payload.is_key_frame, payload.data));
249        }
250
251        let total = payload.obj_size as usize;
252        let entry = self
253            .in_flight
254            .entry(key)
255            .or_insert_with(|| FrameAssembly::new(total, payload.pts_ms, payload.is_key_frame));
256
257        // Update meta (first PTS wins; keyframe if any fragment says so).
258        entry.is_key |= payload.is_key_frame;
259        entry.insert(payload.obj_offset as usize, &payload.data);
260
261        if entry.is_complete() {
262            let assembly = self.in_flight.remove(&key).unwrap();
263            return Some((assembly.pts_ms, assembly.is_key, assembly.data));
264        }
265        None
266    }
267}
268
269/// ASF + WMV2 decoding pipeline.
270///
271/// This type owns the `Read+Seek` source, parses ASF headers, reassembles media objects
272/// and decodes WMV2 frames into `YuvFrame`.
273pub struct AsfWmv2Decoder<R: Read + Seek> {
274    reader: R,
275    asf: AsfFile,
276    video_info: VideoStreamInfo,
277    assembler: FrameAssembler,
278    decoder: Wmv2Decoder,
279}
280
281/// ASF + WMA (v1/v2) decoding pipeline.
282///
283/// This type owns the `Read+Seek` source, parses ASF headers, reassembles media objects
284/// and decodes WMA packets into PCM.
285pub struct AsfWmaDecoder<R: Read + Seek> {
286    reader: R,
287    asf: AsfFile,
288    audio_stream_number: u8,
289    decoder: WmaDecoder,
290    assembler: FrameAssembler,
291    last_pts_ms: u32,
292    flushed_eof: bool,
293}
294
295impl<R: Read + Seek> AsfWmaDecoder<R> {
296    /// Open an ASF/WMV stream and initialize the WMA decoder.
297    ///
298    /// The decoder selects the first audio stream with format tag 0x0160 (WMAv1)
299    /// or 0x0161 (WMAv2).
300    pub fn open(mut reader: R) -> Result<Self> {
301        let asf = AsfFile::open(&mut reader)?;
302        let mut chosen = None;
303        for a in asf.audio_streams.iter() {
304            if matches!(a.format_tag, 0x0160 | 0x0161) {
305                chosen = Some(a.clone());
306                break;
307            }
308        }
309        let Some(audio_info) = chosen else {
310            return Err(DecoderError::Unsupported(
311                "No supported WMA (0x0160/0x0161) audio stream found".into(),
312            ));
313        };
314
315        reader.seek(SeekFrom::Start(asf.data_offset))?;
316        let decoder = WmaDecoder::new(&audio_info)?;
317
318        Ok(Self {
319            reader,
320            asf,
321            audio_stream_number: audio_info.stream_number,
322            decoder,
323            assembler: FrameAssembler::default(),
324            last_pts_ms: 0,
325            flushed_eof: false,
326        })
327    }
328
329    pub fn sample_rate(&self) -> u32 {
330        self.decoder.sample_rate()
331    }
332
333    pub fn channels(&self) -> u16 {
334        self.decoder.channels()
335    }
336
337    /// Decode the next audio frame.
338    ///
339    /// Returns `Ok(None)` on end-of-stream.
340    pub fn next_frame(&mut self) -> Result<Option<DecodedAudioFrame>> {
341        loop {
342            let payloads = match self.asf.read_packet(&mut self.reader) {
343                Ok(p) => p,
344                Err(DecoderError::EndOfStream) => {
345                    if self.flushed_eof {
346                        return Ok(None);
347                    }
348                    self.flushed_eof = true;
349                    if let Some(frame) = self.decoder.decode_packet(&[], self.last_pts_ms)? {
350                        return Ok(Some(DecodedAudioFrame {
351                            pts_ms: frame.pts_ms,
352                            frame,
353                        }));
354                    }
355                    return Ok(None);
356                }
357                Err(e) => return Err(e),
358            };
359
360            for payload in payloads {
361                if payload.stream_number != self.audio_stream_number {
362                    continue;
363                }
364                let Some((pts_ms, _is_key, data)) = self.assembler.push(payload) else {
365                    continue;
366                };
367                self.last_pts_ms = pts_ms;
368                if let Some(frame) = self.decoder.decode_packet(&data, pts_ms)? {
369                    return Ok(Some(DecodedAudioFrame { pts_ms, frame }));
370                }
371            }
372        }
373    }
374}
375
376impl<R: Read + Seek> AsfWmv2Decoder<R> {
377    /// Open an ASF/WMV stream and initialize the WMV2 decoder.
378    ///
379    /// The decoder selects the first video stream whose FourCC is WMV2 or WMV1.
380    pub fn open(mut reader: R) -> Result<Self> {
381        let asf = AsfFile::open(&mut reader)?;
382        let mut video_info: Option<VideoStreamInfo> = None;
383        for v in asf.video_streams.iter() {
384            let four_cc = std::str::from_utf8(&v.codec_four_cc)
385                .unwrap_or("")
386                .to_uppercase();
387            if matches!(four_cc.as_str(), "WMV2" | "WMV1") {
388                video_info = Some(v.clone());
389                break;
390            }
391        }
392        let Some(video_info) = video_info else {
393            return Err(DecoderError::Unsupported("No WMV2/WMV1 video stream found".into()));
394        };
395
396        reader.seek(SeekFrom::Start(asf.data_offset))?;
397
398        let decoder = Wmv2Decoder::new(video_info.width, video_info.height, &video_info.extra_data);
399
400        Ok(Self {
401            reader,
402            asf,
403            video_info,
404            assembler: FrameAssembler::default(),
405            decoder,
406        })
407    }
408
409    /// Return the selected video stream info.
410    pub fn video_stream_info(&self) -> &VideoStreamInfo {
411        &self.video_info
412    }
413
414    /// Decode the next video frame.
415    ///
416    /// Returns `Ok(None)` on end-of-stream.
417    pub fn next_frame(&mut self) -> Result<Option<DecodedFrame>> {
418        loop {
419            let payloads = match self.asf.read_packet(&mut self.reader) {
420                Ok(p) => p,
421                Err(DecoderError::EndOfStream) => return Ok(None),
422                Err(e) => return Err(e),
423            };
424
425            for payload in payloads {
426                if payload.stream_number != self.video_info.stream_number {
427                    continue;
428                }
429                let Some((pts_ms, is_key, data)) = self.assembler.push(payload) else {
430                    continue;
431                };
432
433                if let Some(frame) = self.decoder.decode_frame_owned(&data, is_key)? {
434                    return Ok(Some(DecodedFrame {
435                        pts_ms,
436                        is_key_frame: is_key,
437                        frame,
438                    }));
439                }
440            }
441        }
442    }
443}
wmv_decoder/api.rs

wmv_decoder/
api.rs
