[ntuple] add support for leaf count arrays in classes

tree/ntuple/inc/ROOT/RField.hxx

@@ -36,8 +36,10 @@
 #include <vector>
 
 class TClass;
+class TDataMember;
 class TEnum;
 class TObject;
+class TRealData;
 class TVirtualStreamerInfo;
 
 namespace ROOT {
@@ -197,6 +199,9 @@ private:
    /// Fields may not have an on-disk representation (e.g., when inserted by schema evolution), in which case the passed
    /// field descriptor is nullptr.
    std::vector<const TSchemaRule *> FindRules(const ROOT::RFieldDescriptor *fieldDesc);
+   /// Checks if the data member dm in fClass is a leaf count array. If so, returns a pointer to the data member
+   /// corresponding to the count leaf (which may be in a base class).
+   TRealData *IsLeafCountArray(const TDataMember &dm) const;
 
 protected:
    std::unique_ptr<RFieldBase> CloneImpl(std::string_view newName) const final;

tree/ntuple/inc/ROOT/RField/RFieldSequenceContainer.hxx

@@ -417,6 +417,51 @@ public:
    void AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const final;
 };
 
+/**
+\class ROOT::RLeafCountArrayField
+\brief A field for an RVec field as a class member on disk that is represented as leaf count array` in memory.
+\ingroup NTuple
+*/
+class RLeafCountArrayField final : public RFieldBase {
+   friend class RClassField; // only class fields are able to construct this
+
+private:
+   std::size_t fItemSize{0};                  ///< The size of a child field's item
+   ROOT::Internal::RColumnIndex fNWritten{0}; ///< Number of items written in the current cluster
+   std::ptrdiff_t fCountLeafDelta{0};         ///< Offset delta of the count leaf int member relative to the array
+   bool fHasPersistentCountLeaf{true};        ///< The count leaf may be only transient in RNTuple
+
+   RLeafCountArrayField(std::string_view fieldName, std::unique_ptr<RFieldBase> itemField,
+                        std::ptrdiff_t countLeafDelta, bool hasPersistentCountLeaf);
+
+protected:
+   std::unique_ptr<RFieldBase> CloneImpl(std::string_view newName) const final;
+
+   const RColumnRepresentations &GetColumnRepresentations() const final;
+   void GenerateColumns() final;
+   void GenerateColumns(const ROOT::RNTupleDescriptor &desc) final;
+
+   void ConstructValue(void *where) const final { *reinterpret_cast<void **>(where) = nullptr; }
+
+   std::size_t AppendImpl(const void *from) final;
+   void ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to) final;
+
+   void ReconcileOnDiskField(const RNTupleDescriptor &desc) final;
+
+   void CommitClusterImpl() final { fNWritten = 0; }
+
+public:
+   RLeafCountArrayField(RLeafCountArrayField &&other) = default;
+   RLeafCountArrayField &operator=(RLeafCountArrayField &&other) = default;
+   ~RLeafCountArrayField() final = default;
+
+   std::vector<RValue> SplitValue(const RValue &value) const final;
+
+   size_t GetValueSize() const final { return sizeof(void *); }
+   size_t GetAlignment() const final { return std::alignment_of<void *>(); }
+   void AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const final;
+};
+
 } // namespace ROOT
 
 #endif

tree/ntuple/inc/ROOT/RFieldVisitor.hxx

@@ -91,6 +91,7 @@ public:
    virtual void VisitVectorField(const ROOT::RVectorField &field) { VisitField(field); }
    virtual void VisitVectorBoolField(const ROOT::RField<std::vector<bool>> &field) { VisitField(field); }
    virtual void VisitRVecField(const ROOT::RRVecField &field) { VisitField(field); }
+   virtual void VisitLeafCountArrayField(const RLeafCountArrayField &field) { VisitField(field); }
 }; // class RFieldVisitor
 
 } // namespace Detail
@@ -245,6 +246,7 @@ public:
    void VisitVectorField(const ROOT::RVectorField &field) final;
    void VisitVectorBoolField(const ROOT::RField<std::vector<bool>> &field) final;
    void VisitRVecField(const ROOT::RRVecField &field) final;
+   void VisitLeafCountArrayField(const RLeafCountArrayField &field) final;
    void VisitBitsetField(const ROOT::RBitsetField &field) final;
    void VisitNullableField(const ROOT::RNullableField &field) final;
    void VisitEnumField(const ROOT::REnumField &field) final;

tree/ntuple/src/RFieldMeta.cxx

@@ -205,7 +205,19 @@ ROOT::RClassField::RClassField(std::string_view fieldName, TClass *classp)
          }
       }
 
-      auto subField = RFieldBase::Create(dataMember->GetName(), typeName).Unwrap();
+      const auto memberName = dataMember->GetName();
+      std::unique_ptr<RFieldBase> subField;
+
+      if (auto realMember = IsLeafCountArray(*dataMember)) {
+         assert(typeName.length() > 0);
+         assert(*typeName.rbegin() == '*');
+         auto itemField = RFieldBase::Create("_0", typeName.substr(0, typeName.length() - 1)).Unwrap();
+         const std::ptrdiff_t offsetDelta = realMember->GetThisOffset() - dataMember->GetOffset();
+         subField = std::unique_ptr<RLeafCountArrayField>(new RLeafCountArrayField(
+            memberName, std::move(itemField), offsetDelta, true /* fHasPersistentCountLeaf */));
+      } else {
+         subField = RFieldBase::Create(memberName, typeName).Unwrap();
+      }
 
       const auto normTypeName = ROOT::Internal::GetNormalizedUnresolvedTypeName(origTypeName);
       if (normTypeName == subField->GetTypeName()) {
@@ -239,6 +251,43 @@ void ROOT::RClassField::Attach(std::unique_ptr<RFieldBase> child, RSubFieldInfo
    RFieldBase::Attach(std::move(child));
 }
 
+TRealData *ROOT::RClassField::IsLeafCountArray(const TDataMember &dm) const
+{
+   if (!(dm.Property() & kIsPointer) || !dm.GetTitle())
+      return nullptr;
+
+   std::string title = dm.GetTitle();
+   if (title.length() == 0 || title[0] != '[')
+      return nullptr;
+
+   auto idxRight = title.find_first_of("]", 1);
+   if (idxRight == std::string::npos)
+      return nullptr;
+
+   auto countLeaf = title.substr(1, idxRight - 1);
+   for (auto realMember : ROOT::Detail::TRangeStaticCast<TRealData>(*fClass->GetListOfRealData())) {
+      if (countLeaf != realMember->GetName())
+         continue;
+
+      const auto dmCountLeaf = realMember->GetDataMember();
+
+      const auto dtCountLeaf = dmCountLeaf->GetDataType();
+      if (!dtCountLeaf || ((dtCountLeaf->GetType() != kInt_t) && (dtCountLeaf->GetType() != kUInt_t))) {
+         throw ROOT::RException(
+            R__FAIL(std::string("invalid count leaf type: ") + GetTypeName() + "." + realMember->GetName()));
+      }
+
+      if (realMember->GetThisOffset() >= dm.GetOffset()) {
+         throw ROOT::RException(R__FAIL(std::string("count leaf member defined after array: ") + GetTypeName() + "." +
+                                        realMember->GetName()));
+      }
+
+      return realMember;
+   }
+
+   throw ROOT::RException(R__FAIL(std::string("invalid count leaf name in: ") + GetTypeName() + "." + dm.GetName()));
+}
+
 std::vector<const ROOT::TSchemaRule *> ROOT::RClassField::FindRules(const ROOT::RFieldDescriptor *fieldDesc)
 {
    ROOT::Detail::TSchemaRuleSet::TMatches rules;

tree/ntuple/src/RFieldSequenceContainer.cxx

@@ -1038,3 +1038,129 @@ void ROOT::RArrayAsVectorField::AcceptVisitor(ROOT::Detail::RFieldVisitor &visit
 {
    visitor.VisitArrayAsVectorField(*this);
 }
+
+//------------------------------------------------------------------------------
+
+ROOT::RLeafCountArrayField::RLeafCountArrayField(std::string_view fieldName, std::unique_ptr<RFieldBase> itemField,
+                                                 std::ptrdiff_t countLeafDelta, bool hasPersistentCountLeaf)
+   : ROOT::RFieldBase(fieldName, "ROOT::VecOps::RVec<" + itemField->GetTypeName() + ">",
+                      ROOT::ENTupleStructure::kCollection, false /* isSimple */),
+     fItemSize(itemField->GetValueSize()),
+     fCountLeafDelta(countLeafDelta),
+     fHasPersistentCountLeaf(hasPersistentCountLeaf)
+{
+   static_assert(sizeof(int) == sizeof(std::int32_t)); // implicit assumption in this field
+   assert(fCountLeafDelta < 0);                        // the count leaf must be defined before the array
+
+   if (!(itemField->GetTraits() & kTraitTriviallyDestructible)) {
+      throw RException(R__FAIL("RLeafCountArrayField only supports trivially destructible item types"));
+   }
+   // Note that we expect the array pointer to be deleted by the user.
+   fTraits |= kTraitTriviallyDestructible | (itemField->GetTraits() & kTraitTriviallyConstructible);
+   Attach(std::move(itemField));
+}
+
+std::unique_ptr<ROOT::RFieldBase> ROOT::RLeafCountArrayField::CloneImpl(std::string_view newName) const
+{
+   auto newItemField = fSubfields[0]->Clone(fSubfields[0]->GetFieldName());
+   return std::unique_ptr<RLeafCountArrayField>(
+      new RLeafCountArrayField(newName, std::move(newItemField), fCountLeafDelta, fHasPersistentCountLeaf));
+}
+
+const ROOT::RFieldBase::RColumnRepresentations &ROOT::RLeafCountArrayField::GetColumnRepresentations() const
+{
+   static RColumnRepresentations representations({{ENTupleColumnType::kSplitIndex64},
+                                                  {ENTupleColumnType::kIndex64},
+                                                  {ENTupleColumnType::kSplitIndex32},
+                                                  {ENTupleColumnType::kIndex32}},
+                                                 {});
+   return representations;
+}
+
+void ROOT::RLeafCountArrayField::GenerateColumns()
+{
+   GenerateColumnsImpl<ROOT::Internal::RColumnIndex>();
+}
+
+void ROOT::RLeafCountArrayField::GenerateColumns(const ROOT::RNTupleDescriptor &desc)
+{
+   GenerateColumnsImpl<ROOT::Internal::RColumnIndex>(desc);
+}
+
+std::size_t ROOT::RLeafCountArrayField::AppendImpl(const void *from)
+{
+   auto arrPtr = *reinterpret_cast<const unsigned char *const *>(from);
+   auto count = *reinterpret_cast<const std::uint32_t *>(static_cast<const unsigned char *>(from) + fCountLeafDelta);
+   std::size_t nbytes = 0;
+
+   if (fSubfields[0]->IsSimple() && count) {
+      GetPrincipalColumnOf(*fSubfields[0])->AppendV(arrPtr, count);
+      nbytes += count * GetPrincipalColumnOf(*fSubfields[0])->GetElement()->GetPackedSize();
+   } else {
+      for (unsigned i = 0; i < count; ++i) {
+         nbytes += CallAppendOn(*fSubfields[0], arrPtr + (i * fItemSize));
+      }
+   }
+
+   fNWritten += count;
+   fPrincipalColumn->Append(&fNWritten);
+   return nbytes + fPrincipalColumn->GetElement()->GetPackedSize();
+}
+
+void ROOT::RLeafCountArrayField::ReadGlobalImpl(ROOT::NTupleSize_t globalIndex, void *to)
+{
+   auto typedValue = reinterpret_cast<unsigned char **>(to);
+   auto countPtr = reinterpret_cast<std::uint32_t *>(static_cast<unsigned char *>(to) + fCountLeafDelta);
+
+   ROOT::NTupleSize_t nItems;
+   RNTupleLocalIndex collectionStart;
+   fPrincipalColumn->GetCollectionInfo(globalIndex, &collectionStart, &nItems);
+   if (nItems > std::numeric_limits<std::uint32_t>::max()) {
+      throw RException(R__FAIL("count leaf overflow"));
+   }
+
+   if (fHasPersistentCountLeaf) {
+      if (nItems != *countPtr) {
+         throw RException(R__FAIL("count leaf value different from collection size on disk"));
+      }
+   } else {
+      *countPtr = nItems;
+   }
+
+   operator delete(*typedValue);
+   *typedValue = static_cast<unsigned char *>(operator new(nItems * fItemSize));
+
+   if (fSubfields[0]->IsSimple()) {
+      if (nItems)
+         GetPrincipalColumnOf(*fSubfields[0])->ReadV(collectionStart, nItems, *typedValue);
+   } else {
+      for (std::size_t i = 0; i < nItems; ++i) {
+         CallReadOn(*fSubfields[0], collectionStart + i, *typedValue + (i * fItemSize));
+      }
+   }
+}
+
+void ROOT::RLeafCountArrayField::ReconcileOnDiskField(const RNTupleDescriptor &desc)
+{
+   EnsureMatchingOnDiskField(desc, kDiffTypeName).ThrowOnError();
+}
+
+std::vector<ROOT::RFieldBase::RValue> ROOT::RLeafCountArrayField::SplitValue(const RValue &value) const
+{
+   auto arrPtr = *reinterpret_cast<unsigned char **>(value.GetPtr<void>().get());
+   auto count =
+      *reinterpret_cast<std::uint32_t *>(static_cast<unsigned char *>(value.GetPtr<void>().get()) + fCountLeafDelta);
+
+   std::vector<RValue> result;
+   result.reserve(count);
+   for (std::uint32_t i = 0; i < count; ++i) {
+      result.emplace_back(
+         fSubfields[0]->BindValue(std::shared_ptr<void>(value.GetPtr<void>(), arrPtr + i * fItemSize)));
+   }
+   return result;
+}
+
+void ROOT::RLeafCountArrayField::AcceptVisitor(ROOT::Detail::RFieldVisitor &visitor) const
+{
+   visitor.VisitLeafCountArrayField(*this);
+}

tree/ntuple/src/RFieldVisitor.cxx

@@ -412,6 +412,11 @@ void ROOT::Internal::RPrintValueVisitor::VisitRVecField(const ROOT::RRVecField &
    PrintCollection(field);
 }
 
+void ROOT::Internal::RPrintValueVisitor::VisitLeafCountArrayField(const ROOT::RLeafCountArrayField &field)
+{
+   PrintCollection(field);
+}
+
 //---------------------------- RNTupleFormatter --------------------------------
 
 std::string ROOT::Internal::RNTupleFormatter::FitString(const std::string &str, int availableSpace)

tree/ntuple/test/CustomStruct.hxx

@@ -460,4 +460,28 @@ struct ExampleMC {
 };
 } // namespace v2
 
+struct LeafCountInClassBase {
+   Int_t fSize = 0;
+};
+
+struct LeafCountInClass : public LeafCountInClassBase {
+   unsigned char *fPayload1 = nullptr; //[fSize];
+   unsigned char *fPayload2 = nullptr; //[fSize];
+};
+
+struct LeafCountInClassFail1 {
+   int fSize = 0;
+   unsigned char *fPayload = nullptr; //[fTypo];
+};
+
+struct LeafCountInClassFail2 {
+   char fSize = 0;                    // wrong count leaf type
+   unsigned char *fPayload = nullptr; //[fSize];
+};
+
+struct LeafCountInClassFail3 {
+   unsigned char *fPayload = nullptr; //[fSize];
+   int fSize = 0;                     // declared after the array
+};
+
 #endif

tree/ntuple/test/CustomStructLinkDef.h

@@ -175,4 +175,10 @@
 #pragma read sourceClass = "v1::ExampleMC" source = "v1::Vector3D fSpin" version="[1-]" targetClass = \
    "v2::ExampleMC" target = "fHelicity" code = "{ fHelicity = onfile.fSpin.fZ; }"
 
+#pragma link C++ class LeafCountInClassBase+;
+#pragma link C++ class LeafCountInClass+;
+#pragma link C++ class LeafCountInClassFail1+;
+#pragma link C++ class LeafCountInClassFail2+;
+#pragma link C++ class LeafCountInClassFail3+;
+
 #endif