From 69eff1075923180fec5044ee1fd5f6e0456b9212 Mon Sep 17 00:00:00 2001
From: Emerson Knapp <537409+emersonknapp@users.noreply.github.com>
Date: Fri, 14 Oct 2022 17:15:04 -0700
Subject: [PATCH] Add set_read_order reader API (#54)

* Implement new set_read_order API for storage

Signed-off-by: Emerson Knapp <emerson.b.knapp@gmail.com>
Signed-off-by: James Smith <james@foxglove.dev>
---
 rosbag2_storage_mcap/CMakeLists.txt       |  1 +
 rosbag2_storage_mcap/src/mcap_storage.cpp | 32 +++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/rosbag2_storage_mcap/CMakeLists.txt b/rosbag2_storage_mcap/CMakeLists.txt
index 1a4f2498ef..71c1a49654 100644
--- a/rosbag2_storage_mcap/CMakeLists.txt
+++ b/rosbag2_storage_mcap/CMakeLists.txt
@@ -57,6 +57,7 @@ endif()
 # COMPATIBILITY(foxy, galactic, humble) - 0.17.x is the Rolling release.
 if(${rosbag2_storage_VERSION} VERSION_GREATER_EQUAL 0.17.0)
   list(APPEND MCAP_COMPILE_DEFS ROSBAG2_STORAGE_MCAP_HAS_STORAGE_FILTER_TOPIC_REGEX)
+  list(APPEND MCAP_COMPILE_DEFS ROSBAG2_STORAGE_MCAP_HAS_SET_READ_ORDER)
 endif()
 
 target_compile_definitions(${PROJECT_NAME} PRIVATE ${MCAP_COMPILE_DEFS})
diff --git a/rosbag2_storage_mcap/src/mcap_storage.cpp b/rosbag2_storage_mcap/src/mcap_storage.cpp
index af2b0129c3..23069231ef 100644
--- a/rosbag2_storage_mcap/src/mcap_storage.cpp
+++ b/rosbag2_storage_mcap/src/mcap_storage.cpp
@@ -174,6 +174,9 @@ class MCAPStorage : public rosbag2_storage::storage_interfaces::ReadWriteInterfa
   std::string get_storage_identifier() const override;
 
   /** BaseReadInterface **/
+#ifdef ROSBAG2_STORAGE_MCAP_HAS_SET_READ_ORDER
+  void set_read_order(const rosbag2_storage::ReadOrder&) override;
+#endif
   bool has_next() override;
   std::shared_ptr<rosbag2_storage::SerializedBagMessage> read_next() override;
   std::vector<rosbag2_storage::TopicMetadata> get_all_topics_and_types() override;
@@ -458,6 +461,35 @@ void MCAPStorage::ensure_summary_read() {
   }
 }
 
+#ifdef ROSBAG2_STORAGE_MCAP_HAS_SET_READ_ORDER
+void MCAPStorage::set_read_order(const rosbag2_storage::ReadOrder& read_order) {
+  auto next_read_order = read_order_;
+  switch (read_order.sort_by) {
+    case rosbag2_storage::ReadOrder::ReceivedTimestamp:
+      if (read_order.reverse) {
+        next_read_order = mcap::ReadMessageOptions::ReadOrder::ReverseLogTimeOrder;
+      } else {
+        next_read_order = mcap::ReadMessageOptions::ReadOrder::LogTimeOrder;
+      }
+      break;
+    case rosbag2_storage::ReadOrder::File:
+      if (!read_order.reverse) {
+        next_read_order = mcap::ReadMessageOptions::ReadOrder::FileOrder;
+      } else {
+        throw std::runtime_error("Reverse file order reading not implemented.");
+      }
+      break;
+    case rosbag2_storage::ReadOrder::PublishedTimestamp:
+      throw std::runtime_error("PublishedTimestamp read order not yet implemented in ROS 2");
+      break;
+  }
+  if (next_read_order != read_order_) {
+    read_order_ = next_read_order;
+    reset_iterator();
+  }
+}
+#endif
+
 bool MCAPStorage::has_next() {
   if (!linear_iterator_) {
     return false;