Skip to main content

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}