Skip to content

Commit cde8d52

Browse files
silverweedsilverweed
committed
[io] Change TFile::WalkTKeys() to return an iterable
1 parent 45adcda commit cde8d52

File tree

2 files changed

+143
-78
lines changed

2 files changed

+143
-78
lines changed

‎io/io/inc/TFile.h

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include "TDirectoryFile.h"
3030
#include "TUrl.h"
3131
#include "ROOT/RConcurrentHashColl.hxx"
32+
#include <optional>
3233

3334
// Not a part of TFile interface; provide a forward declaration instead of #include.
3435
// #ifndef R__LESS_INCLUDES
@@ -78,11 +79,59 @@ struct TKeyMapNode {
7879
TKeyMapNode() = default;
7980
TKeyMapNode(std::uint64_t addr, EType type, std::uint32_t len = 0) : fAddr(addr), fType(type), fLen(len) {}
8081
};
82+
83+
class TKeyMapIterable {
84+
TFile *fFile;
85+
86+
public:
87+
class TIterator {
88+
TFile *fFile;
89+
std::optional<TKeyMapNode> fCur;
90+
std::uint64_t fCurAddr;
91+
92+
std::optional<TKeyMapNode> Next();
93+
94+
void Advance()
95+
{
96+
fCur = Next();
97+
// If we encounter an error key, skip to the end of the file so we don't try reading additional ones.
98+
if (!fCur || fCur->fType == TKeyMapNode::kError)
99+
fCurAddr = -1;
100+
}
101+
102+
public:
103+
using iterator = TIterator;
104+
using iterator_category = std::forward_iterator_tag;
105+
using difference_type = std::ptrdiff_t;
106+
using value_type = TKeyMapNode;
107+
using pointer = value_type *;
108+
using reference = value_type &;
109+
110+
TIterator(TFile *file, std::uint64_t addr);
111+
112+
iterator operator++()
113+
{
114+
Advance();
115+
return *this;
116+
}
117+
reference operator*() { return fCur.value(); }
118+
pointer operator->() { return fCur ? &*fCur : nullptr; }
119+
bool operator!=(const iterator &rh) const { return !(*this == rh); }
120+
bool operator==(const iterator &rh) const { return fFile == rh.fFile && fCurAddr == rh.fCurAddr; }
121+
};
122+
123+
explicit TKeyMapIterable(TFile *file) : fFile(file) {}
124+
125+
TIterator begin() const { return TIterator(fFile, 0); }
126+
TIterator end() const { return TIterator(fFile, -1); }
127+
};
128+
81129
} // namespace ROOT::Detail
82130

83131
class TFile : public TDirectoryFile {
84132
friend class TDirectoryFile;
85133
friend class TFilePrefetch;
134+
friend class ROOT::Detail::TKeyMapIterable::TIterator;
86135
// TODO: We need to make sure only one TBasket is being written at a time
87136
// if we are writing multiple baskets in parallel.
88137
#ifdef R__USE_IMT
@@ -295,7 +344,7 @@ class TFile : public TDirectoryFile {
295344
Option_t *option="new"); // *MENU*
296345

297346
/// Traverses all TKeys in the TFile and returns information about them.
298-
std::vector<ROOT::Detail::TKeyMapNode> WalkTKeys();
347+
ROOT::Detail::TKeyMapIterable WalkTKeys();
299348

300349
virtual void Map(Option_t *opt); // *MENU*
301350
virtual void Map() { Map(""); }; // *MENU*

‎io/io/src/TFile.cxx

Lines changed: 93 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1574,8 +1574,10 @@ void TFile::Map(Option_t *opt)
15741574

15751575
const unsigned char nDigits = std::log10(fEND) + 1;
15761576

1577+
std::optional<ROOT::Detail::TKeyMapNode> lastNode;
15771578
const auto tkeyInfos = WalkTKeys();
15781579
for (const auto &key : tkeyInfos) {
1580+
lastNode = key;
15791581
switch (key.fType) {
15801582
case ROOT::Detail::TKeyMapNode::kError:
15811583
Printf("Address = %" PRIu64 "\tNbytes = %u\t=====E R R O R=======", key.fAddr, key.fLen);
@@ -1616,7 +1618,7 @@ void TFile::Map(Option_t *opt)
16161618
}
16171619

16181620
if (!forComp) {
1619-
Int_t datime = tkeyInfos.empty() ? 0 : tkeyInfos.back().fDatime;
1621+
Int_t datime = lastNode ? lastNode->fDatime : 0;
16201622
Int_t date, time;
16211623
TDatime::GetDateTime(datime, date, time);
16221624
Printf("%d/%06d At:%-*lld N=%-8d %-14s", date, time, nDigits + 1, fEND, 1, "END");
@@ -1625,94 +1627,108 @@ void TFile::Map(Option_t *opt)
16251627
}
16261628
}
16271629

1628-
std::vector<ROOT::Detail::TKeyMapNode> TFile::WalkTKeys()
1630+
ROOT::Detail::TKeyMapIterable TFile::WalkTKeys()
16291631
{
1630-
std::vector<ROOT::Detail::TKeyMapNode> res;
1631-
1632-
constexpr int headerSize = 512;
1633-
1634-
std::uint64_t idcur = fBEGIN;
1635-
const std::uint64_t end = fEND;
1636-
while (idcur < end) {
1637-
Seek(idcur);
1638-
auto nread = headerSize;
1639-
if (idcur + nread >= end)
1640-
nread = fEND - idcur - 1;
1641-
1642-
char header[headerSize];
1643-
if (ReadBuffer(header, nread)) {
1644-
// ReadBuffer returns kTRUE in case of failure.
1645-
res.push_back({idcur, ROOT::Detail::TKeyMapNode::kError});
1646-
break;
1647-
}
1632+
return ROOT::Detail::TKeyMapIterable(this);
1633+
}
16481634

1649-
char *buffer = header;
1650-
Int_t nbytes;
1651-
frombuf(buffer, &nbytes);
1652-
if (!nbytes) {
1653-
res.push_back({idcur, ROOT::Detail::TKeyMapNode::kError});
1654-
break;
1655-
}
1656-
if (nbytes < 0) {
1657-
// free slot
1658-
res.push_back({idcur, ROOT::Detail::TKeyMapNode::kGap, static_cast<std::uint32_t>(-nbytes)});
1659-
idcur -= nbytes;
1660-
continue;
1661-
}
1635+
ROOT::Detail::TKeyMapIterable::TIterator::TIterator(TFile *file, std::uint64_t addr) : fFile(file), fCurAddr(addr)
1636+
{
1637+
if (addr == 0)
1638+
fCurAddr = fFile->fBEGIN;
1639+
Advance();
1640+
}
16621641

1663-
auto &node = res.emplace_back(idcur, ROOT::Detail::TKeyMapNode::kKey, nbytes);
1642+
std::optional<ROOT::Detail::TKeyMapNode> ROOT::Detail::TKeyMapIterable::TIterator::Next()
1643+
{
1644+
constexpr int headerSize = 512;
16641645

1665-
frombuf(buffer, &node.fKeyVersion);
1666-
frombuf(buffer, &node.fObjLen);
1667-
frombuf(buffer, &node.fDatime);
1668-
frombuf(buffer, &node.fKeyLen);
1669-
frombuf(buffer, &node.fCycle);
1670-
if (node.fKeyVersion > 1000) {
1671-
frombuf(buffer, &node.fSeekKey);
1672-
frombuf(buffer, &node.fSeekPdir);
1646+
const std::uint64_t idcur = fCurAddr;
1647+
const std::uint64_t end = fFile->fEND;
1648+
if (idcur >= end)
1649+
return std::nullopt;
1650+
1651+
fFile->Seek(idcur);
1652+
auto nread = headerSize;
1653+
if (idcur + nread >= end)
1654+
nread = end - idcur - 1;
1655+
1656+
char header[headerSize];
1657+
if (fFile->ReadBuffer(header, nread)) {
1658+
// ReadBuffer returns kTRUE in case of failure.
1659+
auto node = ROOT::Detail::TKeyMapNode{idcur, ROOT::Detail::TKeyMapNode::kError};
1660+
fCurAddr = end;
1661+
return node;
1662+
}
1663+
1664+
char *buffer = header;
1665+
Int_t nbytes;
1666+
frombuf(buffer, &nbytes);
1667+
if (!nbytes) {
1668+
auto node = ROOT::Detail::TKeyMapNode{idcur, ROOT::Detail::TKeyMapNode::kError};
1669+
fCurAddr = end;
1670+
return node;
1671+
}
1672+
1673+
if (nbytes < 0) {
1674+
// free slot
1675+
auto node =
1676+
ROOT::Detail::TKeyMapNode{idcur, ROOT::Detail::TKeyMapNode::kGap, static_cast<std::uint32_t>(-nbytes)};
1677+
fCurAddr -= nbytes;
1678+
return node;
1679+
}
1680+
1681+
auto node = ROOT::Detail::TKeyMapNode{idcur, ROOT::Detail::TKeyMapNode::kKey, static_cast<std::uint32_t>(nbytes)};
1682+
frombuf(buffer, &node.fKeyVersion);
1683+
frombuf(buffer, &node.fObjLen);
1684+
frombuf(buffer, &node.fDatime);
1685+
frombuf(buffer, &node.fKeyLen);
1686+
frombuf(buffer, &node.fCycle);
1687+
if (node.fKeyVersion > 1000) {
1688+
frombuf(buffer, &node.fSeekKey);
1689+
frombuf(buffer, &node.fSeekPdir);
1690+
} else {
1691+
Int_t skey, sdir;
1692+
frombuf(buffer, &skey);
1693+
frombuf(buffer, &sdir);
1694+
node.fSeekKey = static_cast<Long64_t>(skey);
1695+
node.fSeekPdir = static_cast<Long64_t>(sdir);
1696+
}
1697+
1698+
const auto readString = [&buffer, &header](bool skipCheck = false) {
1699+
char stringLen;
1700+
char str[256];
1701+
if (!skipCheck && ((buffer - header) >= headerSize)) {
1702+
stringLen = 0;
16731703
} else {
1674-
Int_t skey, sdir;
1675-
frombuf(buffer, &skey);
1676-
frombuf(buffer, &sdir);
1677-
node.fSeekKey = static_cast<Long64_t>(skey);
1678-
node.fSeekPdir = static_cast<Long64_t>(sdir);
1679-
}
1680-
1681-
const auto readString = [&buffer, &header] (bool skipCheck = false) {
1682-
char stringLen;
1683-
char str[256];
1684-
if (!skipCheck && ((buffer - header) >= headerSize)) {
1704+
frombuf(buffer, &stringLen);
1705+
if (stringLen < 0)
16851706
stringLen = 0;
1686-
} else {
1687-
frombuf(buffer, &stringLen);
1688-
if (stringLen < 0)
1689-
stringLen = 0;
1690-
else if ((buffer - header) + stringLen > headerSize)
1691-
stringLen = headerSize - (buffer - header);
1692-
}
1693-
for (int i = 0; i < stringLen; ++i)
1694-
frombuf(buffer, &str[i]);
1695-
str[static_cast<int>(stringLen)] = 0;
1707+
else if ((buffer - header) + stringLen > headerSize)
1708+
stringLen = headerSize - (buffer - header);
1709+
}
1710+
for (int i = 0; i < stringLen; ++i)
1711+
frombuf(buffer, &str[i]);
1712+
str[static_cast<int>(stringLen)] = 0;
16961713

1697-
return std::string(str, stringLen);
1698-
};
1714+
return std::string(str, stringLen);
1715+
};
16991716

1700-
node.fClassName = readString(true);
1717+
node.fClassName = readString(true);
17011718

1702-
if (idcur == static_cast<std::uint64_t>(fSeekFree))
1703-
node.fClassName = "FreeSegments";
1704-
else if (idcur == static_cast<std::uint64_t>(fSeekInfo))
1705-
node.fClassName = "StreamerInfo";
1706-
else if (idcur == static_cast<std::uint64_t>(fSeekKeys))
1707-
node.fClassName = "KeysList";
1719+
if (idcur == static_cast<std::uint64_t>(fFile->fSeekFree))
1720+
node.fClassName = "FreeSegments";
1721+
else if (idcur == static_cast<std::uint64_t>(fFile->fSeekInfo))
1722+
node.fClassName = "StreamerInfo";
1723+
else if (idcur == static_cast<std::uint64_t>(fFile->fSeekKeys))
1724+
node.fClassName = "KeysList";
17081725

1709-
node.fKeyName = readString();
1710-
node.fKeyTitle = readString();
1726+
node.fKeyName = readString();
1727+
node.fKeyTitle = readString();
17111728

1712-
idcur += nbytes;
1713-
}
1729+
fCurAddr += nbytes;
17141730

1715-
return res;
1731+
return node;
17161732
}
17171733

17181734
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)