Skip to content

Common: allow to write c-style arrays in TreeStream #14098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 22, 2025
Merged

Common: allow to write c-style arrays in TreeStream #14098

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 79 additions & 83 deletions Common/Utils/include/CommonUtils/TreeStream.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include <TString.h>
#include <TTree.h>
#include <vector>
#include <type_traits>
#include <concepts>
#include "GPUCommonDef.h"

class TBranch;
Expand All @@ -39,10 +41,79 @@ namespace utils
///
/// See testTreeStream.cxx for functional example
///
namespace details
{
template <typename T>
struct IsTrivialRootType {
static constexpr bool value =
std::is_same_v<T, Float_t> || // Float_t
std::is_same_v<T, Double_t> || // Double_t
std::is_same_v<T, ULong64_t> || std::is_same_v<T, ULong_t> || // ULong64_t or ULong_t
std::is_same_v<T, Long64_t> || std::is_same_v<T, Long_t> || // Long64_t or Long_t
std::is_same_v<T, UInt_t> || // UInt_t
std::is_same_v<T, Int_t> || // Int_t
std::is_same_v<T, UShort_t> || // UShort_t
std::is_same_v<T, Short_t> || // Short_t
std::is_same_v<T, UChar_t> || // UChar_t
std::is_same_v<T, Char_t> || std::is_same_v<T, int8_t> || std::is_same_v<T, Bool_t>; // Char_t, int8_t, or Bool_t
};

template <typename T>
struct IsTrivialRootType<T[]> {
static constexpr bool value = IsTrivialRootType<T>::value;
};

template <typename T, std::size_t N>
struct IsTrivialRootType<T[N]> {
static constexpr bool value = IsTrivialRootType<T>::value;
};

template <typename T>
concept TrivialRootType = IsTrivialRootType<T>::value;

template <typename T>
concept ComplexRootType = !IsTrivialRootType<T>::value;

template <TrivialRootType T>
static constexpr char getRootTypeCode()
{
if constexpr (std::is_array_v<T>) {
return getRootTypeCode<std::remove_all_extents_t<T>>();
} else if constexpr (std::is_same_v<T, Float_t>) {
return 'F';
} else if constexpr (std::is_same_v<T, Double_t>) {
return 'D';
} else if constexpr (std::is_same_v<T, ULong64_t> ||
std::is_same_v<T, ULong_t>) {
return 'l';
} else if constexpr (std::is_same_v<T, Long64_t> ||
std::is_same_v<T, Long_t>) {
return 'L';
} else if constexpr (std::is_same_v<T, UInt_t>) {
return 'i';
} else if constexpr (std::is_same_v<T, Int_t>) {
return 'I';
} else if constexpr (std::is_same_v<T, UShort_t>) {
return 's';
} else if constexpr (std::is_same_v<T, Short_t>) {
return 'S';
} else if constexpr (std::is_same_v<T, UChar_t>) {
return 'b';
} else if constexpr (std::is_same_v<T, Char_t> ||
std::is_same_v<T, int8_t> ||
std::is_same_v<T, Bool_t>) {
return 'B';
} else {
static_assert(false, "unsupported type!");
}
}
} // namespace details

class TreeStream
{
public:
struct TreeDataElement {
int arsize = 1; ///< size of array
char type = 0; ///< type of data element
const TClass* cls = nullptr; ///< data type pointer
const void* ptr = nullptr; ///< pointer to element
Expand All @@ -64,87 +135,10 @@ class TreeStream
void setID(int id) { mID = id; }
int getID() const { return mID; }

TreeStream& operator<<(const Bool_t& b)
{
CheckIn('B', &b);
return *this;
}

TreeStream& operator<<(const Char_t& c)
{
CheckIn('B', &c);
return *this;
}

TreeStream& operator<<(const int8_t& i)
{
CheckIn('B', &i);
return *this;
}

TreeStream& operator<<(const UChar_t& c)
{
CheckIn('b', &c);
return *this;
}

TreeStream& operator<<(const Short_t& h)
{
CheckIn('S', &h);
return *this;
}

TreeStream& operator<<(const UShort_t& h)
{
CheckIn('s', &h);
return *this;
}

TreeStream& operator<<(const Int_t& i)
{
CheckIn('I', &i);
return *this;
}

TreeStream& operator<<(const UInt_t& i)
{
CheckIn('i', &i);
return *this;
}

TreeStream& operator<<(const Long_t& l)
{
CheckIn('L', &l);
return *this;
}

TreeStream& operator<<(const ULong_t& l)
{
CheckIn('l', &l);
return *this;
}

TreeStream& operator<<(const Long64_t& l)
{
CheckIn('L', &l);
return *this;
}

TreeStream& operator<<(const ULong64_t& l)
{
CheckIn('l', &l);
return *this;
}

TreeStream& operator<<(const Float_t& f)
{
CheckIn('F', &f);
return *this;
}

TreeStream& operator<<(const Double_t& d)
template <details::TrivialRootType T>
TreeStream& operator<<(const T& t)
{
CheckIn('D', &d);
CheckIn(details::getRootTypeCode<T>(), &t);
return *this;
}

Expand All @@ -157,7 +151,7 @@ class TreeStream
return *this;
}

template <class T, typename std::enable_if<!std::is_pointer<GPUgeneric() T>::value, bool>::type* = nullptr>
template <details::ComplexRootType T, typename std::enable_if<!std::is_pointer<GPUgeneric() T>::value, bool>::type* = nullptr>
TreeStream& operator<<(const T& obj)
{
CheckIn(&obj);
Expand All @@ -175,6 +169,7 @@ class TreeStream
int mCurrentIndex = 0; ///< index of current element
int mID = -1; ///< identifier of layout
int mNextNameCounter = 0; ///< next name counter
int mNextArraySize = 0; ///< next array size
int mStatus = 0; ///< status of the layout
TString mNextName; ///< name for next entry

Expand All @@ -191,8 +186,7 @@ Int_t TreeStream::CheckIn(const T* obj)
}

if (mCurrentIndex >= static_cast<int>(mElements.size())) {
mElements.emplace_back();
auto& element = mElements.back();
auto& element = mElements.emplace_back();
element.cls = pClass;
TString name = mNextName;
if (name.Length()) {
Expand All @@ -204,6 +198,8 @@ Int_t TreeStream::CheckIn(const T* obj)
}
element.name = name.Data();
element.ptr = obj;
element.arsize = mNextArraySize;
mNextArraySize = 1; // reset
} else {
auto& element = mElements[mCurrentIndex];
if (!element.cls) {
Expand Down
38 changes: 30 additions & 8 deletions Common/Utils/src/TreeStream.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,7 @@ int TreeStream::CheckIn(Char_t type, const void* pointer)
// Insert object

if (mCurrentIndex >= static_cast<int>(mElements.size())) {
mElements.emplace_back();
auto& element = mElements.back();
auto& element = mElements.emplace_back();
element.type = type;
TString name = mNextName;
if (name.Length()) {
Expand All @@ -42,6 +41,8 @@ int TreeStream::CheckIn(Char_t type, const void* pointer)
}
element.name = name.Data();
element.ptr = pointer;
element.arsize = mNextArraySize;
mNextArraySize = 1; // reset
} else {
auto& element = mElements[mCurrentIndex];
if (element.type != type) {
Expand Down Expand Up @@ -89,7 +90,13 @@ void TreeStream::BuildTree()
}

if (element.type > 0) {
TString nameC = TString::Format("%s/%c", name.Data(), element.type);
TString nameC;
if (element.arsize > 1) {
nameC = TString::Format("%s[%d]/%c", name.Data(), element.arsize,
element.type);
} else {
nameC = TString::Format("%s/%c", name.Data(), element.type);
}
br = mTree.Branch(name.Data(), const_cast<void*>(element.ptr), nameC.Data());
if (entriesFilled) {
br->SetAddress(nullptr);
Expand Down Expand Up @@ -148,28 +155,43 @@ TreeStream& TreeStream::Endl()
TreeStream& TreeStream::operator<<(const Char_t* name)
{
// Stream the branch name
//
if (name[0] == '\n') {
return Endl();
}
//

// if tree was already defined ignore
if (mTree.GetEntries() > 0) {
return *this;
}

int arsize = 1;

// check branch name if tree was not
//
Int_t last = 0;
for (last = 0;; last++) {
if (name[last] == 0) {
break;
}
}

if (last > 0 && name[last - 1] == '=') {
mNextName = name;
mNextName[last - 1] = 0;
mNextName[last - 1] = 0; // remove '=' from string
mNextNameCounter = 0;

TString inName{name};
auto brkStaPos = inName.Index('[');

if (brkStaPos != kNPOS) {
auto brkEndPos = inName.Index(']');
if (brkEndPos != kNPOS && brkEndPos > brkStaPos + 1) {
TString size = inName(brkStaPos + 1, brkEndPos - brkStaPos - 1);
arsize = size.Atoi();
mNextName = inName(0, brkStaPos); // use parsed name
}
}
}

mNextArraySize = arsize;

return *this;
}
42 changes: 39 additions & 3 deletions Common/Utils/test/testTreeStream.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,28 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
tstStream << "TrackTreeR"
<< "id=" << i << "x=" << x << "track=" << trc << "\n";
}

// test for c-arrays
int iArray[6] = {1, 2, 3, 4, 5, 6};
float fArray[6] = {1.1f, 2.2f, 3.3f, 4.4f, 5.5f, 6.6f};
for (int i{0}; i < nit; ++i) {
for (int j{0}; j < 6; ++j) {
iArray[j] += i;
fArray[j] += (float)i;
}
tstStream << "ArrayTree"
<< "id=" << i
<< "iArray[6]=" << iArray
<< "fArray[6]=" << fArray
<< "\n";
}

// on destruction of tstTreem the trees will be stored, but we can also force it by
tstStream.Close();
}
//
LOG(info) << "Testing reading back tree maid by the TreeStream ";
// read back tracks
LOG(info) << "Testing reading back tree made by the TreeStream ";
// read back tracks and arrays
{
TFile inpf(outFName.data());
BOOST_CHECK(!inpf.IsZombie());
Expand All @@ -80,6 +96,27 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
trc->printParam();
BOOST_CHECK(std::abs(x - trc->getX()) < 1e-4);
}

// check arrays
tree = (TTree*)inpf.GetObjectChecked("ArrayTree", "TTree");
BOOST_CHECK(tree);
nent = tree->GetEntries();
BOOST_CHECK(nent == nit);
int iArray[6];
float fArray[6];
BOOST_CHECK(!tree->SetBranchAddress("id", &id));
BOOST_CHECK(!tree->SetBranchAddress("iArray", iArray));
BOOST_CHECK(!tree->SetBranchAddress("fArray", fArray));
for (int i = 0; i < nit; i++) {
BOOST_CHECK(tree->GetEntry(i) > 0);
BOOST_CHECK(id == i);
for (int j = 0; j < 6; j++) {
BOOST_CHECK(iArray[j] == (1 + j + i * (i + 1) / 2));
}
for (int j = 0; j < 6; j++) {
BOOST_CHECK_CLOSE(fArray[j], (1.f + j + i * (i + 1) / 2.f + 0.1 * (j + 1)), 1e-5);
}
}
}

LOG(info) << "Testing loading tree via RootChain";
Expand All @@ -104,7 +141,6 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
nit = 1000;
BOOST_CHECK(UnitTestSparse(0.5, nit));
BOOST_CHECK(UnitTestSparse(0.1, nit));
//
}

//_________________________________________________
Expand Down