#![allow(dead_code)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MergeStrategy {
PreferFirst,
PreferLast,
KeepAll,
DropOnConflict,
}
impl MergeStrategy {
pub fn name(&self) -> &'static str {
match self {
MergeStrategy::PreferFirst => "prefer_first",
MergeStrategy::PreferLast => "prefer_last",
MergeStrategy::KeepAll => "keep_all",
MergeStrategy::DropOnConflict => "drop_on_conflict",
}
}
}
#[derive(Debug, Clone)]
pub struct SubtitleEntry {
pub track: usize,
pub start_ms: i64,
pub end_ms: i64,
pub text: String,
pub conflicted: bool,
}
impl SubtitleEntry {
pub fn new(track: usize, start_ms: i64, end_ms: i64, text: impl Into<String>) -> Self {
Self {
track,
start_ms,
end_ms,
text: text.into(),
conflicted: false,
}
}
pub fn overlaps(&self, other: &SubtitleEntry) -> bool {
self.start_ms < other.end_ms && self.end_ms > other.start_ms
}
pub fn duration_ms(&self) -> i64 {
self.end_ms - self.start_ms
}
pub fn mark_conflicted(&mut self) {
self.conflicted = true;
}
}
#[derive(Debug)]
pub struct SubtitleMerger {
strategy: MergeStrategy,
tracks: Vec<Vec<SubtitleEntry>>,
}
impl SubtitleMerger {
pub fn new(strategy: MergeStrategy) -> Self {
Self {
strategy,
tracks: Vec::new(),
}
}
pub fn add_track(&mut self, entries: Vec<SubtitleEntry>) {
self.tracks.push(entries);
}
pub fn track_count(&self) -> usize {
self.tracks.len()
}
pub fn merge(&self) -> MergeResult {
let mut merged: Vec<SubtitleEntry> = Vec::new();
let mut conflicts: usize = 0;
for (track_idx, track) in self.tracks.iter().enumerate() {
for entry in track {
let mut new_entry = entry.clone();
new_entry.track = track_idx;
let overlap_count = merged.iter().filter(|e| e.overlaps(&new_entry)).count();
if overlap_count > 0 {
conflicts += 1;
match self.strategy {
MergeStrategy::PreferFirst => {
continue;
}
MergeStrategy::PreferLast => {
merged.retain(|e| !e.overlaps(&new_entry));
merged.push(new_entry);
}
MergeStrategy::KeepAll => {
new_entry.mark_conflicted();
merged.push(new_entry);
}
MergeStrategy::DropOnConflict => {
continue;
}
}
} else {
merged.push(new_entry);
}
}
}
merged.sort_by_key(|e| e.start_ms);
MergeResult {
entries: merged,
conflict_count: conflicts,
}
}
}
#[derive(Debug)]
pub struct MergeResult {
pub entries: Vec<SubtitleEntry>,
pub conflict_count: usize,
}
impl MergeResult {
pub fn entry_count(&self) -> usize {
self.entries.len()
}
pub fn conflict_count(&self) -> usize {
self.conflict_count
}
pub fn conflicted_entries(&self) -> Vec<&SubtitleEntry> {
self.entries.iter().filter(|e| e.conflicted).collect()
}
pub fn span_ms(&self) -> i64 {
let start = self.entries.first().map(|e| e.start_ms).unwrap_or(0);
let end = self.entries.last().map(|e| e.end_ms).unwrap_or(0);
end - start
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_entry(track: usize, start: i64, end: i64, text: &str) -> SubtitleEntry {
SubtitleEntry::new(track, start, end, text)
}
#[test]
fn test_merge_strategy_name() {
assert_eq!(MergeStrategy::PreferFirst.name(), "prefer_first");
assert_eq!(MergeStrategy::PreferLast.name(), "prefer_last");
assert_eq!(MergeStrategy::KeepAll.name(), "keep_all");
assert_eq!(MergeStrategy::DropOnConflict.name(), "drop_on_conflict");
}
#[test]
fn test_subtitle_entry_overlaps_true() {
let a = make_entry(0, 1000, 4000, "A");
let b = make_entry(1, 3000, 6000, "B");
assert!(a.overlaps(&b));
assert!(b.overlaps(&a));
}
#[test]
fn test_subtitle_entry_overlaps_false() {
let a = make_entry(0, 1000, 3000, "A");
let b = make_entry(1, 3000, 5000, "B");
assert!(!a.overlaps(&b));
}
#[test]
fn test_subtitle_entry_duration_ms() {
let e = make_entry(0, 1000, 4500, "Hi");
assert_eq!(e.duration_ms(), 3500);
}
#[test]
fn test_subtitle_entry_mark_conflicted() {
let mut e = make_entry(0, 0, 1000, "X");
assert!(!e.conflicted);
e.mark_conflicted();
assert!(e.conflicted);
}
#[test]
fn test_merger_add_track() {
let mut merger = SubtitleMerger::new(MergeStrategy::PreferFirst);
merger.add_track(vec![make_entry(0, 0, 1000, "A")]);
assert_eq!(merger.track_count(), 1);
}
#[test]
fn test_merge_no_overlap() {
let mut merger = SubtitleMerger::new(MergeStrategy::PreferFirst);
merger.add_track(vec![make_entry(0, 0, 1000, "A")]);
merger.add_track(vec![make_entry(1, 2000, 3000, "B")]);
let result = merger.merge();
assert_eq!(result.entry_count(), 2);
assert_eq!(result.conflict_count(), 0);
}
#[test]
fn test_merge_prefer_first_drops_conflict() {
let mut merger = SubtitleMerger::new(MergeStrategy::PreferFirst);
merger.add_track(vec![make_entry(0, 0, 3000, "First")]);
merger.add_track(vec![make_entry(1, 1000, 4000, "Conflict")]);
let result = merger.merge();
assert_eq!(result.entry_count(), 1);
assert_eq!(result.entries[0].text, "First");
}
#[test]
fn test_merge_prefer_last_replaces() {
let mut merger = SubtitleMerger::new(MergeStrategy::PreferLast);
merger.add_track(vec![make_entry(0, 0, 3000, "First")]);
merger.add_track(vec![make_entry(1, 1000, 4000, "Last")]);
let result = merger.merge();
assert_eq!(result.entry_count(), 1);
assert_eq!(result.entries[0].text, "Last");
}
#[test]
fn test_merge_keep_all_marks_conflicts() {
let mut merger = SubtitleMerger::new(MergeStrategy::KeepAll);
merger.add_track(vec![make_entry(0, 0, 3000, "A")]);
merger.add_track(vec![make_entry(1, 1000, 4000, "B")]);
let result = merger.merge();
assert_eq!(result.entry_count(), 2);
assert_eq!(result.conflict_count(), 1);
assert_eq!(result.conflicted_entries().len(), 1);
}
#[test]
fn test_merge_drop_on_conflict() {
let mut merger = SubtitleMerger::new(MergeStrategy::DropOnConflict);
merger.add_track(vec![make_entry(0, 0, 3000, "A")]);
merger.add_track(vec![make_entry(1, 1000, 4000, "B")]);
let result = merger.merge();
assert_eq!(result.entry_count(), 1);
assert_eq!(result.entries[0].text, "A");
}
#[test]
fn test_merge_result_span_ms() {
let mut merger = SubtitleMerger::new(MergeStrategy::KeepAll);
merger.add_track(vec![
make_entry(0, 1000, 3000, "A"),
make_entry(0, 5000, 8000, "B"),
]);
let result = merger.merge();
assert_eq!(result.span_ms(), 7000);
}
#[test]
fn test_merge_sorted_output() {
let mut merger = SubtitleMerger::new(MergeStrategy::KeepAll);
merger.add_track(vec![
make_entry(0, 5000, 7000, "Late"),
make_entry(0, 1000, 3000, "Early"),
]);
let result = merger.merge();
assert_eq!(result.entries[0].start_ms, 1000);
assert_eq!(result.entries[1].start_ms, 5000);
}
}