remuxer/src/import/seek_index.rs


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use anyhow::Result;
use jellycommon::{BlockIndex, SeekIndex};
use jellymatroska::{
    block::Block,
    read::EbmlReader,
    unflatten::{IterWithPos, Unflat, Unflatten},
    MatroskaTag,
};
use log::{debug, info, trace, warn};
use std::collections::BTreeMap;

pub fn import_seek_index(input: &mut EbmlReader) -> Result<BTreeMap<u64, SeekIndex>> {
    let mut seek_index = BTreeMap::new();
    while let Some(item) = input.next() {
        let item = match item {
            Ok(item) => item,
            Err(e) => {
                warn!("{e}");
                break;
            }
        };
        match item {
            MatroskaTag::Segment(_) => {
                info!("segment start");
                let mut children = Unflatten::new_with_end(input, item);
                import_seek_index_segment(&mut children, &mut seek_index)?;
                info!("segment end");
            }
            _ => debug!("(r) tag ignored: {item:?}"),
        }
    }
    Ok(seek_index)
}

fn import_seek_index_segment(
    segment: &mut Unflatten,
    seek_index: &mut BTreeMap<u64, SeekIndex>,
) -> Result<()> {
    while let Some(Ok(Unflat { children, item, .. })) = segment.n() {
        match item {
            MatroskaTag::SeekHead(_) => {}
            MatroskaTag::Info(_) => {}
            MatroskaTag::Tags(_) => {}
            MatroskaTag::Cues(_) => {}
            MatroskaTag::Chapters(_) => {}
            MatroskaTag::Tracks(_) => {}
            MatroskaTag::Cluster(_) => {
                let mut children = children.unwrap();
                let mut pts = 0;
                let mut position = children.position();

                loop {
                    if let Some(Ok(Unflat { children, item, .. })) = children.n() {
                        match item {
                            MatroskaTag::Timestamp(ts) => pts = ts,
                            MatroskaTag::BlockGroup(_) => {
                                debug!("group");
                                let mut children = children.unwrap();
                                // let position = children.position(); //? TODO where should this point to? cluster or block? // probably block
                                while let Some(Ok(Unflat {
                                    children: _,
                                    item,
                                    position,
                                })) = children.n()
                                {
                                    match item {
                                        MatroskaTag::Block(ref buf) => {
                                            let block = Block::parse(buf)?;
                                            debug!(
                                                "block: track={} tso={}",
                                                block.track, block.timestamp_off
                                            );
                                            seek_index
                                                .entry(block.track)
                                                .or_insert(SeekIndex { blocks: vec![] })
                                                .blocks
                                                .push(BlockIndex {
                                                    pts: pts + block.timestamp_off as u64,
                                                    source_off: position,
                                                    size: block.data.len(),
                                                });
                                        }
                                        _ => trace!("{item:?}"),
                                    }
                                }
                            }
                            MatroskaTag::SimpleBlock(buf) => {
                                let block = Block::parse(&buf)?;
                                debug!(
                                    "simple block: track={} tso={}",
                                    block.track, block.timestamp_off
                                );
                                debug!("{pts} {}", block.timestamp_off);
                                seek_index
                                    .entry(block.track)
                                    .or_insert(SeekIndex { blocks: vec![] })
                                    .blocks
                                    .push(BlockIndex {
                                        pts: (pts as i64 + block.timestamp_off as i64) as u64,
                                        source_off: position,
                                        size: block.data.len(),
                                    });
                            }
                            _ => debug!("(rsc) tag ignored: {item:?}"),
                        }
                    } else {
                        break;
                    }
                    position = children.position();
                }
            }
            _ => debug!("(rs) tag ignored: {item:?}"),
        };
    }
    Ok(())
}