podio/detail/LinkCollectionImpl.h

0001 #ifndef PODIO_DETAIL_LINKCOLLECTIONIMPL_H
0002 #define PODIO_DETAIL_LINKCOLLECTIONIMPL_H
0003
0004 #include "podio/ObjectID.h"
0005 #include "podio/detail/Link.h"
0006 #include "podio/detail/LinkCollectionData.h"
0007 #include "podio/detail/LinkCollectionIterator.h"
0008 #include "podio/detail/LinkFwd.h"
0009 #include "podio/detail/LinkObj.h"
0010
0011 #include "podio/CollectionBase.h"
0012 #include "podio/CollectionBufferFactory.h"
0013 #include "podio/CollectionBuffers.h"
0014 #include "podio/DatamodelRegistry.h"
0015 #include "podio/ICollectionProvider.h"
0016 #include "podio/SchemaEvolution.h"
0017 #include "podio/detail/Pythonizations.h"
0018 #include "podio/utilities/MaybeSharedPtr.h"
0019 #include "podio/utilities/StaticConcatenate.h"
0020 #include "podio/utilities/TypeHelpers.h"
0021
0022 #include <iterator>
0023 #include <string_view>
0024
0025 #ifdef PODIO_JSON_OUTPUT
0026   #include "nlohmann/json.hpp"
0027 #endif
0028
0029 #include <iomanip>
0030 #include <memory>
0031 #include <mutex>
0032 #include <ostream>
0033 #include <string>
0034 #include <type_traits>
0035
0036 namespace podio {
0037
0038 template <typename FromT, typename ToT>
0039 class LinkCollection : public podio::CollectionBase {
0040   static_assert(std::is_same_v<FromT, detail::GetDefaultHandleType<FromT>>,
0041                 "Links need to be instantiated with the default types!");
0042   static_assert(std::is_same_v<ToT, detail::GetDefaultHandleType<ToT>>,
0043                 "Links need to be instantiated with the default types!");
0044
0045   // convenience typedefs
0046   using CollectionDataT = podio::LinkCollectionData<FromT, ToT>;
0047
0048 public:
0049   using from_type = FromT;
0050   using to_type = ToT;
0051   using value_type = Link<FromT, ToT>;
0052   using mutable_type = MutableLink<FromT, ToT>;
0053   using const_iterator = LinkCollectionIterator<FromT, ToT>;
0054   using iterator = LinkMutableCollectionIterator<FromT, ToT>;
0055   using difference_type = ptrdiff_t;
0056   using size_type = size_t;
0057   using const_reverse_iterator = std::reverse_iterator<const_iterator>;
0058   using reverse_iterator = std::reverse_iterator<iterator>;
0059
0060   LinkCollection() = default;
0061
0062   LinkCollection(CollectionDataT&& data, bool isSubsetColl) :
0063       m_isSubsetColl(isSubsetColl), m_collectionID(0), m_storage(std::move(data)) {
0064   }
0065
0066   // Move-only type
0067   LinkCollection(const LinkCollection&) = delete;
0068   LinkCollection& operator=(const LinkCollection&) = delete;
0069   LinkCollection(LinkCollection&&) = default;
0070   LinkCollection& operator=(LinkCollection&&) = default;
0071
0072   ~LinkCollection() override {
0073     // Need to tell the storage how to clean up
0074     m_storage.clear(m_isSubsetColl);
0075   }
0076
0077   /// Append a new link to the collection and return this object
0078   mutable_type create() {
0079     if (m_isSubsetColl) {
0080       throw std::logic_error("Cannot create new elements on a subset collection");
0081     }
0082
0083     auto obj = m_storage.entries.emplace_back(new LinkObj<FromT, ToT>());
0084     obj->id = {int(m_storage.entries.size() - 1), m_collectionID};
0085     return mutable_type(podio::utils::MaybeSharedPtr(obj));
0086   }
0087
0088   /// Returns the immutable object of given index
0089   value_type operator[](std::size_t index) const {
0090     return value_type(m_storage.entries[index]);
0091   }
0092   /// Returns the mutable object of given index
0093   mutable_type operator[](std::size_t index) {
0094     return mutable_type(podio::utils::MaybeSharedPtr(m_storage.entries[index]));
0095   }
0096   /// Returns the immutable object of given index
0097   value_type at(std::size_t index) const {
0098     return value_type(m_storage.entries.at(index));
0099   }
0100   /// Returns the mutable object of given index
0101   mutable_type at(std::size_t index) {
0102     return mutable_type(podio::utils::MaybeSharedPtr(m_storage.entries.at(index)));
0103   }
0104
0105   void push_back(mutable_type object) {
0106     // We have to do different things here depending on whether this is a
0107     // subset collection or not. A normal collection cannot collect objects
0108     // that are already part of another collection, while a subset collection
0109     // can only collect such objects
0110     if (!m_isSubsetColl) {
0111       auto obj = object.m_obj;
0112       if (obj->id.index == podio::ObjectID::untracked) {
0113         obj->id = {static_cast<int>(m_storage.entries.size()), m_collectionID};
0114         m_storage.entries.push_back(obj.release());
0115       } else {
0116         throw std::invalid_argument("Object already in a collection. Cannot add it to a second collection");
0117       }
0118
0119     } else {
0120       push_back(value_type(object));
0121     }
0122   }
0123
0124   void push_back(value_type object) {
0125     if (!m_isSubsetColl) {
0126       throw std::invalid_argument("Can only add immutable objects to subset collections");
0127     }
0128     auto obj = object.m_obj;
0129     if (obj->id.index < 0) {
0130       throw std::invalid_argument(
0131           "Object needs to be tracked by another collection in order for it to be storable in a subset collection");
0132     }
0133     m_storage.entries.push_back(obj.release());
0134   }
0135
0136   /// Number of elements in the collection
0137   size_t size() const override {
0138     return m_storage.entries.size();
0139   }
0140
0141   /// maximal number of elements in the collection
0142   std::size_t max_size() const override {
0143     const auto maxStorage = m_storage.entries.max_size();
0144     if (!m_isSubsetColl) {
0145       // non-subset collections shouldn't have more elements than the maximum index of ObjectID
0146       const auto maxIndex = std::numeric_limits<decltype(ObjectID::index)>::max();
0147       return std::min<size_t>(maxIndex, maxStorage);
0148     }
0149     return maxStorage;
0150   }
0151
0152   /// Is the collection empty
0153   bool empty() const override {
0154     return m_storage.entries.empty();
0155   }
0156
0157   void clear() override {
0158     m_storage.clear(m_isSubsetColl);
0159     m_isPrepared = false;
0160   }
0161
0162   void print(std::ostream& os = std::cout, bool flush = true) const override {
0163     os << *this;
0164     if (flush) {
0165       os.flush();
0166     }
0167   }
0168
0169   // support for the iterator protocol
0170   const_iterator begin() const {
0171     return const_iterator(0, &m_storage.entries);
0172   }
0173   const_iterator cbegin() const {
0174     return begin();
0175   }
0176   const_iterator end() const {
0177     return const_iterator(m_storage.entries.size(), &m_storage.entries);
0178   }
0179   const_iterator cend() const {
0180     return end();
0181   }
0182   iterator begin() {
0183     return iterator(0, &m_storage.entries);
0184   }
0185   iterator end() {
0186     return iterator(m_storage.entries.size(), &m_storage.entries);
0187   }
0188   // reverse iterators
0189   reverse_iterator rbegin() {
0190     return reverse_iterator(end());
0191   }
0192   const_reverse_iterator rbegin() const {
0193     return const_reverse_iterator(end());
0194   }
0195   const_reverse_iterator crbegin() const {
0196     return rbegin();
0197   }
0198   reverse_iterator rend() {
0199     return reverse_iterator(begin());
0200   }
0201   const_reverse_iterator rend() const {
0202     return const_reverse_iterator(begin());
0203   }
0204   const_reverse_iterator crend() const {
0205     return rend();
0206   }
0207
0208   /// check if the collection has a valid ID
0209   bool hasID() const override {
0210     return getID() != static_cast<uint32_t>(podio::ObjectID::untracked) &&
0211         getID() != static_cast<uint32_t>(podio::ObjectID::invalid);
0212   }
0213
0214   [[deprecated("isValid will be removed, use hasID() if you want to check if it has an ID, otherwise assume the "
0215                "collection is valid")]]
0216   bool isValid() const override {
0217     return hasID();
0218   }
0219
0220   podio::CollectionWriteBuffers getBuffers() override {
0221     return m_storage.getCollectionBuffers(m_isSubsetColl);
0222   }
0223
0224   constexpr static std::string_view typeName =
0225       podio::utils::static_concatenate_v<podio::detail::link_coll_name_prefix, FromT::typeName,
0226                                          podio::detail::link_name_infix, ToT::typeName,
0227                                          podio::detail::link_name_suffix>;
0228
0229   constexpr static std::string_view valueTypeName = value_type::typeName;
0230   constexpr static std::string_view dataTypeName = "podio::LinkData";
0231
0232   const std::string_view getTypeName() const override {
0233     return typeName;
0234   }
0235
0236   const std::string_view getValueTypeName() const override {
0237     return valueTypeName;
0238   }
0239
0240   const std::string_view getDataTypeName() const override {
0241     return dataTypeName;
0242   }
0243
0244   /// Cppyy protocol to setup the pythonizations for this class. Not to be called directly.
0245   static void __cppyy_pythonize__(PyObject* klass, const std::string& name) {
0246     podio::detail::pythonizations::pythonize_subscript(klass, name);
0247   }
0248
0249   bool isSubsetCollection() const override {
0250     return m_isSubsetColl;
0251   }
0252
0253   void setSubsetCollection(bool setSubset = true) override {
0254     if (m_isSubsetColl != setSubset && !m_storage.entries.empty()) {
0255       throw std::logic_error("Cannot change the character of a collection that already contains elements");
0256     }
0257
0258     if (setSubset) {
0259       m_storage.makeSubsetCollection();
0260     }
0261     m_isSubsetColl = setSubset;
0262   }
0263
0264   void setID(uint32_t id) override {
0265     m_collectionID = id;
0266     if (!m_isSubsetColl) {
0267       std::ranges::for_each(m_storage.entries, [id](auto* obj) { obj->id = {obj->id.index, id}; });
0268     }
0269   }
0270
0271   uint32_t getID() const override {
0272     return m_collectionID;
0273   }
0274
0275   void prepareForWrite() const override {
0276     std::lock_guard lock{*m_storageMtx};
0277     // by the time we acquire the lock another thread might have already
0278     // succeeded, so we need to check again now
0279     if (m_isPrepared) {
0280       return;
0281     }
0282     m_storage.prepareForWrite(m_isSubsetColl);
0283     m_isPrepared = true;
0284   }
0285
0286   void prepareAfterRead() override {
0287     // No need to go through this again if we have already done it
0288     if (m_isPrepared) {
0289       return;
0290     }
0291
0292     if (!m_isSubsetColl) {
0293       // Subset collections do not store any data that would require post-processing
0294       m_storage.prepareAfterRead(m_collectionID);
0295     }
0296     // Preparing a collection doesn't affect the underlying I/O buffers, so this
0297     // collection is still prepared
0298     m_isPrepared = true;
0299   }
0300
0301   bool setReferences(const ICollectionProvider* collectionProvider) override {
0302     return m_storage.setReferences(collectionProvider, m_isSubsetColl);
0303   }
0304
0305   static constexpr SchemaVersionT schemaVersion = 1;
0306
0307   SchemaVersionT getSchemaVersion() const override {
0308     return schemaVersion;
0309   }
0310
0311   size_t getDatamodelRegistryIndex() const override {
0312     return podio::DatamodelRegistry::NoDefinitionNecessary;
0313   }
0314
0315 private:
0316   // For setReferences, we need to give our own CollectionData access to our
0317   // private entries. Otherwise we would need to expose a public member function
0318   // that gives access to the Obj* which is definitely not what we want
0319   friend CollectionDataT;
0320
0321   mutable bool m_isPrepared{false};
0322   bool m_isSubsetColl{false};
0323   uint32_t m_collectionID{0};
0324   mutable std::unique_ptr<std::mutex> m_storageMtx{std::make_unique<std::mutex>()};
0325   mutable CollectionDataT m_storage{};
0326 };
0327
0328 template <typename FromT, typename ToT>
0329 std::ostream& operator<<(std::ostream& o, const LinkCollection<FromT, ToT>& v) {
0330   const auto old_flags = o.flags();
0331   o << "          id:      weight:" << '\n';
0332   for (const auto&& el : v) {
0333     o << std::scientific << std::showpos << std::setw(12) << el.id() << " " << std::setw(12) << " " << el.getWeight()
0334       << '\n';
0335
0336     o << "     from : ";
0337     o << el.getFrom().id() << std::endl;
0338     o << "       to : ";
0339     o << el.getTo().id() << std::endl;
0340   }
0341
0342   o.flags(old_flags);
0343   return o;
0344 }
0345
0346 namespace detail {
0347   template <typename FromT, typename ToT>
0348   podio::CollectionReadBuffers createLinkBuffers(bool subsetColl) {
0349     auto readBuffers = podio::CollectionReadBuffers{};
0350     readBuffers.type = podio::LinkCollection<FromT, ToT>::typeName;
0351     readBuffers.schemaVersion = podio::LinkCollection<FromT, ToT>::schemaVersion;
0352     readBuffers.data = subsetColl ? nullptr : new LinkDataContainer();
0353
0354     // Either it is a subset collection or we have two relations
0355     const auto nRefs = subsetColl ? 1 : 2;
0356     readBuffers.references = new podio::CollRefCollection(nRefs);
0357     for (auto& ref : *readBuffers.references) {
0358       // Make sure to place usable buffer pointers here
0359       ref = std::make_unique<std::vector<podio::ObjectID>>();
0360     }
0361
0362     readBuffers.createCollection = [](const podio::CollectionReadBuffers& buffers, bool isSubsetColl) {
0363       LinkCollectionData<FromT, ToT> data(buffers, isSubsetColl);
0364       return std::make_unique<LinkCollection<FromT, ToT>>(std::move(data), isSubsetColl);
0365     };
0366
0367     readBuffers.recast = [](podio::CollectionReadBuffers& buffers) {
0368       if (buffers.data) {
0369         buffers.data = podio::CollectionWriteBuffers::asVector<float>(buffers.data);
0370       }
0371     };
0372
0373     readBuffers.deleteBuffers = [](const podio::CollectionReadBuffers& buffers) {
0374       if (buffers.data) {
0375         // If we have data then we are not a subset collection and we have
0376         // to clean up all type erased buffers by casting them back to
0377         // something that we can delete
0378         delete static_cast<LinkDataContainer*>(buffers.data);
0379       }
0380       delete buffers.references;
0381       delete buffers.vectorMembers;
0382     };
0383
0384     return readBuffers;
0385   }
0386
0387   template <typename FromT, typename ToT>
0388   bool registerLinkCollection(const std::string_view linkTypeName) {
0389     const static auto reg = [&linkTypeName]() {
0390       const auto schemaVersion = LinkCollection<FromT, ToT>::schemaVersion;
0391
0392       auto& factory = CollectionBufferFactory::mutInstance();
0393       factory.registerCreationFunc(std::string(linkTypeName), schemaVersion, createLinkBuffers<FromT, ToT>);
0394
0395       // For now passing the same schema version for from and current version
0396       // simply to make SchemaEvolution aware of LinkCollections
0397       podio::SchemaEvolution::mutInstance().registerEvolutionFunc(std::string(linkTypeName), schemaVersion,
0398                                                                   schemaVersion, SchemaEvolution::noOpSchemaEvolution,
0399                                                                   SchemaEvolution::Priority::AutoGenerated);
0400
0401       return true;
0402     }();
0403     return reg;
0404   }
0405 } // namespace detail
0406
0407 #if defined(PODIO_JSON_OUTPUT) && !defined(__CLING__)
0408 template <typename FromT, typename ToT>
0409 void to_json(nlohmann::json& j, const podio::LinkCollection<FromT, ToT>& collection) {
0410   j = nlohmann::json::array();
0411   for (auto&& elem : collection) {
0412     j.emplace_back(elem);
0413   }
0414 }
0415 #endif
0416
0417 } // namespace podio
0418
0419 #endif // PODIO_DETAIL_LINKCOLLECTIONIMPL_H