Skip to main content

libbitsub_core/vobsub/
vobsub_parser.rs

1//! VobSub parser.
2
3use memchr::memchr;
4use std::collections::HashMap;
5
6use super::{
7    DebandConfig, ExtractedVobSub, IdxParseResult, SubtitlePacket, VobSubPalette, VobSubTimestamp,
8    apply_deband, decode_vobsub_rle, extract_vobsub_from_mks, parse_idx, parse_subtitle_packet,
9};
10use crate::utils::binary_search_timestamp;
11
12/// VobSub subtitle parser and renderer.
13pub struct VobSubParser {
14    /// Parsed IDX data
15    idx_data: Option<IdxParseResult>,
16    /// Raw SUB file data
17    sub_data: Option<Vec<u8>>,
18    /// Timestamps in milliseconds for quick lookup
19    timestamps_ms: Vec<u32>,
20    /// Cache for decoded subtitle packets
21    packet_cache: HashMap<usize, Option<SubtitlePacket>>,
22    /// Debanding configuration
23    deband_config: DebandConfig,
24    /// Whether the parser was loaded from IDX metadata.
25    loaded_from_idx: bool,
26    /// Last non-fatal render issue for diagnostics.
27    last_render_issue: Option<String>,
28}
29
30impl VobSubParser {
31    /// Create a new VobSub parser.
32    pub fn new() -> Self {
33        Self {
34            idx_data: None,
35            sub_data: None,
36            timestamps_ms: Vec::new(),
37            packet_cache: HashMap::new(),
38            deband_config: DebandConfig::default(),
39            loaded_from_idx: false,
40            last_render_issue: None,
41        }
42    }
43
44    /// Load VobSub from IDX content and SUB data.
45    pub fn load_from_data(&mut self, idx_content: &str, sub_data: Vec<u8>) {
46        self.dispose();
47        self.apply_loaded_data(parse_idx(idx_content), sub_data, true);
48    }
49
50    /// Load VobSub from a Matroska subtitle container with embedded S_VOBSUB tracks.
51    pub fn load_from_mks(&mut self, mks_data: &[u8]) -> Result<(), String> {
52        self.dispose();
53
54        let ExtractedVobSub {
55            idx_content,
56            sub_data,
57            language,
58            track_id,
59        } = extract_vobsub_from_mks(mks_data)?;
60
61        let mut idx = parse_idx(&idx_content);
62        if language.is_some() {
63            idx.metadata.language = language;
64        }
65        if track_id.is_some() {
66            idx.metadata.id = track_id;
67        }
68
69        self.apply_loaded_data(idx, sub_data, true);
70        Ok(())
71    }
72
73    /// Load VobSub from SUB file only (scans for timestamps).
74    pub fn load_from_sub_only(&mut self, sub_data: Vec<u8>) {
75        self.dispose();
76
77        // Default palette
78        let palette = VobSubPalette::default();
79
80        // Pre-allocate with estimate (roughly 1 subtitle per 10KB)
81        let estimated_count = (sub_data.len() / 10000).max(32);
82        let mut timestamps: Vec<VobSubTimestamp> = Vec::with_capacity(estimated_count);
83        let mut offset = 0;
84        let len = sub_data.len();
85
86        // Look for MPEG-2 PS pack start code
87        while offset < len.saturating_sub(4) {
88            // Find next potential start code (0x00 0x00 0x01 0xBA)
89            if let Some(pos) = memchr(0x00, &sub_data[offset..]) {
90                let candidate = offset + pos;
91
92                // Check for full start code: 00 00 01 BA
93                if candidate + 3 < len
94                    && sub_data[candidate + 1] == 0x00
95                    && sub_data[candidate + 2] == 0x01
96                    && sub_data[candidate + 3] == 0xBA
97                    && let Some((packet, _)) = parse_subtitle_packet(&sub_data, candidate, &palette)
98                    && packet.width > 0
99                    && packet.height > 0
100                {
101                    timestamps.push(VobSubTimestamp {
102                        timestamp_ms: packet.timestamp_ms,
103                        file_position: candidate as u64,
104                    });
105                }
106                offset = candidate + 1;
107            } else {
108                // No more 0x00 bytes found
109                break;
110            }
111        }
112
113        // Sort and store
114        timestamps.sort_by_key(|t| t.timestamp_ms);
115        self.timestamps_ms = timestamps.iter().map(|t| t.timestamp_ms).collect();
116
117        let idx = IdxParseResult {
118            palette,
119            timestamps,
120            metadata: Default::default(),
121        };
122        self.apply_loaded_data(idx, sub_data, false);
123    }
124
125    /// Dispose of all resources.
126    pub fn dispose(&mut self) {
127        self.idx_data = None;
128        self.sub_data = None;
129        self.timestamps_ms.clear();
130        self.packet_cache.clear();
131        self.deband_config = DebandConfig::default();
132        self.loaded_from_idx = false;
133        self.last_render_issue = None;
134    }
135
136    /// Get the last non-fatal render issue for diagnostics.
137    pub fn last_render_issue(&self) -> String {
138        self.last_render_issue.clone().unwrap_or_default()
139    }
140
141    /// Get the number of subtitle entries.
142    pub fn count(&self) -> usize {
143        self.timestamps_ms.len()
144    }
145
146    /// Get the presentation width for this subtitle track.
147    pub fn screen_width(&self) -> u16 {
148        self.idx_data
149            .as_ref()
150            .map_or(0, |idx_data| idx_data.metadata.width)
151    }
152
153    /// Get the presentation height for this subtitle track.
154    pub fn screen_height(&self) -> u16 {
155        self.idx_data
156            .as_ref()
157            .map_or(0, |idx_data| idx_data.metadata.height)
158    }
159
160    /// Get the declared language code from IDX metadata.
161    pub fn language(&self) -> String {
162        self.idx_data
163            .as_ref()
164            .and_then(|idx_data| idx_data.metadata.language.clone())
165            .unwrap_or_default()
166    }
167
168    /// Get the declared subtitle track ID from IDX metadata.
169    pub fn track_id(&self) -> String {
170        self.idx_data
171            .as_ref()
172            .and_then(|idx_data| idx_data.metadata.id.clone())
173            .unwrap_or_default()
174    }
175
176    /// Check whether IDX metadata was used to load the parser.
177    pub fn has_idx_metadata(&self) -> bool {
178        self.loaded_from_idx
179    }
180
181    /// Get all timestamps in milliseconds.
182    pub fn get_timestamps(&self) -> Vec<f64> {
183        self.timestamps_ms.iter().map(|&ts| ts as f64).collect()
184    }
185
186    /// Find the subtitle index for a given timestamp in milliseconds.
187    /// Returns -1 if no subtitle should be displayed at this time.
188    pub fn find_index_at_timestamp(&mut self, time_ms: f64) -> i32 {
189        if self.timestamps_ms.is_empty() {
190            return -1;
191        }
192
193        let time_ms_u32 = time_ms as u32;
194        let index = binary_search_timestamp(&self.timestamps_ms, time_ms_u32);
195
196        // Get the start time from IDX (what we searched against)
197        let start_time = self.timestamps_ms[index];
198
199        // Don't show if we're before this subtitle's start time
200        if time_ms_u32 < start_time {
201            return -1;
202        }
203
204        // Calculate end time
205        let end_time = self.calculate_end_time(index, start_time);
206
207        if time_ms_u32 < end_time {
208            return index as i32;
209        }
210
211        // Current time is past the subtitle's duration
212        -1
213    }
214
215    /// Get the cue start time in milliseconds.
216    pub fn get_cue_start_time(&self, index: usize) -> f64 {
217        self.timestamps_ms
218            .get(index)
219            .copied()
220            .map_or(-1.0, |ts| ts as f64)
221    }
222
223    /// Get the cue end time in milliseconds.
224    pub fn get_cue_end_time(&mut self, index: usize) -> f64 {
225        let Some(&start_time) = self.timestamps_ms.get(index) else {
226            return -1.0;
227        };
228
229        self.calculate_end_time(index, start_time) as f64
230    }
231
232    /// Get the cue duration in milliseconds.
233    pub fn get_cue_duration(&mut self, index: usize) -> f64 {
234        let Some(&start_time) = self.timestamps_ms.get(index) else {
235            return -1.0;
236        };
237
238        self.calculate_end_time(index, start_time)
239            .saturating_sub(start_time) as f64
240    }
241
242    /// Get the cue file position in the SUB file.
243    pub fn get_cue_file_position(&self, index: usize) -> f64 {
244        self.idx_data
245            .as_ref()
246            .and_then(|idx_data| idx_data.timestamps.get(index).copied())
247            .map_or(-1.0, |timestamp| timestamp.file_position as f64)
248    }
249
250    fn apply_loaded_data(
251        &mut self,
252        idx_data: IdxParseResult,
253        sub_data: Vec<u8>,
254        loaded_from_idx: bool,
255    ) {
256        self.timestamps_ms = idx_data.timestamps.iter().map(|t| t.timestamp_ms).collect();
257        self.idx_data = Some(idx_data);
258        self.sub_data = Some(sub_data);
259        self.loaded_from_idx = loaded_from_idx;
260    }
261
262    /// Calculate the end time for a subtitle at the given index.
263    fn calculate_end_time(&mut self, index: usize, start_time: u32) -> u32 {
264        // Maximum duration for the last subtitle (no next subtitle to clamp to)
265        const MAX_LAST_DURATION_MS: u32 = 5000;
266
267        // Try to get explicit duration from control sequence first
268        self.ensure_packet_cached(index);
269        let explicit_duration = self
270            .cached_packet(index)
271            .filter(|p| p.duration_ms > 0 && p.duration_ms != 5000)
272            .map(|p| p.duration_ms);
273
274        // Check if we have a next subtitle
275        if index + 1 < self.timestamps_ms.len() {
276            let next_start = self.timestamps_ms[index + 1];
277
278            if let Some(duration) = explicit_duration {
279                let explicit_end = start_time.saturating_add(duration);
280                return explicit_end.min(next_start);
281            }
282
283            next_start
284        } else {
285            // Last subtitle - use explicit duration if valid, otherwise default
286            if let Some(duration) = explicit_duration {
287                return start_time.saturating_add(duration);
288            }
289            // Default duration for last subtitle
290            start_time.saturating_add(MAX_LAST_DURATION_MS)
291        }
292    }
293
294    fn ensure_packet_cached(&mut self, index: usize) -> Option<()> {
295        let idx_data = self.idx_data.as_ref()?;
296        if index >= idx_data.timestamps.len() {
297            return None;
298        }
299
300        if self.packet_cache.contains_key(&index) {
301            return Some(());
302        }
303
304        let packet = {
305            let idx_data = self.idx_data.as_ref()?;
306            let sub_data = self.sub_data.as_ref()?;
307            let timestamp = idx_data.timestamps.get(index)?;
308
309            parse_subtitle_packet(
310                sub_data,
311                timestamp.file_position as usize,
312                &idx_data.palette,
313            )
314            .map(|(p, _)| p)
315        };
316
317        self.packet_cache.insert(index, packet);
318        Some(())
319    }
320
321    fn cached_packet(&self, index: usize) -> Option<&SubtitlePacket> {
322        self.packet_cache
323            .get(&index)
324            .and_then(|packet| packet.as_ref())
325    }
326
327    /// Render subtitle at the given index and return RGBA data.
328    pub fn render_at_index(&mut self, index: usize) -> Option<VobSubFrame> {
329        self.last_render_issue = None;
330
331        if index >= self.timestamps_ms.len() {
332            self.last_render_issue = Some("INDEX_OUT_OF_RANGE".to_string());
333            return None;
334        }
335
336        if self.ensure_packet_cached(index).is_none() {
337            self.last_render_issue = Some("NO_DATA".to_string());
338            return None;
339        }
340
341        let Some(idx_data) = self.idx_data.as_ref() else {
342            self.last_render_issue = Some("NO_DATA".to_string());
343            return None;
344        };
345        let Some(sub_data) = self.sub_data.as_ref() else {
346            self.last_render_issue = Some("NO_DATA".to_string());
347            return None;
348        };
349        let Some(packet) = self.cached_packet(index) else {
350            self.last_render_issue = Some("INVALID_PACKET".to_string());
351            return None;
352        };
353
354        Some(self.render_packet(packet, sub_data, &idx_data.palette, &idx_data.metadata))
355    }
356
357    /// Render a packet to a frame.
358    fn render_packet(
359        &self,
360        packet: &SubtitlePacket,
361        sub_data: &[u8],
362        palette: &VobSubPalette,
363        metadata: &super::VobSubMetadata,
364    ) -> VobSubFrame {
365        let mut rgba = decode_vobsub_rle(packet, sub_data, palette);
366
367        // Apply debanding if enabled
368        if self.deband_config.enabled {
369            rgba = apply_deband(
370                &rgba,
371                packet.width as usize,
372                packet.height as usize,
373                &self.deband_config,
374            );
375        }
376
377        VobSubFrame {
378            screen_width: metadata.width,
379            screen_height: metadata.height,
380            x: packet.x,
381            y: packet.y,
382            width: packet.width,
383            height: packet.height,
384            rgba,
385        }
386    }
387
388    /// Clear the internal cache.
389    pub fn clear_cache(&mut self) {
390        self.packet_cache.clear();
391    }
392
393    /// Enable or disable debanding.
394    pub fn set_deband_enabled(&mut self, enabled: bool) {
395        self.deband_config.enabled = enabled;
396    }
397
398    /// Set the deband threshold (0.0-255.0, default: 64.0).
399    pub fn set_deband_threshold(&mut self, threshold: f32) {
400        self.deband_config.threshold = threshold.clamp(0.0, 255.0);
401    }
402
403    /// Set the deband sample range in pixels (default: 15).
404    pub fn set_deband_range(&mut self, range: u32) {
405        self.deband_config.range = range.clamp(1, 64);
406    }
407
408    /// Check if debanding is enabled.
409    pub fn deband_enabled(&self) -> bool {
410        self.deband_config.enabled
411    }
412}
413
414impl Default for VobSubParser {
415    fn default() -> Self {
416        Self::new()
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn dispose_restores_default_deband_config() {
426        let mut parser = VobSubParser::new();
427
428        parser.set_deband_enabled(false);
429        parser.set_deband_threshold(12.0);
430        parser.set_deband_range(3);
431
432        parser.dispose();
433
434        assert!(parser.deband_enabled());
435        assert_eq!(
436            parser.deband_config.threshold,
437            DebandConfig::default().threshold
438        );
439        assert_eq!(parser.deband_config.range, DebandConfig::default().range);
440    }
441}
442
443/// A VobSub subtitle frame.
444pub struct VobSubFrame {
445    pub screen_width: u16,
446    pub screen_height: u16,
447    pub x: u16,
448    pub y: u16,
449    pub width: u16,
450    pub height: u16,
451    pub rgba: Vec<u8>,
452}
453
454impl VobSubFrame {
455    pub fn screen_width(&self) -> u16 {
456        self.screen_width
457    }
458    pub fn screen_height(&self) -> u16 {
459        self.screen_height
460    }
461    pub fn x(&self) -> u16 {
462        self.x
463    }
464    pub fn y(&self) -> u16 {
465        self.y
466    }
467    pub fn width(&self) -> u16 {
468        self.width
469    }
470    pub fn height(&self) -> u16 {
471        self.height
472    }
473
474    /// Get RGBA pixel data.
475    pub fn get_rgba(&self) -> &[u8] {
476        &self.rgba
477    }
478}