From 11fcf1ffd589eb8d6233a3d8ceb3d6115e3f23cb Mon Sep 17 00:00:00 2001
From: Sandro Wenzel <sandro.wenzel@cern.ch>
Date: Mon, 1 Mar 2021 16:09:51 +0100
Subject: [PATCH] MCReader: release early

---
 Steer/include/Steer/MCKinematicsReader.h | 22 ++++++++++-------
 Steer/src/MCKinematicsReader.cxx         | 30 +++++++++++++++++-------
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/Steer/include/Steer/MCKinematicsReader.h b/Steer/include/Steer/MCKinematicsReader.h
index 03b3709423b1c..0f604884fa48a 100644
--- a/Steer/include/Steer/MCKinematicsReader.h
+++ b/Steer/include/Steer/MCKinematicsReader.h
@@ -75,6 +75,9 @@ class MCKinematicsReader
   /// variant returning all tracks for source and event at once
   std::vector<MCTrack> const& getTracks(int source, int event) const;
 
+  /// API to ask releasing tracks (freeing memory) for source + event
+  void releaseTracksForSourceAndEvent(int source, int event);
+
   /// variant returning all tracks for source and event at once
   std::vector<MCTrack> const& getTracks(int event) const;
 
@@ -106,7 +109,8 @@ class MCKinematicsReader
   }
 
  private:
-  void loadTracksForSource(int source) const;
+  void initTracksForSource(int source) const;
+  void loadTracksForSourceAndEvent(int source, int eventID) const;
   void loadHeadersForSource(int source) const;
   void loadTrackRefsForSource(int source) const;
   void initIndexedTrackRefs(std::vector<o2::TrackReference>& refs, o2::dataformats::MCTruthContainer<o2::TrackReference>& indexedrefs) const;
@@ -117,7 +121,7 @@ class MCKinematicsReader
   std::vector<TChain*> mInputChains;
 
   // a vector of tracks foreach source and each collision
-  mutable std::vector<std::vector<std::vector<o2::MCTrack>>> mTracks;                                        // the in-memory track container
+  mutable std::vector<std::vector<std::vector<o2::MCTrack>*>> mTracks;                                       // the in-memory track container
   mutable std::vector<std::vector<o2::dataformats::MCEventHeader>> mHeaders;                                 // the in-memory header container
   mutable std::vector<std::vector<o2::dataformats::MCTruthContainer<o2::TrackReference>>> mIndexedTrackRefs; // the in-memory track ref container
 
@@ -134,10 +138,7 @@ inline MCTrack const* MCKinematicsReader::getTrack(o2::MCCompLabel const& label)
 
 inline MCTrack const* MCKinematicsReader::getTrack(int source, int event, int track) const
 {
-  if (mTracks[source].size() == 0) {
-    loadTracksForSource(source);
-  }
-  return &mTracks[source][event][track];
+  return &getTracks(source, event)[track];
 }
 
 inline MCTrack const* MCKinematicsReader::getTrack(int event, int track) const
@@ -148,9 +149,12 @@ inline MCTrack const* MCKinematicsReader::getTrack(int event, int track) const
 inline std::vector<MCTrack> const& MCKinematicsReader::getTracks(int source, int event) const
 {
   if (mTracks[source].size() == 0) {
-    loadTracksForSource(source);
+    initTracksForSource(source);
+  }
+  if (mTracks[source][event] == nullptr) {
+    loadTracksForSourceAndEvent(source, event);
   }
-  return mTracks[source][event];
+  return *mTracks[source][event];
 }
 
 inline std::vector<MCTrack> const& MCKinematicsReader::getTracks(int event) const
@@ -195,7 +199,7 @@ inline size_t MCKinematicsReader::getNSources() const
 inline size_t MCKinematicsReader::getNEvents(int source) const
 {
   if (mTracks[source].size() == 0) {
-    loadTracksForSource(source);
+    initTracksForSource(source);
   }
   return mTracks[source].size();
 }
diff --git a/Steer/src/MCKinematicsReader.cxx b/Steer/src/MCKinematicsReader.cxx
index 3f0b39cee693e..e5979e47987e6 100644
--- a/Steer/src/MCKinematicsReader.cxx
+++ b/Steer/src/MCKinematicsReader.cxx
@@ -49,7 +49,17 @@ void MCKinematicsReader::initIndexedTrackRefs(std::vector<o2::TrackReference>& r
   }
 }
 
-void MCKinematicsReader::loadTracksForSource(int source) const
+void MCKinematicsReader::initTracksForSource(int source) const
+{
+  auto chain = mInputChains[source];
+  if (chain) {
+    // todo: get name from NameConfig
+    auto br = chain->GetBranch("MCTrack");
+    mTracks[source].resize(br->GetEntries(), nullptr);
+  }
+}
+
+void MCKinematicsReader::loadTracksForSourceAndEvent(int source, int event) const
 {
   auto chain = mInputChains[source];
   if (chain) {
@@ -58,18 +68,22 @@ void MCKinematicsReader::loadTracksForSource(int source) const
     if (br) {
       std::vector<MCTrack>* loadtracks = nullptr;
       br->SetAddress(&loadtracks);
-      // load all kinematics
-      mTracks[source].resize(br->GetEntries());
-      for (int event = 0; event < br->GetEntries(); ++event) {
-        br->GetEntry(event);
-        mTracks[source][event] = *loadtracks;
-      }
+      br->GetEntry(event);
+      mTracks[source][event] = new std::vector<o2::MCTrack>;
+      *mTracks[source][event] = *loadtracks;
       delete loadtracks;
-      loadtracks = nullptr;
     }
   }
 }
 
+void MCKinematicsReader::releaseTracksForSourceAndEvent(int source, int eventID)
+{
+  if (mTracks.at(source).at(eventID) != nullptr) {
+    delete mTracks[source][eventID];
+    mTracks[source][eventID] = nullptr;
+  }
+}
+
 void MCKinematicsReader::loadHeadersForSource(int source) const
 {
   auto chain = mInputChains[source];