GroupBy.h

// This file is part of the Acts project.
//
// Copyright (C) 2020 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
#pragma once
 
#include <algorithm>
#include <iterator>
#include <utility>
 
#include "Tracking/Sim/Range.h"
 
namespace ActsExamples {
 
template <typename Iterator, typename KeyGetter>
class GroupBy {
 public:
  using Key = std::decay_t<decltype(KeyGetter()(*Iterator()))>;
  using Group = std::pair<Key, Range<Iterator>>;
  using GroupEndIterator = Iterator;
  class GroupIterator {
   public:
    using iterator_category = std::input_iterator_tag;
    using value_type = Group;
    using difference_type = std::ptrdiff_t;
    using pointer = Group*;
    using reference = Group&;
 
    constexpr GroupIterator(const GroupBy& groupBy, Iterator groupBegin)
        : m_groupBy(groupBy),
          m_groupBegin(groupBegin),
          m_groupEnd(groupBy.findEndOfGroup(groupBegin)) {}
    constexpr GroupIterator& operator++() {
      // make the current end the new group beginning
      std::swap(m_groupBegin, m_groupEnd);
      // find the end of the next group starting from the new beginning
      m_groupEnd = m_groupBy.findEndOfGroup(m_groupBegin);
      return *this;
    }
    constexpr GroupIterator operator++(int) {
      GroupIterator retval = *this;
      ++(*this);
      return retval;
    }
    constexpr Group operator*() const {
      const Key key = (m_groupBegin != m_groupEnd)
                          ? m_groupBy.m_keyGetter(*m_groupBegin)
                          : Key();
      return {key, makeRange(m_groupBegin, m_groupEnd)};
    }
 
   private:
    const GroupBy& m_groupBy;
    Iterator m_groupBegin;
    Iterator m_groupEnd;
 
    friend constexpr bool operator==(const GroupIterator& lhs,
                                     const GroupEndIterator& rhs) {
      return lhs.m_groupBegin == rhs;
    }
    friend constexpr bool operator!=(const GroupIterator& lhs,
                                     const GroupEndIterator& rhs) {
      return not(lhs == rhs);
    }
  };
 
  constexpr GroupBy(Iterator begin, Iterator end,
                    KeyGetter keyGetter = KeyGetter())
      : m_begin(begin), m_end(end), m_keyGetter(std::move(keyGetter)) {}
  constexpr GroupIterator begin() const {
    return GroupIterator(*this, m_begin);
  }
  constexpr GroupEndIterator end() const { return m_end; }
  constexpr bool empty() const { return m_begin == m_end; }
 
 private:
  Iterator m_begin;
  Iterator m_end;
  KeyGetter m_keyGetter;
 
  constexpr Iterator findEndOfGroup(Iterator start) const {
    // check for end so we can safely dereference the start iterator.
    if (start == m_end) {
      return start;
    }
    // search the first element that does not share a key with the start.
    return std::find_if_not(std::next(start), m_end,
                            [this, start](const auto& x) {
                              return m_keyGetter(x) == m_keyGetter(*start);
                            });
  }
};
 
template <typename Container, typename KeyGetter>
auto makeGroupBy(const Container& container, KeyGetter keyGetter)
    -> GroupBy<decltype(std::begin(container)), KeyGetter> {
  return {std::begin(container), std::end(container), std::move(keyGetter)};
}
 
}  // namespace ActsExamples