DPL: refactor InputSpan

Now that navigating the MessageSet does not have an index anymore
it does not make sense to use a random accessor to implement the
InputSpan::Iterator, because that results in an O(N^2) loop for
normal linear iterations. It is much better to simply implement
a proper advancement operator, so that looping stays O(N).

This resumes the performance when iterating on old-style multiparts
which are apparently still used in some cases.

General modernization of InputSpan also done.

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;

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()

Framework/Core/include/Framework/InputSpan.h

@@ -14,8 +14,8 @@
 #include "Framework/DataRef.h"
 #include <functional>
 
-extern template class std::function<o2::framework::DataRef(size_t)>;
-extern template class std::function<o2::framework::DataRef(size_t, size_t)>;
+extern template class std::function<o2::framework::DataRef(size_t, o2::framework::DataRefIndices)>;
+extern template class std::function<o2::framework::DataRefIndices(size_t, o2::framework::DataRefIndices)>;
 
 namespace o2::framework
 {
@@ -32,37 +32,48 @@ class InputSpan
   InputSpan(InputSpan const&) = delete;
   InputSpan(InputSpan&&) = default;
 
-  /// @a getter is the mapping between an element of the span referred by
-  /// index and the buffer associated.
-  /// @a size is the number of elements in the span.
-  InputSpan(std::function<DataRef(size_t)> getter, size_t size);
+  /// Navigate the message store via the DataRefIndices protocol.
+  /// get_next_pair (DataModelViews.h) provides O(1) sequential advancement for nextIndicesGetter.
+  InputSpan(std::function<size_t(size_t)> nofPartsGetter,
+            std::function<int(size_t)> refCountGetter,
+            std::function<DataRef(size_t, DataRefIndices)> indicesGetter,
+            std::function<DataRefIndices(size_t, DataRefIndices)> nextIndicesGetter,
+            size_t size);
 
-  /// @a getter is the mapping between an element of the span referred by
-  /// index and the buffer associated.
-  /// @a size is the number of elements in the span.
-  InputSpan(std::function<DataRef(size_t, size_t)> getter, size_t size);
+  /// @a i-th element of the InputSpan (O(partidx) sequential scan via indices protocol)
+  [[nodiscard]] DataRef get(size_t i, size_t partidx = 0) const
+  {
+    DataRefIndices idx{0, 1};
+    for (size_t p = 0; p < partidx; ++p) {
+      idx = mNextIndicesGetter(i, idx);
+    }
+    return mIndicesGetter(i, idx);
+  }
 
-  /// @a getter is the mapping between an element of the span referred by
-  /// index and the buffer associated.
-  /// @nofPartsGetter is the getter for the number of parts associated with an index
-  /// @a size is the number of elements in the span.
-  InputSpan(std::function<DataRef(size_t, size_t)> getter, std::function<size_t(size_t)> nofPartsGetter, std::function<int(size_t)> refCountGetter, size_t size);
+  /// Return the DataRef for the part described by @a indices in slot @a slotIdx in O(1).
+  [[nodiscard]] DataRef getAtIndices(size_t slotIdx, DataRefIndices indices) const
+  {
+    return mIndicesGetter(slotIdx, indices);
+  }
 
-  /// @a i-th element of the InputSpan
-  [[nodiscard]] DataRef get(size_t i, size_t partidx = 0) const
+  /// Advance from @a current to the indices of the next part in slot @a slotIdx in O(1).
+  [[nodiscard]] DataRefIndices nextIndices(size_t slotIdx, DataRefIndices current) const
   {
-    return mGetter(i, partidx);
+    return mNextIndicesGetter(slotIdx, current);
   }
 
+  // --- slot-level Iterator protocol (headerIdx doubles as slot position) ---
+  [[nodiscard]] DataRefIndices initialIndices() const { return {0, 0}; }
+  [[nodiscard]] DataRefIndices endIndices() const { return {mSize, 0}; }
+  [[nodiscard]] DataRef getAtIndices(DataRefIndices indices) const { return mIndicesGetter(indices.headerIdx, {0, 1}); }
+  [[nodiscard]] DataRefIndices nextIndices(DataRefIndices current) const { return {current.headerIdx + 1, 0}; }
+
   /// @a number of parts in the i-th element of the InputSpan
   [[nodiscard]] size_t getNofParts(size_t i) const
   {
     if (i >= mSize) {
       return 0;
     }
-    if (!mNofPartsGetter) {
-      return 1;
-    }
     return mNofPartsGetter(i);
   }
 
@@ -94,7 +105,8 @@ class InputSpan
     return get(i).payload;
   }
 
-  /// an iterator class working on position within the a parent class
+  /// An iterator over the elements of a parent container using the DataRefIndices protocol.
+  /// ParentT must provide: initialIndices(), getAtIndices(DataRefIndices), nextIndices(DataRefIndices).
   template <typename ParentT, typename T>
   class Iterator
   {
@@ -110,23 +122,23 @@ class InputSpan
 
     Iterator() = delete;
 
-    Iterator(ParentType const* parent, size_t position = 0, size_t size = 0)
-      : mPosition(position), mSize(size > position ? size : position), mParent(parent), mElement{}
+    Iterator(ParentType const* parent, bool isEnd = false)
+      : mParent(parent),
+        mCurrentIndices(isEnd ? parent->endIndices() : parent->initialIndices()),
+        mElement{}
     {
-      if (mPosition < mSize) {
-        mElement = mParent->get(mPosition);
+      if (mCurrentIndices != mParent->endIndices()) {
+        mElement = mParent->getAtIndices(mCurrentIndices);
       }
     }
 
-    ~Iterator() = default;
-
     // prefix increment
     SelfType& operator++()
     {
-      if (mPosition < mSize && ++mPosition < mSize) {
-        mElement = mParent->get(mPosition);
+      mCurrentIndices = mParent->nextIndices(mCurrentIndices);
+      if (mCurrentIndices != mParent->endIndices()) {
+        mElement = mParent->getAtIndices(mCurrentIndices);
       } else {
-        // reset the element to the default value of the type
         mElement = ElementType{};
       }
       return *this;
@@ -145,16 +157,14 @@ class InputSpan
       return mElement;
     }
 
-    // comparison
     bool operator==(const SelfType& rh) const
     {
-      return mPosition == rh.mPosition;
+      return mCurrentIndices == rh.mCurrentIndices;
     }
 
-    // comparison
-    bool operator!=(const SelfType& rh) const
+    auto operator<=>(const SelfType& rh) const
     {
-      return mPosition != rh.mPosition;
+      return mCurrentIndices <=> rh.mCurrentIndices;
     }
 
     // return pointer to parent instance
@@ -163,22 +173,21 @@ class InputSpan
       return mParent;
     }
 
-    // return current position
+    // return current position (headerIdx serves as the slot index for slot-level iteration)
     [[nodiscard]] size_t position() const
     {
-      return mPosition;
+      return mCurrentIndices.headerIdx;
     }
 
    private:
-    size_t mPosition;
-    size_t mSize;
     ParentType const* mParent;
+    DataRefIndices mCurrentIndices;
     ElementType mElement;
   };
 
   /// @class InputSpanIterator
-  /// An iterator over the input slots
-  /// It supports an iterator interface to access the parts in the slot
+  /// An iterator over the input slots.
+  /// It supports an iterator interface to access the parts in the slot.
   template <typename T>
   class InputSpanIterator : public Iterator<InputSpan, T>
   {
@@ -192,24 +201,26 @@ class InputSpan
     using iterator = Iterator<SelfType, T>;
     using const_iterator = Iterator<SelfType, const T>;
 
-    InputSpanIterator(InputSpan const* parent, size_t position = 0, size_t size = 0)
-      : BaseType(parent, position, size)
+    InputSpanIterator(InputSpan const* parent, bool isEnd = false)
+      : BaseType(parent, isEnd)
     {
     }
 
-    /// Get element at {slotindex, partindex}
-    [[nodiscard]] ElementType get(size_t pos) const
+    /// 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)}; }
+
+    /// Get element at the given raw message indices in O(1).
+    [[nodiscard]] ElementType getAtIndices(DataRefIndices indices) const
     {
-      return this->parent()->get(this->position(), pos);
+      return this->parent()->getAtIndices(this->position(), indices);
     }
 
-    /// Check if slot is valid, index of part is not used
-    [[nodiscard]] bool isValid(size_t = 0) const
+    /// Advance @a current to the next part's indices in O(1).
+    [[nodiscard]] DataRefIndices nextIndices(DataRefIndices current) const
     {
-      if (this->position() < this->parent()->size()) {
-        return this->parent()->isValid(this->position());
-      }
-      return false;
+      return this->parent()->nextIndices(this->position(), current);
     }
 
     /// Get number of parts in input slot
@@ -218,15 +229,14 @@ class InputSpan
       return this->parent()->getNofParts(this->position());
     }
 
-    // iterator for the part access
     [[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);
     }
   };
 
@@ -236,19 +246,19 @@ class InputSpan
   // supporting read-only access and returning const_iterator
   [[nodiscard]] const_iterator begin() const
   {
-    return {this, 0, size()};
+    return {this, false};
   }
 
-  // supporting read-only access and returning const_iterator
   [[nodiscard]] const_iterator end() const
   {
-    return {this, size()};
+    return {this, true};
   }
 
  private:
-  std::function<DataRef(size_t, size_t)> mGetter;
   std::function<size_t(size_t)> mNofPartsGetter;
   std::function<int(size_t)> mRefCountGetter;
+  std::function<DataRef(size_t, DataRefIndices)> mIndicesGetter;
+  std::function<DataRefIndices(size_t, DataRefIndices)> mNextIndicesGetter;
   size_t mSize;
 };