diff --git a/include/podio/Frame.h b/include/podio/Frame.h index c3b4b4763..ca615c022 100644 --- a/include/podio/Frame.h +++ b/include/podio/Frame.h @@ -9,6 +9,7 @@ #include "podio/SchemaEvolution.h" #include "podio/utilities/TypeHelpers.h" +#include #include #include #include @@ -91,6 +92,7 @@ class Frame { virtual ~FrameConcept() = default; virtual const podio::CollectionBase* get(const std::string& name) const = 0; virtual const podio::CollectionBase* put(std::unique_ptr coll, const std::string& name) = 0; + virtual std::unique_ptr extract(const std::string& name) = 0; virtual podio::GenericParameters& parameters() = 0; virtual const podio::GenericParameters& parameters() const = 0; @@ -122,6 +124,10 @@ class Frame { /// a nullptr const podio::CollectionBase* put(std::unique_ptr coll, const std::string& name) final; + /// Try to extract a collection from the internal storage and return + /// ownership to the caller. If not found return a nullptr + std::unique_ptr extract(const std::string& name) final; + /// Get a reference to the internally used GenericParameters podio::GenericParameters& parameters() override { return *m_parameters; @@ -219,6 +225,30 @@ class Frame { /// if it is not const podio::CollectionBase* get(const std::string& name) const; + /// Extract a collection pointer from the Frame by name + /// + /// @note This does not check if there are any existing references to this + /// collection already. All existing references to the extracted collection + /// and their elements (!) will be invalidated. + /// + /// @returns An owning pointer to the collection if it is available or a + /// nullptr if it is not + std::unique_ptr extract(const std::string& name); + + /// Extract a collection from the Frame by name + /// + /// @tparam CollT The type of the desired collection + /// @param name The name of the collection + /// + /// @note This does not check if there are any existing references to this + /// collection already. All existing references to the extracted collection + /// and their elements (!) will be invalidated. + /// + /// @returns The extracted collection if it existed inside the Frame or + /// otherwise a default initialized empty collection + template + CollT extract(const std::string& name); + /// (Destructively) move a collection into the Frame and get a reference to /// the inserted collection back for further use. /// @@ -406,6 +436,21 @@ inline const podio::CollectionBase* Frame::get(const std::string& name) const { return m_self->get(name); } +inline std::unique_ptr Frame::extract(const std::string& name) { + return m_self->extract(name); +} + +template +CollT Frame::extract(const std::string& name) { + auto coll = extract(name); + auto typedColl = dynamic_cast(coll.get()); + if (typedColl) { + return std::move(*static_cast(coll.release())); + } + + return CollT{}; +} + inline void Frame::put(std::unique_ptr coll, const std::string& name) { const auto* retColl = m_self->put(std::move(coll), name); if (!retColl) { @@ -442,6 +487,21 @@ const podio::CollectionBase* Frame::FrameModel::get(const std::strin return doGet(name); } +template +std::unique_ptr Frame::FrameModel::extract(const std::string& name) { + // Let get do the heavy lifting of setting relations, etc. + if (get(name) == nullptr) { + return nullptr; + } + + { + std::lock_guard lock{*m_mapMtx}; + const auto collHandle = m_collections.extract(name); + assert(!collHandle.empty()); + return std::move(collHandle.mapped()); + } +} + template podio::CollectionBase* Frame::FrameModel::doGet(const std::string& name, bool setReferences) const { { diff --git a/tests/unittests/frame.cpp b/tests/unittests/frame.cpp index 39517f1f5..b9ed7a1a4 100644 --- a/tests/unittests/frame.cpp +++ b/tests/unittests/frame.cpp @@ -443,3 +443,19 @@ TEST_CASE("EIC-Jana2 cleanup use case", "[memory-management][492][174]") { } delete clone; } + +TEST_CASE("Frame extract", "[frame]") { + auto event = podio::Frame{}; + + auto clusters = ExampleClusterCollection(); + clusters.create(3.14f); + clusters.create(42.0f); + event.put(std::move(clusters), "clusters"); + + auto extractedClusters = event.extract("clusters"); + REQUIRE(extractedClusters[0].energy() == 3.14f); + REQUIRE(extractedClusters[1].energy() == 42.0f); + + // Ensure the original collection is no longer there + REQUIRE(event.get("clusters") == nullptr); +}