Skip to content

Commit d0b375b

Browse files
ariostasariostas
committed
Fixed memory leak
1 parent 6faa946 commit d0b375b

File tree

2 files changed

+89
-63
lines changed

2 files changed

+89
-63
lines changed

‎RecoTracker/LSTCore/standalone/code/rooutil/treeutil.cc‎

Lines changed: 33 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ void TreeUtil::Init(TTree* treeIn) {
1515
void TreeUtil::GetEntry(unsigned int idx) {
1616
index = idx;
1717
for (auto& pair : data) {
18-
pair.second.isLoaded = false;
18+
pair.second->isLoaded = false;
1919
}
2020
}
2121

@@ -82,46 +82,52 @@ void TreeUtil::loadAllBranches() {
8282
}
8383

8484
template <typename T>
85-
T const& TreeUtil::get(std::string name) {
86-
auto search = data.find(name);
87-
if (search == data.end()) {
85+
const T& TreeUtil::get(const std::string& name) {
86+
auto it = data.find(name);
87+
if (it == data.end()) {
8888
tree->SetMakeClass(1);
89-
TBranch* branch = tree->GetBranch(name.c_str());
90-
if (branch == nullptr)
89+
TBranch* br = tree->GetBranch(name.c_str());
90+
if (!br)
9191
throw std::out_of_range("Branch " + name + " does not exist!");
92-
search = data.emplace(name, std::move(BranchData{branch, false, static_cast<T*>(nullptr)})).first;
93-
branch->SetAddress(&search->second.ptr);
92+
it = data.emplace(name, std::make_unique<BranchDataHolder<T>>(br)).first;
9493
tree->SetMakeClass(0);
9594
}
96-
if (!search->second.isLoaded) {
97-
search->second.branch->GetEntry(index);
98-
search->second.isLoaded = true;
95+
96+
auto* bd = static_cast<BranchDataHolder<T>*>(it->second.get());
97+
if (!bd->isLoaded) {
98+
bd->branch->GetEntry(index);
99+
bd->adoptIfNeeded();
100+
bd->isLoaded = true;
99101
}
100-
return *std::get<T*>(search->second.ptr);
102+
return *static_cast<const T*>(bd->getRaw());
101103
}
102104

103-
short const& TreeUtil::getS(std::string name) { return get<short>(name); }
104-
unsigned short const& TreeUtil::getUS(std::string name) { return get<unsigned short>(name); }
105-
int const& TreeUtil::getI(std::string name) { return get<int>(name); }
106-
unsigned int const& TreeUtil::getU(std::string name) { return get<unsigned int>(name); }
107-
float const& TreeUtil::getF(std::string name) { return get<float>(name); }
108-
std::vector<short> const& TreeUtil::getVS(std::string name) { return get<std::vector<short>>(name); }
109-
std::vector<unsigned short> const& TreeUtil::getVUS(std::string name) { return get<std::vector<unsigned short>>(name); }
110-
std::vector<int> const& TreeUtil::getVI(std::string name) { return get<std::vector<int>>(name); }
111-
std::vector<unsigned int> const& TreeUtil::getVU(std::string name) { return get<std::vector<unsigned int>>(name); }
112-
std::vector<float> const& TreeUtil::getVF(std::string name) { return get<std::vector<float>>(name); }
113-
std::vector<std::vector<short>> const& TreeUtil::getVVS(std::string name) {
105+
short const& TreeUtil::getS(std::string const& name) { return get<short>(name); }
106+
unsigned short const& TreeUtil::getUS(std::string const& name) { return get<unsigned short>(name); }
107+
int const& TreeUtil::getI(std::string const& name) { return get<int>(name); }
108+
unsigned int const& TreeUtil::getU(std::string const& name) { return get<unsigned int>(name); }
109+
float const& TreeUtil::getF(std::string const& name) { return get<float>(name); }
110+
std::vector<short> const& TreeUtil::getVS(std::string const& name) { return get<std::vector<short>>(name); }
111+
std::vector<unsigned short> const& TreeUtil::getVUS(std::string const& name) {
112+
return get<std::vector<unsigned short>>(name);
113+
}
114+
std::vector<int> const& TreeUtil::getVI(std::string const& name) { return get<std::vector<int>>(name); }
115+
std::vector<unsigned int> const& TreeUtil::getVU(std::string const& name) {
116+
return get<std::vector<unsigned int>>(name);
117+
}
118+
std::vector<float> const& TreeUtil::getVF(std::string const& name) { return get<std::vector<float>>(name); }
119+
std::vector<std::vector<short>> const& TreeUtil::getVVS(std::string const& name) {
114120
return get<std::vector<std::vector<short>>>(name);
115121
}
116-
std::vector<std::vector<unsigned short>> const& TreeUtil::getVVUS(std::string name) {
122+
std::vector<std::vector<unsigned short>> const& TreeUtil::getVVUS(std::string const& name) {
117123
return get<std::vector<std::vector<unsigned short>>>(name);
118124
}
119-
std::vector<std::vector<int>> const& TreeUtil::getVVI(std::string name) {
125+
std::vector<std::vector<int>> const& TreeUtil::getVVI(std::string const& name) {
120126
return get<std::vector<std::vector<int>>>(name);
121127
}
122-
std::vector<std::vector<unsigned int>> const& TreeUtil::getVVU(std::string name) {
128+
std::vector<std::vector<unsigned int>> const& TreeUtil::getVVU(std::string const& name) {
123129
return get<std::vector<std::vector<unsigned int>>>(name);
124130
}
125-
std::vector<std::vector<float>> const& TreeUtil::getVVF(std::string name) {
131+
std::vector<std::vector<float>> const& TreeUtil::getVVF(std::string const& name) {
126132
return get<std::vector<std::vector<float>>>(name);
127133
}

‎RecoTracker/LSTCore/standalone/code/rooutil/treeutil.h‎

Lines changed: 56 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -11,54 +11,74 @@
1111

1212
class TreeUtil {
1313
private:
14-
using DataTypes = std::variant<short*,
15-
unsigned short*,
16-
int*,
17-
unsigned int*,
18-
float*,
19-
std::vector<short>*,
20-
std::vector<unsigned short>*,
21-
std::vector<int>*,
22-
std::vector<unsigned int>*,
23-
std::vector<float>*,
24-
std::vector<std::vector<short>>*,
25-
std::vector<std::vector<unsigned short>>*,
26-
std::vector<std::vector<int>>*,
27-
std::vector<std::vector<unsigned int>>*,
28-
std::vector<std::vector<float>>*>;
29-
30-
struct BranchData {
14+
struct BranchDataBase {
3115
TBranch* branch;
32-
bool isLoaded;
33-
DataTypes ptr;
16+
bool isLoaded = false;
17+
explicit BranchDataBase(TBranch* b) : branch(b) {}
18+
virtual ~BranchDataBase() = default;
19+
virtual const void* getRaw() const = 0;
20+
virtual void adoptIfNeeded() = 0;
21+
};
22+
23+
template <typename T>
24+
struct BranchDataHolder : BranchDataBase {
25+
std::unique_ptr<T> buffer;
26+
mutable T* raw = nullptr;
27+
mutable std::unique_ptr<T> owner;
28+
29+
explicit BranchDataHolder(TBranch* b) : BranchDataBase(b) {
30+
if constexpr (std::is_fundamental_v<T>) {
31+
buffer = std::make_unique<T>();
32+
branch->SetAddress(buffer.get());
33+
} else {
34+
branch->SetAddress(&raw);
35+
branch->SetAutoDelete(false);
36+
}
37+
}
38+
39+
const void* getRaw() const override {
40+
if constexpr (std::is_fundamental_v<T>) {
41+
return buffer.get();
42+
} else {
43+
return owner ? static_cast<const void*>(owner.get()) : static_cast<const void*>(raw);
44+
}
45+
}
46+
47+
void adoptIfNeeded() override {
48+
if constexpr (!std::is_fundamental_v<T>) {
49+
if (!owner && raw) {
50+
owner.reset(raw);
51+
}
52+
}
53+
}
3454
};
3555

3656
TTree* tree;
3757
unsigned int index;
38-
std::unordered_map<std::string, BranchData> data;
58+
std::unordered_map<std::string, std::unique_ptr<BranchDataBase>> data;
3959

4060
public:
4161
void Init(TTree* tree);
4262
void GetEntry(unsigned int idx);
4363
static void progress(int nEventsTotal, int nEventsChain);
4464
void loadAllBranches();
4565
template <typename T>
46-
T const& get(std::string name);
47-
short const& getS(std::string name);
48-
unsigned short const& getUS(std::string name);
49-
int const& getI(std::string name);
50-
unsigned int const& getU(std::string name);
51-
float const& getF(std::string name);
52-
std::vector<short> const& getVS(std::string name);
53-
std::vector<unsigned short> const& getVUS(std::string name);
54-
std::vector<int> const& getVI(std::string name);
55-
std::vector<unsigned int> const& getVU(std::string name);
56-
std::vector<float> const& getVF(std::string name);
57-
std::vector<std::vector<short>> const& getVVS(std::string name);
58-
std::vector<std::vector<unsigned short>> const& getVVUS(std::string name);
59-
std::vector<std::vector<int>> const& getVVI(std::string name);
60-
std::vector<std::vector<unsigned int>> const& getVVU(std::string name);
61-
std::vector<std::vector<float>> const& getVVF(std::string name);
66+
T const& get(std::string const& name);
67+
short const& getS(std::string const& name);
68+
unsigned short const& getUS(std::string const& name);
69+
int const& getI(std::string const& name);
70+
unsigned int const& getU(std::string const& name);
71+
float const& getF(std::string const& name);
72+
std::vector<short> const& getVS(std::string const& name);
73+
std::vector<unsigned short> const& getVUS(std::string const& name);
74+
std::vector<int> const& getVI(std::string const& name);
75+
std::vector<unsigned int> const& getVU(std::string const& name);
76+
std::vector<float> const& getVF(std::string const& name);
77+
std::vector<std::vector<short>> const& getVVS(std::string const& name);
78+
std::vector<std::vector<unsigned short>> const& getVVUS(std::string const& name);
79+
std::vector<std::vector<int>> const& getVVI(std::string const& name);
80+
std::vector<std::vector<unsigned int>> const& getVVU(std::string const& name);
81+
std::vector<std::vector<float>> const& getVVF(std::string const& name);
6282
};
6383

6484
#endif

0 commit comments

Comments
 (0)