erigon_seg/btree.rs
1//! Reader for a `.bt` B-tree index file.
2//!
3//! The `.bt` carries two things a point lookup can use:
4//!
5//! * an Elias-Fano array of the `.kv` byte offset of every key, in key order; and
6//! * (footer layout only) the *di-nodes*: the key at every `M`-th position, stored
7//! uncompressed.
8//!
9//! Without the di-nodes a lookup is a binary search over all `n` keys, where each probe
10//! seeks the `.kv` getter to `offset(i)` and decompresses the key to compare — roughly
11//! `log2(n)` decompressions, each landing on a different part of the file. The di-nodes
12//! narrow that to the one `M`-key block that can contain the key, using only `memcmp`
13//! against keys already in memory: `log2(n/M)` comparisons plus `log2(M)` decompressions,
14//! and those decompressions all land inside a single contiguous block. See
15//! [`Nodes::narrow`].
16//!
17//! Two on-disk layouts are supported:
18//!
19//! * **legacy** — `[EliasFano][nodes…]`; the first byte is `0x00` (the high byte of the
20//! EF `count`), and the EF starts at offset 0. The trailing nodes are not located by
21//! any header, so narrowing is unavailable and lookups use the full binary search.
22//! * **footer** — `[0x01][nodes…][EliasFano][footer][anchor]`; the fixed 16-byte anchor
23//! ends with the magic `erigon\0\0` and carries `footer_len`; the variable footer
24//! holds `keys_count`, `M`, and `ef_offset` locating the EF section.
25
26use std::path::Path;
27use std::sync::{Arc, OnceLock};
28
29use memmap2::Mmap;
30
31use crate::eliasfano::EliasFano;
32use crate::error::{Error, Result};
33use crate::util::{Advice, advise_mmap, lock_mmap, mmap_file, preload_mmap, unlock_mmap};
34
35/// The fixed footer anchor is 16 bytes: `footer_len:u32 | flags:u16 | version:u16 | magic:u64`.
36const ANCHOR_LEN: usize = 16;
37/// The variable footer payload is at least `keys_count(8) | M(8) | ef_offset(8)`.
38const META_LEN: usize = 24;
39/// Trailing magic identifying the footer layout (and proving the file isn't truncated).
40const FOOTER_MAGIC: [u8; 8] = *b"erigon\x00\x00";
41/// First byte of a footer-layout file (a legacy file has `0x00` here).
42const FIRST_BYTE_FOOTER: u8 = 0x01;
43
44/// A `.bt` index: the Elias-Fano offset array plus, when known, the B-tree fanout `M`
45/// and the di-node array used to narrow lookups.
46pub struct BtreeIndex {
47 ef: Option<EliasFano>,
48 m: Option<u64>,
49 /// Kept so the di-nodes can be parsed on first use rather than at open time.
50 mmap: Arc<Mmap>,
51 /// Where the di-nodes live, for the footer layout: `(keys_count, m, ef_offset)`.
52 node_src: Option<(u64, u64, usize)>,
53 /// Parsed lazily by [`BtreeIndex::nodes`]; `None` once parsing has been attempted
54 /// and found the section unusable.
55 nodes: OnceLock<Option<Nodes>>,
56}
57
58/// The `.bt` di-node array: the key at every `M`-th position.
59///
60/// Copied out of the mapping into one compact arena at first use, so the narrowing
61/// search touches only hot, contiguous heap instead of faulting `.bt` pages scattered
62/// across a multi-gigabyte file.
63pub struct Nodes {
64 /// All node keys back to back.
65 arena: Vec<u8>,
66 /// End offset of each node key within `arena`, or empty when every key has the same
67 /// length (`fixed_len`), in which case offsets are computed arithmetically.
68 ends: Vec<u32>,
69 /// `Some(len)` when every node key is `len` bytes — the common case, which lets us
70 /// drop `ends` entirely (4 bytes/node saved).
71 fixed_len: Option<u32>,
72 count: usize,
73 m: u64,
74 key_count: u64,
75}
76
77impl Nodes {
78 /// Number of nodes (`ceil(key_count / M)`).
79 pub fn len(&self) -> usize {
80 self.count
81 }
82
83 /// Whether there are no nodes.
84 pub fn is_empty(&self) -> bool {
85 self.count == 0
86 }
87
88 /// The B-tree fanout `M` these nodes were sampled at.
89 pub fn m(&self) -> u64 {
90 self.m
91 }
92
93 /// Heap held by the arena and (when present) the offset table.
94 pub fn heap_bytes(&self) -> usize {
95 self.arena.len() + self.ends.len() * 4
96 }
97
98 /// The `j`-th node key, which is key number `j * M` in the file.
99 ///
100 /// # Panics
101 /// Panics if `j >= self.len()`.
102 #[inline]
103 pub fn key(&self, j: usize) -> &[u8] {
104 match self.fixed_len {
105 Some(l) => {
106 let l = l as usize;
107 &self.arena[j * l..(j + 1) * l]
108 }
109 None => {
110 let start = if j == 0 { 0 } else { self.ends[j - 1] as usize };
111 &self.arena[start..self.ends[j] as usize]
112 }
113 }
114 }
115
116 /// Narrow a search for `key` to the half-open key-index range that can contain it.
117 ///
118 /// Node `j` is key `j * M`, and keys are sorted, so if `node[j] <= key < node[j+1]`
119 /// then `key`, if present, lies in `[j*M, (j+1)*M)`. Returns an empty range when
120 /// `key` sorts before the very first key, which cannot be in the file at all.
121 #[inline]
122 pub fn narrow(&self, key: &[u8]) -> (u64, u64) {
123 // Index of the first node strictly greater than `key`.
124 let mut lo = 0usize;
125 let mut hi = self.count;
126 while lo < hi {
127 let mid = lo + (hi - lo) / 2;
128 if self.key(mid) <= key {
129 lo = mid + 1;
130 } else {
131 hi = mid;
132 }
133 }
134 if lo == 0 {
135 return (0, 0); // key < node[0] == first key in the file
136 }
137 let start = (lo as u64 - 1) * self.m;
138 (start, (start + self.m).min(self.key_count))
139 }
140}
141
142impl BtreeIndex {
143 /// Open and parse a `.bt` file, auto-detecting the legacy vs footer layout.
144 pub fn open(path: impl AsRef<Path>) -> Result<BtreeIndex> {
145 let path = path.as_ref();
146 let mmap = mmap_file(path)?;
147 let len = mmap.len();
148
149 let mmap = Arc::new(mmap);
150
151 // A zero-length .bt is a valid empty index (0 keys).
152 if len == 0 {
153 return Ok(BtreeIndex {
154 ef: None,
155 m: None,
156 mmap,
157 node_src: None,
158 nodes: OnceLock::from(None),
159 });
160 }
161
162 // Footer layout iff the trailing anchor carries the magic.
163 if len >= ANCHOR_LEN && mmap[len - 8..len] == FOOTER_MAGIC {
164 let anchor = &mmap[len - ANCHOR_LEN..];
165 let footer_len = u32::from_be_bytes(anchor[0..4].try_into().unwrap()) as usize;
166 if footer_len < META_LEN || ANCHOR_LEN + footer_len > len {
167 return Err(Error::format(format!(
168 "{}: corrupt .bt footer (footer_len={footer_len}, file={len})",
169 path.display()
170 )));
171 }
172 let footer_start = len - ANCHOR_LEN - footer_len;
173 let payload = &mmap[footer_start..len - ANCHOR_LEN];
174 let keys_count = u64::from_be_bytes(payload[0..8].try_into().unwrap());
175 let m = u64::from_be_bytes(payload[8..16].try_into().unwrap());
176 let ef_offset = u64::from_be_bytes(payload[16..24].try_into().unwrap()) as usize;
177 if ef_offset >= footer_start {
178 return Err(Error::format(format!(
179 "{}: corrupt .bt footer (ef_offset={ef_offset} >= body={footer_start})",
180 path.display()
181 )));
182 }
183 let ef = EliasFano::open(Arc::clone(&mmap), ef_offset)?;
184 if ef.len() != keys_count {
185 return Err(Error::format(format!(
186 "{}: .bt EF has {} keys, footer says {keys_count}",
187 path.display(),
188 ef.len()
189 )));
190 }
191 return Ok(BtreeIndex {
192 ef: Some(ef),
193 m: Some(m),
194 mmap,
195 node_src: Some((keys_count, m, ef_offset)),
196 nodes: OnceLock::new(),
197 });
198 }
199
200 // No magic: must be the legacy layout, whose first byte is 0x00.
201 if mmap[0] == FIRST_BYTE_FOOTER {
202 return Err(Error::format(format!(
203 "{}: .bt looks like footer layout but the trailing magic is missing (truncated?)",
204 path.display()
205 )));
206 }
207 let ef = EliasFano::open(Arc::clone(&mmap), 0)?;
208 Ok(BtreeIndex {
209 ef: Some(ef),
210 m: None,
211 mmap,
212 // Legacy layout: nothing locates the trailing nodes, so there is no narrowing.
213 node_src: None,
214 nodes: OnceLock::from(None),
215 })
216 }
217
218 /// Number of indexed keys.
219 pub fn key_count(&self) -> u64 {
220 self.ef.as_ref().map_or(0, EliasFano::len)
221 }
222
223 /// The `.kv` byte offset of the `i`-th key (0-based). Returns `None` if out of range.
224 pub fn key_offset(&self, i: u64) -> Option<u64> {
225 let ef = self.ef.as_ref()?;
226 (i < ef.len()).then(|| ef.get(i))
227 }
228
229 /// The B-tree fanout `M`, if the layout records it (footer layout only).
230 pub fn m(&self) -> Option<u64> {
231 self.m
232 }
233
234 /// The di-node array, parsed on first call and cached thereafter.
235 ///
236 /// Returns `None` for the legacy layout, for an empty index, or if the section does
237 /// not parse — in each case lookups simply fall back to the full binary search.
238 ///
239 /// Parsing walks the whole node section once (`ceil(key_count / M)` entries) and
240 /// copies the keys into an arena, so the first call costs one pass over that
241 /// section and holds it in memory. It is deliberately *not* done at open time, so
242 /// opening a file only to scan it — merging, re-encoding — pays nothing.
243 pub fn nodes(&self) -> Option<&Nodes> {
244 self.nodes
245 .get_or_init(|| {
246 let (key_count, m, ef_offset) = self.node_src?;
247 parse_nodes(&self.mmap, key_count, m, ef_offset)
248 })
249 .as_ref()
250 }
251
252 /// Narrow a lookup for `key` to the half-open key-index range that can contain it,
253 /// using the di-nodes. Falls back to the full range when narrowing is unavailable.
254 #[inline]
255 pub fn narrow(&self, key: &[u8]) -> (u64, u64) {
256 match self.nodes() {
257 Some(n) => n.narrow(key),
258 None => (0, self.key_count()),
259 }
260 }
261
262 /// Advise the kernel that this `.bt` is read in random order (point lookups). See
263 /// [`KvReader::advise_random`](crate::KvReader::advise_random).
264 pub fn advise_random(&self) -> std::io::Result<()> {
265 advise_mmap(&self.mmap, Advice::Random)
266 }
267
268 /// Bytes this `.bt` occupies when fully resident — what
269 /// [`preload`](BtreeIndex::preload) or [`lock`](BtreeIndex::lock) would cost.
270 pub fn mapped_bytes(&self) -> u64 {
271 self.mmap.len() as u64
272 }
273
274 /// Read the whole `.bt` into the page cache, returning once it is resident. See
275 /// [`KvReader::preload_index`](crate::KvReader::preload_index).
276 pub fn preload(&self) -> u64 {
277 preload_mmap(&self.mmap) as u64
278 }
279
280 /// Pin the whole `.bt` in RAM with `mlock`. See
281 /// [`KvReader::lock_index`](crate::KvReader::lock_index) for the caveats.
282 pub fn lock(&self) -> std::io::Result<()> {
283 lock_mmap(&self.mmap)
284 }
285
286 /// Release an [`mlock`](BtreeIndex::lock).
287 pub fn unlock(&self) -> std::io::Result<()> {
288 unlock_mmap(&self.mmap)
289 }
290
291 /// Borrow the underlying Elias-Fano offset array, if the index is non-empty.
292 pub fn elias_fano(&self) -> Option<&EliasFano> {
293 self.ef.as_ref()
294 }
295}
296
297/// Parse the di-node array out of a footer-layout `.bt`.
298///
299/// Layout from byte 1 (after the `0x01` marker): `keys_count / M` entries of
300/// `klen:u16-BE | key`, then zero padding up to `ef_offset`. Returns `None` if anything
301/// fails to line up, which only costs the narrowing optimization.
302fn parse_nodes(data: &[u8], key_count: u64, m: u64, ef_offset: usize) -> Option<Nodes> {
303 if key_count == 0 || m == 0 || ef_offset <= 1 || ef_offset > data.len() {
304 return None;
305 }
306 let count = usize::try_from(key_count.div_ceil(m)).ok()?;
307 let mut arena: Vec<u8> = Vec::new();
308 let mut ends: Vec<u32> = Vec::with_capacity(count);
309 let mut fixed_len: Option<u32> = None;
310 let mut uniform = true;
311 let mut p = 1usize;
312 for j in 0..count {
313 let lb = data.get(p..p + 2)?;
314 let klen = u16::from_be_bytes(lb.try_into().ok()?) as usize;
315 p += 2;
316 if p + klen > ef_offset {
317 return None;
318 }
319 match fixed_len {
320 None if j == 0 => fixed_len = Some(klen as u32),
321 Some(l) if l as usize != klen => uniform = false,
322 _ => {}
323 }
324 arena.extend_from_slice(&data[p..p + klen]);
325 p += klen;
326 ends.push(u32::try_from(arena.len()).ok()?);
327 }
328 if uniform {
329 // Every key is the same length, so offsets are `j * len` — drop the table.
330 ends = Vec::new();
331 } else {
332 fixed_len = None;
333 }
334 Some(Nodes {
335 arena,
336 ends,
337 fixed_len,
338 count,
339 m,
340 key_count,
341 })
342}