@@ -59,6 +59,8 @@ struct RecoDecay {
   static constexpr int PdgBosonMax{37};                          // largest boson (gauge or H) PDG code; o2-linter: disable=pdg/explicit-code (H+ does not have a named constant.)
   static constexpr int PdgDivisorMeson{100};                     // order of magnitude of the meson PDG codes
   static constexpr int PdgDivisorBaryon{1000};                   // order of magnitude of the baryon PDG codes
+  static constexpr int PdgDiquarkSpin0Code{1};                   // last digits of a spin-0 diquark PDG code
+  static constexpr int PdgDiquarkSpin1Code{3};                   // last digits of a spin-1 diquark PDG code
 
   // Auxiliary functions
 
@@ -1179,6 +1181,27 @@ struct RecoDecay {
     return OriginType::None;
   }
 
+  /// Check whether the PDG code corresponds to a diquark state (e.g. 4101, 4201, 5101).
+  /// Diquarks are constituents used by the generators during the hadronisation and are not
+  /// hadrons, but their PDG codes satisfy the pdg / 1000 == 4 or 5 condition that is used
+  /// to identify charm and beauty baryons in getParticleOrigin.
+  /// \param pdgCode PDG code of the particle
+  /// \return true if the PDG code belongs to a diquark state
+  static bool isDiquark(const int pdgCode)
+  {
+    // a diquark PDG code is composed as quark1 * PdgDivisorBaryon + quark2 * PdgDivisorMeson +
+    // spinCode, with quark1 >= quark2 the quark flavours and spinCode either PdgDiquarkSpin0Code
+    // or PdgDiquarkSpin1Code
+    const int absPdgCode = std::abs(pdgCode);
+    const int spinCode = absPdgCode % PdgDivisorMeson;
+    if (spinCode != PdgDiquarkSpin0Code && spinCode != PdgDiquarkSpin1Code) {
+      return false;
+    }
+    const int quark1 = absPdgCode / PdgDivisorBaryon;
+    const int quark2 = (absPdgCode % PdgDivisorBaryon) / PdgDivisorMeson;
+    return quark1 >= PDG_t::kDown && quark1 <= PDG_t::kTop && quark2 >= PDG_t::kDown && quark1 >= quark2;
+  }
+
   /// based on getCharmHardronOrigin in order to extend general particle
   /// Finding the origin (from charm hadronisation or beauty-hadron decay) of paritcle (b, c and others)
   /// \param particlesMC  table with MC particles
@@ -1221,6 +1244,17 @@ struct RecoDecay {
             // Check status code
             // auto motherStatusCode = std::abs(mother.getGenStatusCode());
             auto pdgParticleIMother = std::abs(mother.pdgCode()); // PDG code of the mother
+            // in the hadron-level mode, stop the walk when a parton or a diquark is reached: no
+            // heavy-flavour hadron can be found above the partonic level, and a diquark is not a
+            // hadron although its PDG code (e.g. 4101, 4201, 5101) satisfies pdg / 1000 == 4 or 5
+            // and would be misidentified as a charm/beauty baryon below. In the parton-level mode
+            // the walk has to continue through diquarks, since a charm or beauty quark can still be
+            // found above them.
+            if (!searchUpToQuark &&
+                (isDiquark(pdgParticleIMother) || pdgParticleIMother <= PdgQuarkMax ||
+                 (pdgParticleIMother >= PdgBosonMin && pdgParticleIMother <= PdgBosonMax))) {
+              break;
+            }
             // Check mother's PDG code.
             // printf("getMother: ");
             // for (int i = stage; i < 0; i++) // Indent to make the tree look nice.
 
@@ -750,6 +750,7 @@ enum ParticleOriginMCTruth {
   kElse,                         //! none of the above; (NOTE: used to catch bugs. will be removed once MC usage is properly validated)
   kSecondaryFromHF,              //! Kaon from a charm- or beauty-hadron decay chain
   kSecondaryFromLF,              //! Kaon from a light-flavour decay without a heavy-flavour ancestor
+  kSecondaryFromPhi,             //! Kaon from a phi decay chain (phi -> K0/K0bar -> K0S)
   kNOriginMCTruthTypes
 };
 
@@ -771,7 +772,8 @@ static constexpr std::string_view ParticleOriginMCTruthName[kNOriginMCTruthTypes
   "_SecondaryDaughterXistarMinus",
   "_Else",
   "_SecondaryFromHF",
-  "_SecondaryFromLF"};
+  "_SecondaryFromLF",
+  "_SecondaryFromPhi"};
 
 /// Distinguished between reconstructed and truth
 enum MCType {
 
@@ -32,6 +32,7 @@
 #include "PWGHF/Utils/utilsEvSelHf.h"
 #include "PWGLF/DataModel/LFStrangenessTables.h"
 
+#include "Common/Core/RecoDecay.h"
 #include "Common/DataModel/Centrality.h"
 #include "Common/DataModel/EventSelection.h"
 #include "Common/DataModel/Multiplicity.h"
@@ -59,6 +60,7 @@
 
 #include <TH1.h>
 #include <TMCProcess.h>
+#include <TPDGCode.h>
 
 #include <array>
 #include <chrono>
@@ -494,8 +496,47 @@ struct HfProducerCharmHadronsV0FemtoDream {
     }
   }
 
+  /// Classify the MC origin of a V0 candidate.
+  /// The heavy-flavour feed-down is identified with RecoDecay::getParticleOrigin, which also
+  /// covers charm- and beauty-hadron decays flagged as physical primaries. The feed-down from
+  /// a phi decay is light-flavour and kept as a separate class, since the relatively long phi
+  /// lifetime makes it a distinct source.
+  /// \param mcParticles table of the MC particles
+  /// \param particleMc MC particle matched to the V0 candidate
+  /// \param isFromPhi true if the K0S comes from a phi decay chain
+  /// \return the origin according to ParticleOriginMCTruth (kPrimary, kMaterial,
+  /// kSecondaryFromHF, kSecondaryFromLF or kSecondaryFromPhi)
+  template <typename T>
+  static int getK0sOrigin(T const& mcParticles, typename T::iterator const& particleMc, const bool isFromPhi)
+  {
+    using namespace o2::aod::femtodreamMCparticle;
+    constexpr int GenFromTransport = -1;
+    const auto process = particleMc.getProcess();
+    // hadronic interaction in the detector material
+    if (process == TMCProcess::kPHInhelastic && particleMc.getGenStatusCode() == GenFromTransport) {
+      return kMaterial;
+    }
+    // charm- or beauty-hadron anywhere in the mother chain (feed-down can be a physical primary)
+    if (RecoDecay::getParticleOrigin(mcParticles, particleMc, false) != RecoDecay::OriginType::None) {
+      return kSecondaryFromHF;
+    }
+    // phi -> K0/K0bar -> K0S is light-flavour feed-down, but kept as a separate class because
+    // the relatively long phi lifetime makes it a distinct source (also flagged as a physical primary)
+    if (isFromPhi) {
+      return kSecondaryFromPhi;
+    }
+    if (particleMc.isPhysicalPrimary()) {
+      return kPrimary;
+    }
+    // weak decay of a light-flavour hadron during the transport
+    if (process == TMCProcess::kPDecay && particleMc.getGenStatusCode() == GenFromTransport && particleMc.has_mothers()) {
+      return kSecondaryFromLF;
+    }
+    return kElse;
+  }
+
   template <typename CollisionType, typename ParticleType>
-  void fillMcParticle(CollisionType const& col, ParticleType const& particle, o2::aod::femtodreamparticle::ParticleType fdparttype)
+  void fillMcParticle(CollisionType const& col, ParticleType const& particle, o2::aod::femtodreamparticle::ParticleType fdparttype, aod::McParticles const& mcParticles)
   {
     if (particle.has_mcParticle()) {
       // get corresponding MC particle and its info
@@ -510,6 +551,19 @@ struct HfProducerCharmHadronsV0FemtoDream {
       float etaPos{-1.f}, etaNeg{-1.f};
       float phiPos{-1.f}, phiNeg{-1.f};
       int pdgCodePos{0}, pdgCodeNeg{0};
+      // phi -> K0/K0bar -> K0s (or phi -> K0s directly); the heavy-flavour feed-down is
+      // handled in getK0sOrigin through RecoDecay::getParticleOrigin
+      bool isFromPhi = false;
+      if (fdparttype == o2::aod::femtodreamparticle::ParticleType::kV0 && !motherparticlesMc.empty()) {
+        const auto absPdgMother = std::abs(motherparticlesMc.front().pdgCode());
+        if (absPdgMother == o2::constants::physics::Pdg::kPhi) {
+          isFromPhi = true;
+        } else if (absPdgMother == kK0) {
+          const auto& grandMothersMc = motherparticlesMc.front().template mothers_as<aod::McParticles>();
+          isFromPhi = !grandMothersMc.empty() && std::abs(grandMothersMc.front().pdgCode()) == o2::constants::physics::Pdg::kPhi;
+        }
+      }
+
       // check pdg code
       // if this fails, the particle is a fake
       if (std::abs(pdgCode) == std::abs(v0PDGCode.value)) {
@@ -518,6 +572,11 @@ struct HfProducerCharmHadronsV0FemtoDream {
         if ((col.has_mcCollision() && (particleMc.mcCollisionId() != col.mcCollisionId())) || !col.has_mcCollision()) {
           particleOrigin = aod::femtodreamMCparticle::ParticleOriginMCTruth::kWrongCollision;
           // check if particle is primary
+        } else if (fdparttype == o2::aod::femtodreamparticle::ParticleType::kV0) {
+          particleOrigin = getK0sOrigin(mcParticles, particleMc, isFromPhi);
+          if (!motherparticlesMc.empty()) {
+            pdgCodeMother = motherparticlesMc.front().pdgCode();
+          }
         } else if (particleMc.isPhysicalPrimary()) {
           particleOrigin = aod::femtodreamMCparticle::ParticleOriginMCTruth::kPrimary;
           // check if particle is secondary
@@ -612,7 +671,7 @@ struct HfProducerCharmHadronsV0FemtoDream {
   }
 
   template <bool IsMc = false, typename V0Type, typename CollisionType, typename TrackType>
-  bool fillV0sForCharmHadron(CollisionType const& col, V0Type const& fullV0s, TrackType const& /*tracks*/)
+  bool fillV0sForCharmHadron(CollisionType const& col, V0Type const& fullV0s, TrackType const& /*tracks*/, aod::McParticles const* mcParticles = nullptr)
   {
     const bool isK0S = (selectionFlagV0 == V0Channel::K0S);
     const bool isLambda = (selectionFlagV0 == V0Channel::Lambda);
@@ -724,15 +783,15 @@ struct HfProducerCharmHadronsV0FemtoDream {
       }
 
       if constexpr (IsMc) {
-        fillMcParticle(col, v0, o2::aod::femtodreamparticle::ParticleType::kV0);
+        fillMcParticle(col, v0, o2::aod::femtodreamparticle::ParticleType::kV0, *mcParticles);
       }
       fIsV0Filled = true;
     }
     return fIsV0Filled;
   }
 
   template <DecayChannel Channel, bool IsMc, bool UseCharmMl, typename V0Type, typename CollisionType, typename TrackType, typename CandType>
-  void fillCharmHadronTable(CollisionType const& col, TrackType const& tracks, V0Type const& fullV0s, CandType const& candidates)
+  void fillCharmHadronTable(CollisionType const& col, TrackType const& tracks, V0Type const& fullV0s, CandType const& candidates, aod::McParticles const* mcParticles = nullptr)
   {
     const auto vtxZ = col.posZ();
     const auto sizeCand = candidates.size();
@@ -1019,7 +1078,7 @@ struct HfProducerCharmHadronsV0FemtoDream {
         fillTable(2, candidate.isSelDstarToD0Pi(), outputMlDstar.at(0), outputMlDstar.at(1), outputMlDstar.at(2));
       }
     }
-    isV0Filled = fillV0sForCharmHadron<IsMc>(col, fullV0s, tracks);
+    isV0Filled = fillV0sForCharmHadron<IsMc>(col, fullV0s, tracks, mcParticles);
     aod::femtodreamcollision::BitMaskType bitV0 = 0;
     if (isV0Filled) {
       bitV0 |= 1 << 0;
@@ -1248,26 +1307,26 @@ struct HfProducerCharmHadronsV0FemtoDream {
                         aod::BCsWithTimestamps const&,
                         FemtoHFMcTracks const& tracks,
                         soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                        aod::McParticles const&,
+                        aod::McParticles const& mcParticles,
                         CandidateD0Mc const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::D0ToPiK, true, false>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::D0ToPiK, true, false>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcD0ToPiK, "Provide Mc for D0ToPiK", false);
 
   void processMcD0ToPiKWithML(FemtoFullCollisionMc const& col,
                               aod::BCsWithTimestamps const&,
                               FemtoHFMcTracks const& tracks,
                               soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                              aod::McParticles const&,
+                              aod::McParticles const& mcParticles,
                               soa::Join<CandidateD0Mc,
                                         aod::HfMlD0> const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::D0ToPiK, true, true>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::D0ToPiK, true, true>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcD0ToPiKWithML, "Provide Mc for D0ToPiK with ml", false);
 
@@ -1307,26 +1366,26 @@ struct HfProducerCharmHadronsV0FemtoDream {
                             aod::BCsWithTimestamps const&,
                             FemtoHFMcTracks const& tracks,
                             soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                            aod::McParticles const&,
+                            aod::McParticles const& mcParticles,
                             CandidateDstarMc const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::DstarToD0Pi, true, false>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::DstarToD0Pi, true, false>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcDstarToD0Pi, "Provide Mc for DstarToD0Pi", false);
 
   void processMcDstarToD0PiWithML(FemtoFullCollisionMc const& col,
                                   aod::BCsWithTimestamps const&,
                                   FemtoHFMcTracks const& tracks,
                                   soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                                  aod::McParticles const&,
+                                  aod::McParticles const& mcParticles,
                                   soa::Join<CandidateDstarMc,
                                             aod::HfMlDstarToD0Pi> const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::DstarToD0Pi, true, true>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::DstarToD0Pi, true, true>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcDstarToD0PiWithML, "Provide Mc for DstarToD0Pi with ml", false);
 
@@ -1368,26 +1427,26 @@ struct HfProducerCharmHadronsV0FemtoDream {
                              aod::BCsWithTimestamps const&,
                              FemtoHFMcTracks const& tracks,
                              soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                             aod::McParticles const&,
+                             aod::McParticles const& mcParticles,
                              CandidateDplusMc const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, false>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, false>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcDplusToPiKPi, "Provide Mc for DplusToPiKPi", false);
 
   void processMcDplusToPiKPiWithML(FemtoFullCollisionMc const& col,
                                    aod::BCsWithTimestamps const&,
                                    FemtoHFMcTracks const& tracks,
                                    soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                                   aod::McParticles const&,
+                                   aod::McParticles const& mcParticles,
                                    soa::Join<CandidateDplusMc,
                                              aod::HfMlDplusToPiKPi> const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, true>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, true>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcDplusToPiKPiWithML, "Provide Mc for DplusToPiKPi with ml", false);
 
@@ -1429,26 +1488,26 @@ struct HfProducerCharmHadronsV0FemtoDream {
                          aod::BCsWithTimestamps const&,
                          FemtoHFMcTracks const& tracks,
                          soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                         aod::McParticles const&,
+                         aod::McParticles const& mcParticles,
                          CandidateLcMc const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::LcToPKPi, true, false>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi, true, false>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcLcToPKPi, "Provide Mc for lctopkpi", false);
 
   void processMcLcToPKPiWithML(FemtoFullCollisionMc const& col,
                                aod::BCsWithTimestamps const&,
                                FemtoHFMcTracks const& tracks,
                                soa::Join<o2::aod::V0Datas, aod::McV0Labels> const& fullV0s,
-                               aod::McParticles const&,
+                               aod::McParticles const& mcParticles,
                                soa::Join<CandidateLcMc,
                                          aod::HfMlLcToPKPi> const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
-    fillCharmHadronTable<DecayChannel::LcToPKPi, true, true>(col, tracks, fullV0s, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi, true, true>(col, tracks, fullV0s, candidates, &mcParticles);
   }
   PROCESS_SWITCH(HfProducerCharmHadronsV0FemtoDream, processMcLcToPKPiWithML, "Provide Mc for lctopkpi with ml", false);