re_mcap 0.31.3

use super::super::definitions::tf2_msgs::TFMessage;
use re_chunk::{Chunk, ChunkId};

use re_sdk_types::archetypes::Transform3D;
use re_sdk_types::components::{RotationQuat, Translation3D};
use re_sdk_types::datatypes::Quaternion;

use super::super::Ros2MessageParser;
use crate::parsers::ros2msg::definitions::geometry_msgs::{Transform, TransformStamped};
use crate::parsers::ros2msg::definitions::std_msgs::Header;
use crate::parsers::{
    cdr,
    decode::{MessageParser, ParserContext},
};
use crate::util::{TimestampCell, log_and_publish_timepoint_from_msg};

const STATIC_TF_TOPIC: &str = "/tf_static";

fn static_chunk_timelines()
-> re_chunk::external::nohash_hasher::IntMap<re_log_types::TimelineName, re_chunk::TimeColumn> {
    // Chunks without any timelines are treated as static by Rerun.
    re_chunk::external::nohash_hasher::IntMap::default()
}

pub struct TfMessageParser {
    translations: Vec<Translation3D>,
    quaternions: Vec<RotationQuat>,
    parent_frame_ids: Vec<String>,
    child_frame_ids: Vec<String>,
}

impl Ros2MessageParser for TfMessageParser {
    fn new(_num_rows: usize) -> Self {
        // Note that we can't know the number of output rows in advance,
        // as each message can contain a variable amount of transforms.
        Self {
            translations: Vec::new(),
            quaternions: Vec::new(),
            parent_frame_ids: Vec::new(),
            child_frame_ids: Vec::new(),
        }
    }
}

impl MessageParser for TfMessageParser {
    fn get_log_and_publish_timepoints(
        &self,
        msg: &mcap::Message<'_>,
        time_type: re_log_types::TimeType,
    ) -> anyhow::Result<Vec<re_chunk::TimePoint>> {
        // We need a custom implementation of this method because we have a 1-to-N relationship between input messages and output rows.
        // Assign each output row the same log and publish time as the input message.
        let TFMessage { transforms } = cdr::try_decode_message::<TFMessage>(&msg.data)?;
        Ok(vec![
            log_and_publish_timepoint_from_msg(msg, time_type);
            transforms.len()
        ])
    }

    fn append(&mut self, ctx: &mut ParserContext, msg: &mcap::Message<'_>) -> anyhow::Result<()> {
        re_tracing::profile_function!();
        let TFMessage { transforms } = cdr::try_decode_message::<TFMessage>(&msg.data)?;

        // Each transform in the message has its own timestamp.
        for TransformStamped {
            header,
            child_frame_id,
            transform,
        } in transforms
        {
            // Add the header timestamp to the context.
            // `log_time` and `publish_time` are added via `log_and_publish_time_from_msg`.
            let Header { stamp, frame_id } = header;
            ctx.add_timestamp_cell(TimestampCell::from_nanos_ros2(
                stamp.as_nanos() as u64,
                ctx.time_type(),
            ));

            self.parent_frame_ids.push(frame_id);
            self.child_frame_ids.push(child_frame_id);

            let Transform {
                translation,
                rotation,
            } = transform;
            self.translations.push(Translation3D::new(
                translation.x as f32,
                translation.y as f32,
                translation.z as f32,
            ));
            self.quaternions.push(
                Quaternion::from_xyzw([
                    rotation.x as f32,
                    rotation.y as f32,
                    rotation.z as f32,
                    rotation.w as f32,
                ])
                .into(),
            );
        }

        Ok(())
    }

    fn finalize(self: Box<Self>, ctx: ParserContext) -> anyhow::Result<Vec<re_chunk::Chunk>> {
        re_tracing::profile_function!();
        let Self {
            translations,
            quaternions,
            parent_frame_ids,
            child_frame_ids,
        } = *self;

        let entity_path = ctx.entity_path().clone();
        let timelines = if ctx.channel_topic() == STATIC_TF_TOPIC {
            static_chunk_timelines()
        } else {
            ctx.build_timelines()
        };

        let chunk = Chunk::from_auto_row_ids(
            ChunkId::new(),
            entity_path.clone(),
            timelines.clone(),
            Transform3D::update_fields()
                .with_many_translation(translations)
                .with_many_quaternion(quaternions)
                .with_many_child_frame(child_frame_ids)
                .with_many_parent_frame(parent_frame_ids)
                .columns_of_unit_batches()?
                .collect(),
        )?;

        Ok(vec![chunk])
    }
}

#[cfg(test)]
mod tests {
    use re_chunk::TimePoint;
    use re_log_types::{TimeCell, TimeType, TimelineName};

    use super::*;

    fn test_parser() -> TfMessageParser {
        TfMessageParser {
            translations: vec![Translation3D::new(1.0, 2.0, 3.0)],
            quaternions: vec![Quaternion::from_xyzw([0.0, 0.0, 0.0, 1.0]).into()],
            parent_frame_ids: vec!["parent".to_owned()],
            child_frame_ids: vec!["child".to_owned()],
        }
    }

    #[test]
    fn tf_static_topic_produces_static_chunk() {
        let ctx = ParserContext::new("/tf_static".into(), STATIC_TF_TOPIC, TimeType::TimestampNs);
        let chunk = Box::new(test_parser()).finalize(ctx).unwrap().remove(0);

        assert!(chunk.is_static());
    }

    #[test]
    fn non_tf_static_topic_stays_temporal() {
        let mut ctx = ParserContext::new("/tf".into(), "tf", TimeType::TimestampNs);
        ctx.add_timepoint(TimePoint::from([(
            TimelineName::log_time(),
            TimeCell::from_timestamp_nanos_since_epoch(123),
        )]));

        let chunk = Box::new(test_parser()).finalize(ctx).unwrap().remove(0);

        assert!(!chunk.is_static());
    }
}