Merge 75221fb into sapling-pr-archive-ktf

Author: ktf

Framework/Core/include/Framework/DataModelViews.h

@@ -16,6 +16,7 @@
 #include "DomainInfoHeader.h"
 #include "SourceInfoHeader.h"
 #include "Headers/DataHeader.h"
+#include "Framework/DataRef.h"
 #include "Framework/TimesliceSlot.h"
 #include <ranges>
 #include <span>
@@ -80,10 +81,7 @@ struct count_parts {
   }
 };
 
-struct DataRefIndices {
-  size_t headerIdx;
-  size_t payloadIdx;
-};
+// DataRefIndices is defined in Framework/DataRef.h
 
 struct get_pair {
   size_t pairId;
@@ -127,6 +125,43 @@ struct get_pair {
   }
 };
 
+// Advance from a DataRefIndices to the next one in O(1), reading only the
+// current header.  Intended for use in iterators so that ++ is O(1) rather
+// than the O(n) while-loop that get_pair requires.
+//
+// New-style block  (splitPayloadIndex == splitPayloadParts > 1):
+//   layout: [header, payload_0, payload_1, ..., payload_{N-1}]
+//   advance within block while payloads remain, then jump to the next block.
+//
+// Old-style block  (splitPayloadIndex != splitPayloadParts, splitPayloadParts > 1)
+// or single pair   (splitPayloadParts == 0):
+//   layout: [header, payload]  – always advance by two messages.
+struct get_next_pair {
+  DataRefIndices current;
+  template <typename R>
+    requires std::ranges::random_access_range<R> && std::ranges::sized_range<R>
+  friend DataRefIndices operator|(R&& r, get_next_pair self)
+  {
+    size_t hIdx = self.current.headerIdx;
+    auto* header = o2::header::get<o2::header::DataHeader*>(r[hIdx]->GetData());
+    if (!header) {
+      throw std::runtime_error("Not a DataHeader");
+    }
+    if (header->splitPayloadParts > 1 && header->splitPayloadIndex == header->splitPayloadParts) {
+      // New-style block: one header followed by splitPayloadParts contiguous payloads.
+      if (self.current.payloadIdx < hIdx + header->splitPayloadParts) {
+        // More sub-payloads remain in this block.
+        return {hIdx, self.current.payloadIdx + 1};
+      }
+      // Last sub-payload consumed; move to the first pair of the next block.
+      size_t nextHIdx = hIdx + header->splitPayloadParts + 1;
+      return {nextHIdx, nextHIdx + 1};
+    }
+    // Old-style [header, payload] pairs or a single pair: advance by two messages.
+    return {hIdx + 2, hIdx + 3};
+  }
+};
+
 struct get_dataref_indices {
   size_t part;
   size_t subPart;

Framework/Core/include/Framework/DataRef.h

@@ -12,6 +12,7 @@
 #define FRAMEWORK_DATAREF_H
 
 #include <cstddef> // for size_t
+#include <compare>
 
 namespace o2
 {
@@ -29,6 +30,15 @@ struct DataRef {
   size_t payloadSize = 0;
 };
 
+/// Raw indices into the message vector for one (header, payload) pair.
+/// Kept in a lightweight header so InputSpan can use it without pulling in FairMQ.
+struct DataRefIndices {
+  size_t headerIdx;
+  size_t payloadIdx;
+  bool operator==(const DataRefIndices&) const = default;
+  auto operator<=>(const DataRefIndices&) const = default;
+};
+
 } // namespace framework
 } // namespace o2
 

Framework/Core/include/Framework/InputRecord.h

@@ -13,6 +13,7 @@
 
 #include "Framework/DataRef.h"
 #include "Framework/DataRefUtils.h"
+#include "Framework/InputSpan.h"
 #include "Framework/InputRoute.h"
 #include "Framework/TypeTraits.h"
 #include "Framework/TableConsumer.h"
@@ -202,6 +203,15 @@ class InputRecord
 
   [[nodiscard]] size_t getNofParts(int pos) const;
 
+  /// O(1) access to the part described by @a indices in slot @a pos.
+  [[nodiscard]] DataRef getAtIndices(int pos, DataRefIndices indices) const;
+
+  /// O(1) advance from @a current to the next part's indices in slot @a pos.
+  [[nodiscard]] DataRefIndices nextIndices(int pos, DataRefIndices current) const
+  {
+    return mSpan.nextIndices(pos, current);
+  }
+
   // Given a binding by string, return the associated DataRef
   DataRef getDataRefByString(const char* bindingName, int part = 0) const
   {
@@ -568,8 +578,8 @@ class InputRecord
 
     Iterator() = delete;
 
-    Iterator(ParentType const* parent, size_t position = 0, size_t size = 0)
-      : mPosition(position), mSize(size > position ? size : position), mParent(parent), mElement{nullptr, nullptr, nullptr}
+    Iterator(ParentType const* parent, bool isEnd = false)
+      : mPosition(isEnd ? parent->size() : 0), mSize(parent->size()), mParent(parent), mElement{nullptr, nullptr, nullptr}
     {
       if (mPosition < mSize) {
         if (mParent->isValid(mPosition)) {
@@ -678,18 +688,29 @@ class InputRecord
     using reference = typename BaseType::reference;
     using pointer = typename BaseType::pointer;
     using ElementType = typename std::remove_const<value_type>::type;
-    using iterator = Iterator<SelfType, T>;
-    using const_iterator = Iterator<SelfType, const T>;
+    using iterator = InputSpan::Iterator<SelfType, T>;
+    using const_iterator = InputSpan::Iterator<SelfType, const T>;
+
+    InputRecordIterator(InputRecord const* parent, bool isEnd = false)
+      : BaseType(parent, isEnd)
+    {
+    }
+
+    /// Initial indices for part-level iteration: first part starts at {headerIdx=0, payloadIdx=1}.
+    [[nodiscard]] DataRefIndices initialIndices() const { return {0, 1}; }
+    /// Sentinel used by nextIndicesGetter to signal end-of-slot.
+    [[nodiscard]] DataRefIndices endIndices() const { return {size_t(-1), size_t(-1)}; }
 
-    InputRecordIterator(InputRecord const* parent, size_t position = 0, size_t size = 0)
-      : BaseType(parent, position, size)
+    /// Get element at the given raw message indices in O(1).
+    [[nodiscard]] ElementType getAtIndices(DataRefIndices indices) const
     {
+      return this->parent()->getAtIndices(this->position(), indices);
     }
 
-    /// Get element at {slotindex, partindex}
-    [[nodiscard]] ElementType getByPos(size_t pos) const
+    /// Advance @a current to the next part's indices in O(1).
+    [[nodiscard]] DataRefIndices nextIndices(DataRefIndices current) const
     {
-      return this->parent()->getByPos(this->position(), pos);
+      return this->parent()->nextIndices(this->position(), current);
     }
 
     /// Check if slot is valid, index of part is not used
@@ -709,12 +730,12 @@ class InputRecord
 
     [[nodiscard]] const_iterator begin() const
     {
-      return const_iterator(this, 0, size());
+      return const_iterator(this, size() == 0);
     }
 
     [[nodiscard]] const_iterator end() const
     {
-      return const_iterator(this, size());
+      return const_iterator(this, true);
     }
   };
 
@@ -723,12 +744,12 @@ class InputRecord
 
   [[nodiscard]] const_iterator begin() const
   {
-    return {this, 0, size()};
+    return {this, false};
   }
 
   [[nodiscard]] const_iterator end() const
   {
-    return {this, size()};
+    return {this, true};
   }
 
   InputSpan& span()