JANA/Components/JPodioOutput.h

0001 #pragma once
0002 #include <JANA/Utils/JTypeInfo.h>
0003 #include <JANA/JEvent.h>
0004 #include <JANA/Components/JHasDatabundleOutputs.h>
0005 #include <JANA/Components/JPodioDatabundle.h>
0006 #include <podio/Frame.h>
0007 #include <memory>
0008
0009
0010 namespace jana::components {
0011
0012
0013 template <typename PodioT>
0014 class PodioOutput : public JHasDatabundleOutputs::OutputBase {
0015 private:
0016     std::unique_ptr<typename PodioT::collection_type> m_transient_collection;
0017     JPodioDatabundle* m_podio_databundle;
0018 public:
0019     PodioOutput(JHasDatabundleOutputs* owner, std::string default_collection_name="") {
0020         owner->RegisterOutput(this);
0021         auto bundle = std::make_unique<JPodioDatabundle>();
0022         bundle->SetUniqueName(default_collection_name);
0023         bundle->SetTypeName(JTypeInfo::demangle<PodioT>());
0024         m_podio_databundle = bundle.get();
0025         m_databundles.push_back(std::move(bundle));
0026         m_transient_collection = std::move(std::make_unique<typename PodioT::collection_type>());
0027     }
0028
0029     std::unique_ptr<typename PodioT::collection_type>& operator()() { return m_transient_collection; }
0030
0031     JPodioDatabundle* GetDatabundle() const { return m_podio_databundle; }
0032
0033
0034 protected:
0035
0036     void StoreData(const JEvent& event) override {
0037         podio::Frame* frame;
0038         try {
0039             frame = const_cast<podio::Frame*>(event.GetSingle<podio::Frame>());
0040             if (frame == nullptr) {
0041                 frame = new podio::Frame;
0042                 event.Insert<podio::Frame>(frame);
0043             }
0044         }
0045         catch (...) {
0046             frame = new podio::Frame;
0047             event.Insert<podio::Frame>(frame);
0048         }
0049
0050         frame->put(std::move(m_transient_collection), m_podio_databundle->GetUniqueName());
0051         const auto* moved = &frame->template get<typename PodioT::collection_type>(m_podio_databundle->GetUniqueName());
0052         m_transient_collection = nullptr;
0053         m_podio_databundle->SetCollection(moved);
0054     }
0055     void Reset() override {
0056         m_transient_collection = std::move(std::make_unique<typename PodioT::collection_type>());
0057     }
0058 };
0059
0060
0061 template <typename PodioT>
0062 class VariadicPodioOutput : public JHasDatabundleOutputs::OutputBase {
0063 private:
0064     std::vector<std::unique_ptr<typename PodioT::collection_type>> m_collections;
0065     std::vector<JPodioDatabundle*> m_databundles;
0066
0067 public:
0068     VariadicPodioOutput(JHasDatabundleOutputs* owner, std::vector<std::string> default_collection_names={}) {
0069         owner->RegisterOutput(this);
0070         this->m_is_variadic = true;
0071         for (const std::string& name : default_collection_names) {
0072             auto coll = std::make_unique<JPodioDatabundle>();
0073             coll->SetUniqueName(name);
0074             coll->SetTypeName(JTypeInfo::demangle<PodioT>());
0075             m_collections.push_back(std::move(coll));
0076         }
0077         for (auto& coll_name : this->collection_names) {
0078             m_collections.push_back(std::make_unique<typename PodioT::collection_type>());
0079         }
0080     }
0081     void StoreData(const JEvent& event) override {
0082         if (m_collections.size() != this->collection_names.size()) {
0083             throw JException("VariadicPodioOutput InsertCollection failed: Declared %d collections, but provided %d.", this->collection_names.size(), m_collections.size());
0084         }
0085
0086         podio::Frame* frame;
0087         try {
0088             frame = const_cast<podio::Frame*>(event.GetSingle<podio::Frame>());
0089             if (frame == nullptr) {
0090                 frame = new podio::Frame;
0091                 event.Insert<podio::Frame>(frame);
0092             }
0093         }
0094         catch (...) {
0095             frame = new podio::Frame;
0096             event.Insert<podio::Frame>(frame);
0097         }
0098
0099         size_t i = 0;
0100         for (auto& collection : m_collections) {
0101             frame->put(std::move(std::move(collection)), m_databundles[i]->GetUniqueName());
0102             const auto* moved = &frame->template get<typename PodioT::collection_type>(m_databundles[i]->GetUniqueName());
0103             collection = nullptr;
0104             const auto &databundle = dynamic_cast<JPodioDatabundle*>(m_databundles[i]);
0105             databundle->SetCollection(moved);
0106             i += 1;
0107         }
0108     }
0109     void Reset() override {
0110         m_collections.clear();
0111         for (auto& coll : this->m_databundles) {
0112             coll->ClearData();
0113         }
0114         for (auto& coll_name : this->collection_names) {
0115             m_collections.push_back(std::make_unique<typename PodioT::collection_type>());
0116         }
0117     }
0118 };
0119
0120
0121 } // namespace jana::components
0122