Skip to content

Commit 5147a9c

Browse files
f3schf3sch
committed
Common: allow to write c-style arrays in TreeStream
Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>
1 parent 5a7fba3 commit 5147a9c

File tree

3 files changed

+148
-94
lines changed

3 files changed

+148
-94
lines changed

‎Common/Utils/include/CommonUtils/TreeStream.h‎

Lines changed: 79 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
#include <TString.h>
2020
#include <TTree.h>
2121
#include <vector>
22+
#include <type_traits>
23+
#include <concepts>
2224
#include "GPUCommonDef.h"
2325

2426
class TBranch;
@@ -39,10 +41,79 @@ namespace utils
3941
///
4042
/// See testTreeStream.cxx for functional example
4143
///
44+
namespace details
45+
{
46+
template <typename T>
47+
struct IsTrivialRootType {
48+
static constexpr bool value =
49+
std::is_same_v<T, Float_t> || // Float_t
50+
std::is_same_v<T, Double_t> || // Double_t
51+
std::is_same_v<T, ULong64_t> || std::is_same_v<T, ULong_t> || // ULong64_t or ULong_t
52+
std::is_same_v<T, Long64_t> || std::is_same_v<T, Long_t> || // Long64_t or Long_t
53+
std::is_same_v<T, UInt_t> || // UInt_t
54+
std::is_same_v<T, Int_t> || // Int_t
55+
std::is_same_v<T, UShort_t> || // UShort_t
56+
std::is_same_v<T, Short_t> || // Short_t
57+
std::is_same_v<T, UChar_t> || // UChar_t
58+
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
59+
};
60+
61+
template <typename T>
62+
struct IsTrivialRootType<T[]> {
63+
static constexpr bool value = IsTrivialRootType<T>::value;
64+
};
65+
66+
template <typename T, std::size_t N>
67+
struct IsTrivialRootType<T[N]> {
68+
static constexpr bool value = IsTrivialRootType<T>::value;
69+
};
70+
71+
template <typename T>
72+
concept TrivialRootType = IsTrivialRootType<T>::value;
73+
74+
template <typename T>
75+
concept ComplexRootType = !IsTrivialRootType<T>::value;
76+
77+
template <TrivialRootType T>
78+
static constexpr char getRootTypeCode()
79+
{
80+
if constexpr (std::is_array_v<T>) {
81+
return getRootTypeCode<std::remove_all_extents_t<T>>();
82+
} else if constexpr (std::is_same_v<T, Float_t>) {
83+
return 'F';
84+
} else if constexpr (std::is_same_v<T, Double_t>) {
85+
return 'D';
86+
} else if constexpr (std::is_same_v<T, ULong64_t> ||
87+
std::is_same_v<T, ULong_t>) {
88+
return 'l';
89+
} else if constexpr (std::is_same_v<T, Long64_t> ||
90+
std::is_same_v<T, Long_t>) {
91+
return 'L';
92+
} else if constexpr (std::is_same_v<T, UInt_t>) {
93+
return 'i';
94+
} else if constexpr (std::is_same_v<T, Int_t>) {
95+
return 'I';
96+
} else if constexpr (std::is_same_v<T, UShort_t>) {
97+
return 's';
98+
} else if constexpr (std::is_same_v<T, Short_t>) {
99+
return 'S';
100+
} else if constexpr (std::is_same_v<T, UChar_t>) {
101+
return 'b';
102+
} else if constexpr (std::is_same_v<T, Char_t> ||
103+
std::is_same_v<T, int8_t> ||
104+
std::is_same_v<T, Bool_t>) {
105+
return 'B';
106+
} else {
107+
static_assert(false, "unsupported type!");
108+
}
109+
}
110+
} // namespace details
111+
42112
class TreeStream
43113
{
44114
public:
45115
struct TreeDataElement {
116+
int arsize = 1; ///< size of array
46117
char type = 0; ///< type of data element
47118
const TClass* cls = nullptr; ///< data type pointer
48119
const void* ptr = nullptr; ///< pointer to element
@@ -64,87 +135,10 @@ class TreeStream
64135
void setID(int id) { mID = id; }
65136
int getID() const { return mID; }
66137

67-
TreeStream& operator<<(const Bool_t& b)
68-
{
69-
CheckIn('B', &b);
70-
return *this;
71-
}
72-
73-
TreeStream& operator<<(const Char_t& c)
74-
{
75-
CheckIn('B', &c);
76-
return *this;
77-
}
78-
79-
TreeStream& operator<<(const int8_t& i)
80-
{
81-
CheckIn('B', &i);
82-
return *this;
83-
}
84-
85-
TreeStream& operator<<(const UChar_t& c)
86-
{
87-
CheckIn('b', &c);
88-
return *this;
89-
}
90-
91-
TreeStream& operator<<(const Short_t& h)
92-
{
93-
CheckIn('S', &h);
94-
return *this;
95-
}
96-
97-
TreeStream& operator<<(const UShort_t& h)
98-
{
99-
CheckIn('s', &h);
100-
return *this;
101-
}
102-
103-
TreeStream& operator<<(const Int_t& i)
104-
{
105-
CheckIn('I', &i);
106-
return *this;
107-
}
108-
109-
TreeStream& operator<<(const UInt_t& i)
110-
{
111-
CheckIn('i', &i);
112-
return *this;
113-
}
114-
115-
TreeStream& operator<<(const Long_t& l)
116-
{
117-
CheckIn('L', &l);
118-
return *this;
119-
}
120-
121-
TreeStream& operator<<(const ULong_t& l)
122-
{
123-
CheckIn('l', &l);
124-
return *this;
125-
}
126-
127-
TreeStream& operator<<(const Long64_t& l)
128-
{
129-
CheckIn('L', &l);
130-
return *this;
131-
}
132-
133-
TreeStream& operator<<(const ULong64_t& l)
134-
{
135-
CheckIn('l', &l);
136-
return *this;
137-
}
138-
139-
TreeStream& operator<<(const Float_t& f)
140-
{
141-
CheckIn('F', &f);
142-
return *this;
143-
}
144-
145-
TreeStream& operator<<(const Double_t& d)
138+
template <details::TrivialRootType T>
139+
TreeStream& operator<<(const T& t)
146140
{
147-
CheckIn('D', &d);
141+
CheckIn(details::getRootTypeCode<T>(), &t);
148142
return *this;
149143
}
150144

@@ -157,7 +151,7 @@ class TreeStream
157151
return *this;
158152
}
159153

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

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

193188
if (mCurrentIndex >= static_cast<int>(mElements.size())) {
194-
mElements.emplace_back();
195-
auto& element = mElements.back();
189+
auto& element = mElements.emplace_back();
196190
element.cls = pClass;
197191
TString name = mNextName;
198192
if (name.Length()) {
@@ -204,6 +198,8 @@ Int_t TreeStream::CheckIn(const T* obj)
204198
}
205199
element.name = name.Data();
206200
element.ptr = obj;
201+
element.arsize = mNextArraySize;
202+
mNextArraySize = 1; // reset
207203
} else {
208204
auto& element = mElements[mCurrentIndex];
209205
if (!element.cls) {

‎Common/Utils/src/TreeStream.cxx‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,7 @@ int TreeStream::CheckIn(Char_t type, const void* pointer)
2929
// Insert object
3030

3131
if (mCurrentIndex >= static_cast<int>(mElements.size())) {
32-
mElements.emplace_back();
33-
auto& element = mElements.back();
32+
auto& element = mElements.emplace_back();
3433
element.type = type;
3534
TString name = mNextName;
3635
if (name.Length()) {
@@ -42,6 +41,8 @@ int TreeStream::CheckIn(Char_t type, const void* pointer)
4241
}
4342
element.name = name.Data();
4443
element.ptr = pointer;
44+
element.arsize = mNextArraySize;
45+
mNextArraySize = 1; // reset
4546
} else {
4647
auto& element = mElements[mCurrentIndex];
4748
if (element.type != type) {
@@ -89,7 +90,13 @@ void TreeStream::BuildTree()
8990
}
9091

9192
if (element.type > 0) {
92-
TString nameC = TString::Format("%s/%c", name.Data(), element.type);
93+
TString nameC;
94+
if (element.arsize > 1) {
95+
nameC = TString::Format("%s[%d]/%c", name.Data(), element.arsize,
96+
element.type);
97+
} else {
98+
nameC = TString::Format("%s/%c", name.Data(), element.type);
99+
}
93100
br = mTree.Branch(name.Data(), const_cast<void*>(element.ptr), nameC.Data());
94101
if (entriesFilled) {
95102
br->SetAddress(nullptr);
@@ -148,28 +155,43 @@ TreeStream& TreeStream::Endl()
148155
TreeStream& TreeStream::operator<<(const Char_t* name)
149156
{
150157
// Stream the branch name
151-
//
152158
if (name[0] == '\n') {
153159
return Endl();
154160
}
155-
//
161+
156162
// if tree was already defined ignore
157163
if (mTree.GetEntries() > 0) {
158164
return *this;
159165
}
166+
167+
int arsize = 1;
168+
160169
// check branch name if tree was not
161-
//
162170
Int_t last = 0;
163171
for (last = 0;; last++) {
164172
if (name[last] == 0) {
165173
break;
166174
}
167175
}
168-
169176
if (last > 0 && name[last - 1] == '=') {
170177
mNextName = name;
171-
mNextName[last - 1] = 0;
178+
mNextName[last - 1] = 0; // remove '=' from string
172179
mNextNameCounter = 0;
180+
181+
TString inName{name};
182+
auto brkStaPos = inName.Index('[');
183+
184+
if (brkStaPos != kNPOS) {
185+
auto brkEndPos = inName.Index(']');
186+
if (brkEndPos != kNPOS && brkEndPos > brkStaPos + 1) {
187+
TString size = inName(brkStaPos + 1, brkEndPos - brkStaPos - 1);
188+
arsize = size.Atoi();
189+
mNextName = inName(0, brkStaPos); // use parsed name
190+
}
191+
}
173192
}
193+
194+
mNextArraySize = arsize;
195+
174196
return *this;
175197
}

‎Common/Utils/test/testTreeStream.cxx‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,28 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
5353
tstStream << "TrackTreeR"
5454
<< "id=" << i << "x=" << x << "track=" << trc << "\n";
5555
}
56+
57+
// test for c-arrays
58+
int iArray[6] = {1, 2, 3, 4, 5, 6};
59+
float fArray[6] = {1.1f, 2.2f, 3.3f, 4.4f, 5.5f, 6.6f};
60+
for (int i{0}; i < nit; ++i) {
61+
for (int j{0}; j < 6; ++j) {
62+
iArray[j] += i;
63+
fArray[j] += (float)i;
64+
}
65+
tstStream << "ArrayTree"
66+
<< "id=" << i
67+
<< "iArray[6]=" << iArray
68+
<< "fArray[6]=" << fArray
69+
<< "\n";
70+
}
71+
5672
// on destruction of tstTreem the trees will be stored, but we can also force it by
5773
tstStream.Close();
5874
}
5975
//
60-
LOG(info) << "Testing reading back tree maid by the TreeStream ";
61-
// read back tracks
76+
LOG(info) << "Testing reading back tree made by the TreeStream ";
77+
// read back tracks and arrays
6278
{
6379
TFile inpf(outFName.data());
6480
BOOST_CHECK(!inpf.IsZombie());
@@ -80,6 +96,27 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
8096
trc->printParam();
8197
BOOST_CHECK(std::abs(x - trc->getX()) < 1e-4);
8298
}
99+
100+
// check arrays
101+
tree = (TTree*)inpf.GetObjectChecked("ArrayTree", "TTree");
102+
BOOST_CHECK(tree);
103+
nent = tree->GetEntries();
104+
BOOST_CHECK(nent == nit);
105+
int iArray[6];
106+
float fArray[6];
107+
BOOST_CHECK(!tree->SetBranchAddress("id", &id));
108+
BOOST_CHECK(!tree->SetBranchAddress("iArray", iArray));
109+
BOOST_CHECK(!tree->SetBranchAddress("fArray", fArray));
110+
for (int i = 0; i < nit; i++) {
111+
BOOST_CHECK(tree->GetEntry(i) > 0);
112+
BOOST_CHECK(id == i);
113+
for (int j = 0; j < 6; j++) {
114+
BOOST_CHECK(iArray[j] == (1 + j + i * (i + 1) / 2));
115+
}
116+
for (int j = 0; j < 6; j++) {
117+
BOOST_CHECK_CLOSE(fArray[j], (1.f + j + i * (i + 1) / 2.f + 0.1 * (j + 1)), 1e-5);
118+
}
119+
}
83120
}
84121

85122
LOG(info) << "Testing loading tree via RootChain";
@@ -104,7 +141,6 @@ BOOST_AUTO_TEST_CASE(TreeStream_test)
104141
nit = 1000;
105142
BOOST_CHECK(UnitTestSparse(0.5, nit));
106143
BOOST_CHECK(UnitTestSparse(0.1, nit));
107-
//
108144
}
109145

110146
//_________________________________________________

0 commit comments

Comments
 (0)