[wip] Add the possibility to extract collections from a Frame

include/podio/Frame.h

@@ -9,6 +9,7 @@
 #include "podio/SchemaEvolution.h"
 #include "podio/utilities/TypeHelpers.h"
 
+#include <cassert>
 #include <concepts>
 #include <initializer_list>
 #include <memory>
@@ -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<podio::CollectionBase> coll, const std::string& name) = 0;
+    virtual std::unique_ptr<podio::CollectionBase> 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<CollectionBase> 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<podio::CollectionBase> 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<podio::CollectionBase> 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 <CollectionType CollT>
+  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<podio::CollectionBase> Frame::extract(const std::string& name) {
+  return m_self->extract(name);
+}
+
+template <CollectionType CollT>
+CollT Frame::extract(const std::string& name) {
+  auto coll = extract(name);
+  auto typedColl = dynamic_cast<CollT*>(coll.get());
+  if (typedColl) {
+    return std::move(*static_cast<CollT*>(coll.release()));
+  }
+
+  return CollT{};
+}
+
 inline void Frame::put(std::unique_ptr<podio::CollectionBase> 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<FrameDataT>::get(const std::strin
   return doGet(name);
 }
 
+template <typename FrameDataT>
+std::unique_ptr<podio::CollectionBase> Frame::FrameModel<FrameDataT>::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 <typename FrameDataT>
 podio::CollectionBase* Frame::FrameModel<FrameDataT>::doGet(const std::string& name, bool setReferences) const {
   {

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<ExampleClusterCollection>("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);
+}