Skip to content

Commit 2bcccda

Browse files
silverweedsilverweed
committed
[io] Change TFile::WalkTKeys() to return an iterable
1 parent d516f4b commit 2bcccda

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
@@ -1572,8 +1572,10 @@ void TFile::Map(Option_t *opt)
15721572

15731573
const unsigned char nDigits = std::log10(fEND) + 1;
15741574

1575+
std::optional<ROOT::Detail::TKeyMapNode> lastNode;
15751576
const auto tkeyInfos = WalkTKeys();
15761577
for (const auto &key : tkeyInfos) {
1578+
lastNode = key;
15771579
switch (key.fType) {
15781580
case ROOT::Detail::TKeyMapNode::kError:
15791581
Printf("Address = %" PRIu64 "\tNbytes = %u\t=====E R R O R=======", key.fAddr, key.fLen);
@@ -1614,7 +1616,7 @@ void TFile::Map(Option_t *opt)
16141616
}
16151617

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

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

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

1661-
auto &node = res.emplace_back(idcur, ROOT::Detail::TKeyMapNode::kKey, nbytes);
1640+
std::optional<ROOT::Detail::TKeyMapNode> ROOT::Detail::TKeyMapIterable::TIterator::Next()
1641+
{
1642+
static constexpr int headerSize = 512;
16621643

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

1695-
return std::string(str, stringLen);
1696-
};
1712+
return std::string(str, stringLen);
1713+
};
16971714

1698-
node.fClassName = readString(true);
1715+
node.fClassName = readString(true);
16991716

1700-
if (idcur == static_cast<std::uint64_t>(fSeekFree))
1701-
node.fClassName = "FreeSegments";
1702-
else if (idcur == static_cast<std::uint64_t>(fSeekInfo))
1703-
node.fClassName = "StreamerInfo";
1704-
else if (idcur == static_cast<std::uint64_t>(fSeekKeys))
1705-
node.fClassName = "KeysList";
1717+
if (idcur == static_cast<std::uint64_t>(fFile->fSeekFree))
1718+
node.fClassName = "FreeSegments";
1719+
else if (idcur == static_cast<std::uint64_t>(fFile->fSeekInfo))
1720+
node.fClassName = "StreamerInfo";
1721+
else if (idcur == static_cast<std::uint64_t>(fFile->fSeekKeys))
1722+
node.fClassName = "KeysList";
17061723

1707-
node.fKeyName = readString();
1708-
node.fKeyTitle = readString();
1724+
node.fKeyName = readString();
1725+
node.fKeyTitle = readString();
17091726

1710-
idcur += nbytes;
1711-
}
1727+
fCurAddr += nbytes;
17121728

1713-
return res;
1729+
return node;
17141730
}
17151731

17161732
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)