MCTruthContainer.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#ifndef ALICEO2_DATAFORMATS_MCTRUTH_H_
#define ALICEO2_DATAFORMATS_MCTRUTH_H_
 
#include "GPUCommonRtypes.h" // to have the ClassDef macros
#include <cstdint>           // uint8_t etc
#include <cassert>
#include <stdexcept>
#include <gsl/span> // for guideline support library span
#include <type_traits>
#include <cstring> // memmove, memcpy
#include <memory>
#include <vector>
 
// type traits are needed for the compile time consistency check
// maybe to be moved out of Framework first
//#include "Framework/TypeTraits.h"
 
namespace o2
{
class MCCompLabel;
namespace dataformats
{
 
struct MCTruthHeaderElement {
  MCTruthHeaderElement() = default; // for ROOT IO
 
  MCTruthHeaderElement(uint32_t i) : index(i) {}
  uint32_t index = (uint32_t)-1; // the index into the actual MC track storage (-1 if invalid)
  ClassDefNV(MCTruthHeaderElement, 1);
};
 
template <typename TruthElement>
class MCTruthContainer
{
 private:
  // for the moment we require the truth element to be messageable in order to simply flatten the object
  // if it turnes out that other types are required this needs to be extended and method flatten needs to
  // be conditionally added
  // TODO: activate this check
  // static_assert(o2::framework::is_messageable<TruthElement>::value, "truth element type must be messageable");
 
  std::vector<MCTruthHeaderElement> mHeaderArray; // the header structure array serves as an index into the actual storage
  std::vector<TruthElement> mTruthArray;          // the buffer containing the actual truth information
  std::vector<char> mStreamerData; // buffer used for streaming a flat raw buffer
 
  size_t getSize(uint32_t dataindex) const
  {
    // calculate size / number of labels from a difference in pointed indices
    const auto size = (dataindex < getIndexedSize() - 1)
                        ? getMCTruthHeader(dataindex + 1).index - getMCTruthHeader(dataindex).index
                        : getNElements() - getMCTruthHeader(dataindex).index;
    return size;
  }
 
 public:
  // constructor
  MCTruthContainer() = default;
  // destructor
  ~MCTruthContainer() = default;
  // copy constructor
  MCTruthContainer(const MCTruthContainer& other) = default;
  // move constructor
  MCTruthContainer(MCTruthContainer&& other) = default;
  // construct from raw data
  MCTruthContainer(std::vector<MCTruthHeaderElement>& header, std::vector<TruthElement>& truthArray)
  {
    setFrom(header, truthArray);
  }
  // assignment operator
  MCTruthContainer& operator=(const MCTruthContainer& other) = default;
  // move assignment operator
  MCTruthContainer& operator=(MCTruthContainer&& other) = default;
 
  using self_type = MCTruthContainer<TruthElement>;
  struct FlatHeader {
    uint8_t version = 1;
    uint8_t sizeofHeaderElement = sizeof(MCTruthHeaderElement);
    uint8_t sizeofTruthElement = sizeof(TruthElement);
    uint8_t reserved = 0;
    uint32_t nofHeaderElements;
    uint32_t nofTruthElements;
  };
 
  // access
  MCTruthHeaderElement const& getMCTruthHeader(uint32_t dataindex) const { return mHeaderArray[dataindex]; }
  // access the element directly (can be encapsulated better away)... needs proper element index
  // which can be obtained from the MCTruthHeader startposition and size
  TruthElement const& getElement(uint32_t elementindex) const { return mTruthArray[elementindex]; }
  // return the number of original data indexed here
  size_t getIndexedSize() const { return mHeaderArray.size(); }
  // return the number of elements (labels) pointed to in this container
  size_t getNElements() const { return mTruthArray.size(); }
  // return unterlaying vector of elements
  const std::vector<TruthElement>& getTruthArray() const
  {
    return mTruthArray;
  }
 
  // get individual "view" container for a given data index
  // the caller can do modifications on this view (such as sorting)
  gsl::span<TruthElement> getLabels(uint32_t dataindex)
  {
    if (dataindex >= getIndexedSize()) {
      return gsl::span<TruthElement>();
    }
    return gsl::span<TruthElement>(&mTruthArray[getMCTruthHeader(dataindex).index], getSize(dataindex));
  }
 
  // get individual const "view" container for a given data index
  // the caller can't do modifications on this view
  gsl::span<const TruthElement> getLabels(uint32_t dataindex) const
  {
    if (dataindex >= getIndexedSize()) {
      return gsl::span<const TruthElement>();
    }
    return gsl::span<const TruthElement>(&mTruthArray[getMCTruthHeader(dataindex).index], getSize(dataindex));
  }
 
  void clear()
  {
    mHeaderArray.clear();
    mTruthArray.clear();
  }
 
  // clear and force freeing the memory
  void clear_andfreememory()
  {
    clear();
    // this forces the desctructor being called on existing buffers
    mHeaderArray = std::vector<MCTruthHeaderElement>();
    mTruthArray = std::vector<TruthElement>();
    mStreamerData = std::vector<char>();
  }
 
  // add element for a particular dataindex
  // at the moment only strictly consecutive modes are supported
  // noElement indicates that actually no label/element will be added for this dataindex
  void addElement(uint32_t dataindex, TruthElement const& element, bool noElement = false)
  {
    if (dataindex < mHeaderArray.size()) {
      // look if we have something for this dataindex already
      // must currently be the last one
      if (dataindex != (mHeaderArray.size() - 1)) {
        throw std::runtime_error("MCTruthContainer: unsupported code path");
      }
    } else {
      // assert(dataindex == mHeaderArray.size());
 
      // add empty holes
      int holes = dataindex - mHeaderArray.size();
      assert(holes >= 0);
      for (int i = 0; i < holes; ++i) {
        mHeaderArray.emplace_back(mTruthArray.size());
      }
      // add a new one
      mHeaderArray.emplace_back(mTruthArray.size());
    }
    if (!noElement) {
      mTruthArray.emplace_back(element);
    }
  }
 
  void addNoLabelIndex(uint32_t dataindex)
  {
    addElement(dataindex, TruthElement(), true);
  }
 
  // convenience interface to add multiple labels at once
  // can use elements of any assignable type or sub-type
  template <typename CompatibleLabel>
  void addElements(uint32_t dataindex, gsl::span<CompatibleLabel> elements)
  {
    static_assert(std::is_same<TruthElement, CompatibleLabel>::value ||
                    std::is_assignable<TruthElement, CompatibleLabel>::value ||
                    std::is_base_of<TruthElement, CompatibleLabel>::value,
                  "Need to add compatible labels");
    addNoLabelIndex(dataindex);
    for (auto& e : elements) {
      addElement(dataindex, e);
    }
  }
 
  template <typename CompatibleLabel>
  void addElements(uint32_t dataindex, const std::vector<CompatibleLabel>& v)
  {
    using B = typename std::remove_const<CompatibleLabel>::type;
    auto s = gsl::span<CompatibleLabel>(const_cast<B*>(&v[0]), v.size());
    addElements(dataindex, s);
  }
 
  // Add element at last position or for a previous index
  // (at random access position).
  // This might be a slow process since data has to be moved internally
  // so this function should be used with care.
  void addElementRandomAccess(uint32_t dataindex, TruthElement const& element)
  {
    if (dataindex >= mHeaderArray.size()) {
      // a new dataindex -> push element at back
 
      // we still forbid to leave holes
      assert(dataindex == mHeaderArray.size());
 
      mHeaderArray.resize(dataindex + 1);
      mHeaderArray[dataindex] = mTruthArray.size();
      mTruthArray.emplace_back(element);
    } else {
      // if appending at end use fast function
      if (dataindex == mHeaderArray.size() - 1) {
        addElement(dataindex, element);
        return;
      }
 
      // existing dataindex
      // have to:
      // a) move data;
      // b) insert new element;
      // c) adjust indices of all headers right to this
      auto currentindex = mHeaderArray[dataindex].index;
      auto lastindex = currentindex + getSize(dataindex);
      assert(currentindex >= 0);
 
      // resize truth array
      mTruthArray.resize(mTruthArray.size() + 1);
      // move data (have to do from right to left)
      for (int i = mTruthArray.size() - 1; i > lastindex; --i) {
        mTruthArray[i] = mTruthArray[i - 1];
      }
      // insert new element
      mTruthArray[lastindex] = element;
 
      // fix headers
      for (uint32_t i = dataindex + 1; i < mHeaderArray.size(); ++i) {
        auto oldindex = mHeaderArray[i].index;
        mHeaderArray[i].index = (oldindex != (uint32_t)-1) ? oldindex + 1 : oldindex;
      }
    }
  }
 
  // Set container directly from header + truthArray.
  void setFrom(std::vector<MCTruthHeaderElement>& header, std::vector<TruthElement>& truthArray)
  {
    mHeaderArray = std::move(header);
    mTruthArray = std::move(truthArray);
  }
 
  // merge another container to the back of this one
  void mergeAtBack(MCTruthContainer<TruthElement> const& other)
  {
    const auto oldtruthsize = mTruthArray.size();
    const auto oldheadersize = mHeaderArray.size();
 
    // copy from other
    std::copy(other.mHeaderArray.begin(), other.mHeaderArray.end(), std::back_inserter(mHeaderArray));
    std::copy(other.mTruthArray.begin(), other.mTruthArray.end(), std::back_inserter(mTruthArray));
 
    // adjust information of newly attached part
    for (uint32_t i = oldheadersize; i < mHeaderArray.size(); ++i) {
      mHeaderArray[i].index += oldtruthsize;
    }
  }
 
  // merge part of another container ("n" entries starting from "from") to the back of this one
  void mergeAtBack(MCTruthContainer<TruthElement> const& other, size_t from, size_t n)
  {
    const auto oldtruthsize = mTruthArray.size();
    const auto oldheadersize = mHeaderArray.size();
    auto endIdx = from + n;
    assert(endIdx <= other.mHeaderArray.size());
    const auto* headBeg = &other.mHeaderArray[from];
    const auto* headEnd = headBeg + n;
    const auto* trtArrBeg = &other.mTruthArray[other.getMCTruthHeader(from).index];
    const auto* trtArrEnd = (endIdx == other.mHeaderArray.size()) ? (&other.mTruthArray.back()) + 1 : &other.mTruthArray[other.getMCTruthHeader(endIdx).index];
 
    // copy from other
    std::copy(headBeg, headEnd, std::back_inserter(mHeaderArray));
    std::copy(trtArrBeg, trtArrEnd, std::back_inserter(mTruthArray));
    long offset = long(oldtruthsize) - other.getMCTruthHeader(from).index;
    // adjust information of newly attached part
    for (uint32_t i = oldheadersize; i < mHeaderArray.size(); ++i) {
      mHeaderArray[i].index += offset;
    }
  }
 
  template <typename ContainerType>
  size_t flatten_to(ContainerType& container) const
  {
    size_t bufferSize = sizeof(FlatHeader) + sizeof(MCTruthHeaderElement) * mHeaderArray.size() + sizeof(TruthElement) * mTruthArray.size();
    container.resize((bufferSize / sizeof(typename ContainerType::value_type)) + ((bufferSize % sizeof(typename ContainerType::value_type)) > 0 ? 1 : 0));
    char* target = reinterpret_cast<char*>(container.data());
    auto& flatheader = *reinterpret_cast<FlatHeader*>(target);
    target += sizeof(FlatHeader);
    flatheader.version = 1;
    flatheader.sizeofHeaderElement = sizeof(MCTruthHeaderElement);
    flatheader.sizeofTruthElement = sizeof(TruthElement);
    flatheader.reserved = 0;
    flatheader.nofHeaderElements = mHeaderArray.size();
    flatheader.nofTruthElements = mTruthArray.size();
    size_t copySize = flatheader.sizeofHeaderElement * flatheader.nofHeaderElements;
    memcpy(target, mHeaderArray.data(), copySize);
    target += copySize;
    copySize = flatheader.sizeofTruthElement * flatheader.nofTruthElements;
    memcpy(target, mTruthArray.data(), copySize);
    return bufferSize;
  }
 
  void restore_from(const char* buffer, size_t bufferSize)
  {
    if (buffer == nullptr || bufferSize < sizeof(FlatHeader)) {
      return;
    }
    auto* source = buffer;
    auto& flatheader = *reinterpret_cast<FlatHeader const*>(source);
    source += sizeof(FlatHeader);
    if (bufferSize < sizeof(FlatHeader) + flatheader.sizeofHeaderElement * flatheader.nofHeaderElements + flatheader.sizeofTruthElement * flatheader.nofTruthElements) {
      throw std::runtime_error("inconsistent buffer size: too small");
      return;
    }
    if (flatheader.sizeofHeaderElement != sizeof(MCTruthHeaderElement) || flatheader.sizeofTruthElement != sizeof(TruthElement)) {
      // not yet handled
      throw std::runtime_error("member element sizes don't match");
    }
    // TODO: with a spectator memory ressource the vectors can be built directly
    // over the original buffer, there is the implementation for a memory ressource
    // working on a FairMQ message, here we would need two memory resources over
    // the two ranges of the input buffer
    // for now doing a copy
    mHeaderArray.resize(flatheader.nofHeaderElements);
    mTruthArray.resize(flatheader.nofTruthElements);
    size_t copySize = flatheader.sizeofHeaderElement * flatheader.nofHeaderElements;
    memcpy(mHeaderArray.data(), source, copySize);
    source += copySize;
    copySize = flatheader.sizeofTruthElement * flatheader.nofTruthElements;
    memcpy(mTruthArray.data(), source, copySize);
  }
 
  template <typename Stream>
  void print(Stream& stream)
  {
    stream << "MCTruthContainer index = " << getIndexedSize() << " for " << getNElements() << " elements(s), flat buffer size " << mStreamerData.size() << "\n";
  }
 
  void inflate()
  {
    if (mHeaderArray.size() > 0) {
      mStreamerData.clear();
      return;
    }
    restore_from(mStreamerData.data(), mStreamerData.size());
    mStreamerData.clear();
  }
 
  void deflate()
  {
    if (mStreamerData.size() > 0) {
      clear();
      return;
    }
    mStreamerData.clear();
    flatten_to(mStreamerData);
    clear();
  }
 
  ClassDefNV(MCTruthContainer, 2);
}; // end class
 
using MCLabelContainer = o2::dataformats::MCTruthContainer<o2::MCCompLabel>;
 
} // namespace dataformats
} // namespace o2
 
#endif