generate seek index

1 files changed, 69 insertions, 53 deletions

diff --git a/remuxer/src/import/mod.rs b/remuxer/src/import/mod.rs
index 65fc455..04fb88c 100644
--- a/remuxer/src/import/mod.rs
+++ b/remuxer/src/import/mod.rs
@@ -4,15 +4,15 @@
     Copyright (C) 2023 metamuffin <metamuffin.org>
 */
 use anyhow::{anyhow, bail, Result};
-use jellycommon::{ItemInfo, SourceTrack, SourceTrackKind};
+use jellycommon::{BlockIndex, ItemInfo, SeekIndex, SourceTrack, SourceTrackKind};
 use jellymatroska::{
     block::Block,
     matroska::MatroskaTag,
     read::EbmlReader,
-    unflatten::{Unflat, Unflatten},
+    unflatten::{IterWithPos, Unflat, Unflatten},
 };
 use log::{debug, error, info, trace, warn};
-use std::path::PathBuf;
+use std::{collections::HashMap, fs::File, path::PathBuf};
 
 pub fn import_read(path: &PathBuf, input: &mut EbmlReader, iteminfo: &mut ItemInfo) -> Result<()> {
     // TODO dont traverse the entire file, if the tracks are listed at the end
@@ -57,6 +57,8 @@ fn import_read_segment(
     iteminfo: &mut ItemInfo,
 ) -> Result<()> {
     let (mut timestamp_scale, mut duration) = (None, None);
+    let mut seek_index = HashMap::new();
+
     while let Some(Ok(Unflat { children, item })) = children.next() {
         match item {
             MatroskaTag::SeekHead(_) => {}
@@ -70,57 +72,8 @@ fn import_read_segment(
                     }
                 }
             }
-            MatroskaTag::Cluster(_) => {
-                info!("start of cluster found");
-                let mut children = children.unwrap();
-                while let Some(Ok(Unflat { children, item })) = children.next() {
-                    match item {
-                        MatroskaTag::BlockGroup(_) => {
-                            debug!("group");
-                            let mut children = children.unwrap();
-                            while let Some(Ok(Unflat { children: _, item })) = children.next() {
-                                match item {
-                                    MatroskaTag::Block(buf) => {
-                                        let block = Block::parse(&buf)?;
-                                        debug!(
-                                            "block: track={} tso={}",
-                                            block.track, block.timestamp_off
-                                        )
-                                    }
-                                    _ => trace!("{item:?}"),
-                                }
-                            }
-                        }
-                        MatroskaTag::SimpleBlock(buf) => {
-                            let block = Block::parse(&buf)?;
-                            debug!(
-                                "simple block: track={} tso={}",
-                                block.track, block.timestamp_off
-                            )
-                        }
-                        _ => debug!("(rsc) tag ignored: {item:?}"),
-                    }
-                }
-            }
             MatroskaTag::Tags(_) => {}
-            MatroskaTag::Cues(_) => {
-                let mut children = children.unwrap();
-                while let Some(Ok(Unflat { children, item })) = children.next() {
-                    match item {
-                        MatroskaTag::CuePoint(_) => {
-                            let mut children = children.unwrap();
-                            while let Some(Ok(Unflat {
-                                children: _,
-                                item: _,
-                            })) = children.next()
-                            {
-                                // error!("{item:?}")
-                            }
-                        }
-                        _ => (),
-                    }
-                }
-            }
+            MatroskaTag::Cues(_) => {}
             MatroskaTag::Chapters(_) => {}
             MatroskaTag::Tracks(_) => {
                 let mut children = children.unwrap();
@@ -218,9 +171,72 @@ fn import_read_segment(
                     }
                 }
             }
+            MatroskaTag::Cluster(_) => {
+                let mut children = children.unwrap();
+                let mut pts = 0;
+
+                while let Some(Ok(Unflat { children, item })) = children.next() {
+                    match item {
+                        MatroskaTag::Timestamp(ts) => pts = ts,
+                        MatroskaTag::BlockGroup(_) => {
+                            debug!("group");
+                            let mut children = children.unwrap();
+                            let pos = children.position();
+                            while let Some(Ok(Unflat { children: _, item })) = children.next() {
+                                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: pos,
+                                                size: block.data.len(),
+                                            });
+                                    }
+                                    _ => trace!("{item:?}"),
+                                }
+                            }
+                        }
+                        MatroskaTag::SimpleBlock(buf) => {
+                            let block = Block::parse(&buf)?;
+                            debug!(
+                                "simple 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: 0,
+                                    size: block.data.len(),
+                                });
+                        }
+                        _ => debug!("(rsc) tag ignored: {item:?}"),
+                    }
+                }
+            }
+
             _ => debug!("(rs) tag ignored: {item:?}"),
         }
     }
+
+    for (tn, index) in seek_index {
+        bincode::encode_into_std_write(
+            index,
+            &mut File::create(path.with_extension(&format!("si.{}", tn)))?,
+            bincode::config::standard(),
+        )?;
+    }
+
     iteminfo.duration =
         (duration.unwrap() * timestamp_scale.unwrap() as f64) as f64 / 1_000_000_000 as f64;
     Ok(())