feat: Add (optional) forward linking of track states (#2418)

A track can now initiate forward linking of its track states. The track
states have a new property, `next` that optionally makes the sequence a
doubly linked list. This is not filled by default (not even by
`appendTrackState`) because for branching cases, it's not unambiguous.

`linkForward()` on a track will initiate the forward linking and store
the innermost track state index for easy access. `reverseTrackStates`
also does this, since it's only valid for non-branching trackstate
sequences anyway.

Also adds a `innermostTrackState()` getter, which returns a track state proxy in case the track is forward linked and has a `stemIndex`

Blocked by:
- #2425 
- #2426

Core/include/Acts/EventData/MultiTrajectory.hpp

@@ -971,9 +971,11 @@ class TrackStateProxy {
 };
 
 /// Helper type that wraps two iterators
-template <typename trajectory_t, size_t M, bool ReadOnly>
+template <bool reverse, typename trajectory_t, size_t M, bool ReadOnly>
 class TrackStateRange {
   using ProxyType = TrackStateProxy<trajectory_t, M, ReadOnly>;
+  using IndexType = typename ProxyType::IndexType;
+  static constexpr IndexType kInvalid = ProxyType::kInvalid;
 
  public:
   /// Iterator that wraps a track state proxy. The nullopt case signifies the
@@ -991,12 +993,24 @@ class TrackStateRange {
       if (!proxy) {
         return *this;
       }
-      if (proxy->hasPrevious()) {
-        proxy = proxy->trajectory().getTrackState(proxy->previous());
-        return *this;
+      if constexpr (reverse) {
+        if (proxy->hasPrevious()) {
+          proxy = proxy->trajectory().getTrackState(proxy->previous());
+          return *this;
+        } else {
+          proxy = std::nullopt;
+          return *this;
+        }
       } else {
-        proxy = std::nullopt;
-        return *this;
+        IndexType next =
+            proxy->template component<IndexType, hashString("next")>();
+        if (next != kInvalid) {
+          proxy = proxy->trajectory().getTrackState(next);
+          return *this;
+        } else {
+          proxy = std::nullopt;
+          return *this;
+        }
       }
     }
 
@@ -1134,29 +1148,62 @@ class MultiTrajectory {
   /// @note Const version
   auto reverseTrackStateRange(IndexType iendpoint) const {
     using range_t =
-        decltype(detail_lt::TrackStateRange{getTrackState(iendpoint)});
+        detail_lt::TrackStateRange<true, Derived, MeasurementSizeMax, true>;
     if (iendpoint == kInvalid) {
       return range_t{};
     }
 
     return range_t{getTrackState(iendpoint)};
   }
 
-  /// Range for the track states from @p iendpoint to the trajectory start
+  /// Range for the track states from @p iendpoint to the trajectory start,
+  /// i.e from the outside in.
   /// @param iendpoint Trajectory entry point to start from
   /// @return Iterator pair to iterate over
   /// @note Mutable version
   template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
   auto reverseTrackStateRange(IndexType iendpoint) {
     using range_t =
-        decltype(detail_lt::TrackStateRange{getTrackState(iendpoint)});
+        detail_lt::TrackStateRange<true, Derived, MeasurementSizeMax, false>;
     if (iendpoint == kInvalid) {
       return range_t{};
     }
 
     return range_t{getTrackState(iendpoint)};
   }
 
+  /// Range for the track states from @p istartpoint to the trajectory end,
+  /// i.e from inside out
+  /// @param istartpoint Trajectory state index for the innermost track
+  ///        state to start from
+  /// @return Iterator pair to iterate over
+  /// @note Const version
+  auto forwardTrackStateRange(IndexType istartpoint) const {
+    using range_t =
+        detail_lt::TrackStateRange<false, Derived, MeasurementSizeMax, true>;
+    if (istartpoint == kInvalid) {
+      return range_t{};
+    }
+
+    return range_t{getTrackState(istartpoint)};
+  }
+
+  /// Range for the track states from @p istartpoint to the trajectory end,
+  /// i.e from inside out
+  /// @param istartpoint Trajectory state index for the innermost track
+  ///        state to start from
+  /// @return Iterator pair to iterate over
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  auto forwardTrackStateRange(IndexType istartpoint) {
+    using range_t =
+        detail_lt::TrackStateRange<false, Derived, MeasurementSizeMax, false>;
+    if (istartpoint == kInvalid) {
+      return range_t{};
+    }
+
+    return range_t{getTrackState(istartpoint)};
+  }
+
   /// Apply a function to all previous states starting at a given endpoint.
   ///
   /// @param iendpoint  index of the last state

Core/include/Acts/EventData/TrackContainer.hpp

@@ -295,6 +295,23 @@ class TrackContainer {
     return m_traj->reverseTrackStateRange(tip);
   }
 
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  auto forwardTrackStateRange(IndexType itrack) {
+    auto stem = component<IndexType, hashString("stemIndex")>(itrack);
+    if (stem == kInvalid) {
+      throw std::invalid_argument{"Track has no stem index"};
+    }
+    return m_traj->forwardTrackStateRange(stem);
+  }
+
+  auto forwardTrackStateRange(IndexType itrack) const {
+    auto stem = component<IndexType, hashString("stemIndex")>(itrack);
+    if (stem == kInvalid) {
+      throw std::invalid_argument{"Track has no stem index"};
+    }
+    return m_traj->forwardTrackStateRange(stem);
+  }
+
  private:
   template <typename T, bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
   void copyDynamicFrom(IndexType dstIdx, const T& src, IndexType srcIdx) {

Core/include/Acts/EventData/TrackProxy.hpp

@@ -260,6 +260,14 @@ class TrackProxy {
     return component<IndexType>(hashString("tipIndex"));
   }
 
+  /// Index of the stem, i.e. the innermost track state of the track.
+  /// This might be invalid, signifying that the track state is not
+  /// forward-linked.
+  /// @return the stem index
+  IndexType stemIndex() const {
+    return component<IndexType>(hashString("stemIndex"));
+  }
+
   /// Get a mutable reference to the tip index, i.e. the entry point into the
   /// track container
   /// @return mutable reference to the tip index
@@ -268,6 +276,48 @@ class TrackProxy {
     return component<IndexType>(hashString("tipIndex"));
   }
 
+  /// Index of the stem, i.e. the innermost track state of the track.
+  /// This might be invalid, signifying that the track state is not
+  /// forward-linked.
+  /// @return mutable reference to the stem index
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  IndexType& stemIndex() {
+    return component<IndexType>(hashString("stemIndex"));
+  }
+
+  /// Return a const track state proxy to the innermost track state
+  /// @note This is only available, if the track is forward linked
+  /// @return The innermost track state proxy
+  auto innermostTrackState() const {
+    using proxy_t = decltype(m_container->trackStateContainer().getTrackState(
+        std::declval<IndexType>()));
+
+    IndexType stem = component<IndexType>(hashString("stemIndex"));
+    if (stem == kInvalid) {
+      return std::optional<proxy_t>{};
+    } else {
+      return std::optional<proxy_t>{
+          m_container->trackStateContainer().getTrackState(stem)};
+    }
+  }
+
+  /// Return a mutable track state proxy to the innermost track state
+  /// @note This is only available, if the track is forward linked
+  /// @return The innermost track state proxy
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  auto innermostTrackState() {
+    using proxy_t = decltype(m_container->trackStateContainer().getTrackState(
+        std::declval<IndexType>()));
+
+    IndexType stem = component<IndexType>(hashString("stemIndex"));
+    if (stem == kInvalid) {
+      return std::optional<proxy_t>{};
+    } else {
+      return std::optional<proxy_t>{
+          m_container->trackStateContainer().getTrackState(stem)};
+    }
+  }
+
   /// Get the reference surface of the track (e.g. the perigee)
   /// @return the reference surface
   const Surface& referenceSurface() const {
@@ -404,19 +454,50 @@ class TrackProxy {
 
   /// Get a range over the track states of this track. Return value is
   /// compatible with range based for loop. Const version
+  /// @note This range is from the outside inwards!
   /// @return Track state range to iterate over
   auto trackStatesReversed() const {
     return m_container->reverseTrackStateRange(m_index);
   }
 
   /// Get a range over the track states of this track. Return value is
   /// compatible with range based for loop. Mutable version
+  /// @note This range is from the outside inwards!
   /// @return Track state range to iterate over
   template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
   auto trackStatesReversed() {
     return m_container->reverseTrackStateRange(m_index);
   }
 
+  /// Get a range over the track states of this track. Return value is
+  /// compatible with range based for loop. Const version
+  /// @note This range is from the inside out!
+  /// @return Track state range to iterate over
+  auto trackStates() const {
+    return m_container->forwardTrackStateRange(m_index);
+  }
+
+  /// Get a range over the track states of this track. Return value is
+  /// compatible with range based for loop. Mutable version
+  /// @note This range is from the inside out!
+  /// @return Track state range to iterate over
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  auto trackStates() {
+    return m_container->forwardTrackStateRange(m_index);
+  }
+
+  /// Forward connect a track, i.e. set indices from the inside out
+  /// on all track states.
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  void linkForward() {
+    IndexType last = kInvalid;
+    for (auto ts : trackStatesReversed()) {
+      ts.template component<IndexType>(hashString("next")) = last;
+      last = ts.index();
+    }
+    stemIndex() = last;
+  }
+
   /// Append a track state to this track. This will modify the tip index to
   /// point at the newly created track state, which will be directly after the
   /// previous track state at tip index.
@@ -591,15 +672,19 @@ class TrackProxy {
   /// Afterwards, the previous endpoint of the track state sequence will be the
   /// "innermost" track state
   /// @note This is dangerous with branching track state sequences, as it will break them
+  /// @note This also automatically forward-links the track!
   template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
   void reverseTrackStates() {
     IndexType current = tipIndex();
     IndexType next = kInvalid;
     IndexType prev = kInvalid;
 
+    stemIndex() = tipIndex();
+
     while (current != kInvalid) {
       auto ts = m_container->trackStateContainer().getTrackState(current);
       prev = ts.previous();
+      ts.template component<IndexType>(hashString("next")) = prev;
       ts.previous() = next;
       next = current;
       tipIndex() = current;

Core/include/Acts/EventData/VectorMultiTrajectory.hpp

@@ -167,6 +167,7 @@ class VectorMultiTrajectoryBase {
   VectorMultiTrajectoryBase(const VectorMultiTrajectoryBase& other)
       : m_index{other.m_index},
         m_previous{other.m_previous},
+        m_next{other.m_next},
         m_params{other.m_params},
         m_cov{other.m_cov},
         m_meas{other.m_meas},
@@ -208,6 +209,7 @@ class VectorMultiTrajectoryBase {
         return instance.m_sourceLinks[instance.m_index[istate].iuncalibrated]
             .has_value();
       case "previous"_hash:
+      case "next"_hash:
       case "referenceSurface"_hash:
       case "measdim"_hash:
       case "chi2"_hash:
@@ -230,6 +232,8 @@ class VectorMultiTrajectoryBase {
     switch (key) {
       case "previous"_hash:
         return &instance.m_previous[istate];
+      case "next"_hash:
+        return &instance.m_next[istate];
       case "predicted"_hash:
         return &instance.m_index[istate].ipredicted;
       case "filtered"_hash:
@@ -271,6 +275,7 @@ class VectorMultiTrajectoryBase {
       case "jacobian"_hash:
       case "projector"_hash:
       case "previous"_hash:
+      case "next"_hash:
       case "uncalibratedSourceLink"_hash:
       case "referenceSurface"_hash:
       case "measdim"_hash:
@@ -302,6 +307,7 @@ class VectorMultiTrajectoryBase {
   /// index to map track states to the corresponding
   std::vector<IndexData> m_index;
   std::vector<IndexType> m_previous;
+  std::vector<IndexType> m_next;
   std::vector<typename detail_lt::Types<eBoundSize>::Coefficients> m_params;
   std::vector<typename detail_lt::Types<eBoundSize>::Covariance> m_cov;
 

Core/include/Acts/EventData/VectorTrackContainer.hpp

@@ -74,6 +74,8 @@ class VectorTrackContainerBase {
     switch (key) {
       case "tipIndex"_hash:
         return &instance.m_tipIndex[itrack];
+      case "stemIndex"_hash:
+        return &instance.m_stemIndex[itrack];
       case "particleHypothesis"_hash:
         return &instance.m_particleHypothesis[itrack];
       case "params"_hash:
@@ -113,6 +115,8 @@ class VectorTrackContainerBase {
     bool result = true;
     result = result && m_tipIndex.size() == size;
     assert(result);
+    result = result && m_stemIndex.size() == size;
+    assert(result);
     result = result && m_particleHypothesis.size() == size;
     assert(result);
     result = result && m_params.size() == size;
@@ -161,6 +165,7 @@ class VectorTrackContainerBase {
   // END INTERFACE HELPER
 
   std::vector<IndexType> m_tipIndex;
+  std::vector<IndexType> m_stemIndex;
   std::vector<ParticleHypothesis> m_particleHypothesis;
   std::vector<typename detail_lt::Types<eBoundSize>::Coefficients> m_params;
   std::vector<typename detail_lt::Types<eBoundSize>::Covariance> m_cov;

Core/src/EventData/VectorMultiTrajectory.cpp

@@ -31,6 +31,7 @@ auto VectorMultiTrajectory::addTrackState_impl(TrackStatePropMask mask,
   IndexData& p = m_index.back();
   IndexType index = m_index.size() - 1;
   m_previous.emplace_back(iprevious);
+  m_next.emplace_back(kInvalid);
 
   p.allocMask = mask;
 
@@ -169,6 +170,7 @@ void VectorMultiTrajectory::unset_impl(TrackStatePropMask target,
 void VectorMultiTrajectory::clear_impl() {
   m_index.clear();
   m_previous.clear();
+  m_next.clear();
   m_params.clear();
   m_cov.clear();
   m_meas.clear();
@@ -225,6 +227,7 @@ void detail_vmt::VectorMultiTrajectoryBase::Statistics::toStream(
 void VectorMultiTrajectory::reserve(std::size_t n) {
   m_index.reserve(n);
   m_previous.reserve(n);
+  m_next.reserve(n);
   m_params.reserve(n * 2);
   m_cov.reserve(n * 2);
   m_meas.reserve(n * 2);

Core/src/EventData/VectorTrackContainer.cpp

@@ -19,6 +19,7 @@ namespace detail_vtc {
 VectorTrackContainerBase::VectorTrackContainerBase(
     const VectorTrackContainerBase& other)
     : m_tipIndex{other.m_tipIndex},
+      m_stemIndex{other.m_stemIndex},
       m_particleHypothesis{other.m_particleHypothesis},
       m_params{other.m_params},
       m_cov{other.m_cov},
@@ -40,6 +41,7 @@ VectorTrackContainer::IndexType VectorTrackContainer::addTrack_impl() {
   assert(checkConsistency());
 
   m_tipIndex.emplace_back(kInvalid);
+  m_stemIndex.emplace_back(kInvalid);
 
   m_particleHypothesis.emplace_back(ParticleHypothesis::pion());
   m_params.emplace_back();
@@ -74,6 +76,7 @@ void VectorTrackContainer::removeTrack_impl(IndexType itrack) {
   };
 
   erase(m_tipIndex);
+  erase(m_stemIndex);
 
   erase(m_params);
   erase(m_cov);
@@ -117,6 +120,7 @@ void VectorTrackContainer::ensureDynamicColumns_impl(
 
 void VectorTrackContainer::reserve(IndexType size) {
   m_tipIndex.reserve(size);
+  m_stemIndex.reserve(size);
 
   m_particleHypothesis.reserve(size);
   m_params.reserve(size);
@@ -139,6 +143,7 @@ void VectorTrackContainer::reserve(IndexType size) {
 
 void VectorTrackContainer::clear() {
   m_tipIndex.clear();
+  m_stemIndex.clear();
 
   m_particleHypothesis.clear();
   m_params.clear();