rsbag/examples/bag_info.rs

use std::{convert::TryFrom, env::args, fs::File, io};

use log::{error, info, trace};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use regex::Regex;
use ros_message::{MessagePath, Msg};
use rsbag::{
    index::{BagIndex, ConnInfo},
    reader::{BagReader, MmapReader},
};

fn parse_msgdef(message_name: &str, msgdef: &str) -> rsbag::Result<Msg> {
    trace!("message definition: {}", msgdef);
    let path = MessagePath::try_from(message_name).map_err(rsbag::Error::other)?;
    let msgtype = Msg::new(path, msgdef).map_err(rsbag::Error::other)?;
    Ok(msgtype)
}

fn parse_message_definitions(conn: &ConnInfo) -> rsbag::Result<Vec<Msg>> {
    let msgdefs = conn.message_definition()?;
    let boundary_re = Regex::new(r"\r?\n==+\r?\nMSG: ([^\r\n]+)\r?\n").unwrap();
    let mut name = conn.datatype()?;
    let mut begin = 0usize;
    let mut msgs = Vec::new();

    for cap in boundary_re.captures_iter(&msgdefs) {
        let boundary_range = cap.get(0).unwrap();
        let end = boundary_range.start();

        let msgdef = &msgdefs[begin..end];
        let msgtype = parse_msgdef(&name, msgdef)?;
        msgs.push(msgtype);

        name = cap[1].to_string();
        begin = boundary_range.end();
    }

    let msgdef = &msgdefs[begin..];

    let msg = parse_msgdef(&name, msgdef)?;
    msgs.push(msg);

    Ok(msgs)
}

fn read_chunk<R: BagReader>(bag_reader: &mut R, pos: u64) -> rsbag::Result<Vec<u8>> {
    let chunk_header = bag_reader.read_chunk_header(pos)?;

    let compressed_data = bag_reader.read_data().unwrap();
    let mut data = Vec::with_capacity(chunk_header.uncompressed_size as usize);
    let mut decompresor = chunk_header.compression.decompress(compressed_data);
    io::copy(&mut decompresor, &mut data).unwrap();

    Ok(data)
}

fn main() {
    env_logger::init();

    let args: Vec<_> = args().collect();
    if args.len() != 2 {
        eprintln!("Usage: {} <bag path>", args[0]);
        return;
    }

    let bag_path = &args[1];
    let bag_file = File::open(bag_path).expect("Could not open bag file");
    let mut bag_reader = MmapReader::new(bag_file).unwrap();

    let index = match BagIndex::read_all(&mut bag_reader) {
        Ok(index) => index,
        Err(err) => {
            error!("bag parse error: {}", err);
            return;
        }
    };

    for conn in &index.connections {
        match parse_message_definitions(conn) {
            Ok(msgs) => {
                for msg in &msgs {
                    trace!(
                        "message definition parsed: {:#?}",
                        msg.fields()
                            .iter()
                            .filter(|field| !field.is_constant())
                            .map(ToString::to_string)
                            .collect::<Vec<_>>()
                    );
                }
            }
            Err(err) => error!("could not parse message definition: {}", err),
        }
    }

    // let mut total_size = 0;

    // let total_size = index
    //     .chunks
    //     .par_iter()
    //     .try_fold(
    //         || 0u64,
    //         |total_size, chunk| -> rsbag::Result<_> {
    //             let chunk_header = bag_reader.clone().read_chunk_header(chunk.pos)?;

    //             // let data = read_chunk(&mut bag_reader.clone(), chunk.pos)?;
    //             // chunks.push(data);
    //             Ok(total_size + chunk_header.uncompressed_size as u64)
    //         },
    //     )
    //     .reduce(
    //         // || Ok(Vec::new()),
    //         || Ok(0),
    //         |a, b| a.and_then(|a| b.map(|b| a + b)),
    //     ).unwrap();

    let total_size = index
        .chunks
        .par_iter()
        .try_fold(
            || 0u64,
            |total_size, chunk| -> rsbag::Result<_> {
                let data = read_chunk(&mut bag_reader.clone(), chunk.pos)?;
                Ok(total_size + data.len() as u64)
            },
        )
        .reduce(
            // || Ok(Vec::new()),
            || Ok(0),
            |a, b| a.and_then(|a| b.map(|b| a + b)),
        ).unwrap();

    info!("total uncompressed size: {}", total_size);
}