erigon_seg/writer/
bt_writer.rs1use std::fs::File;
14use std::io::Write;
15use std::path::Path;
16
17use super::ef_builder::EfBuilder;
18use crate::error::{Error, Result};
19use crate::seg::Seg;
20
21pub const DEFAULT_BTREE_M: u64 = 256;
23
24const BT_EF_ALIGN: usize = 4096;
25const BT_FOOTER_ALIGN: usize = 8;
26const BT_VERSION: u16 = 1;
27const BT_METADATA_LEN: u32 = 24;
28const FOOTER_MAGIC: [u8; 8] = *b"erigon\x00\x00";
29const FIRST_BYTE_FOOTER: u8 = 0x01;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum BtLayout {
34 Legacy,
36 Footer,
38}
39
40#[derive(Debug, Clone, Copy)]
42pub struct BtOptions {
43 pub layout: BtLayout,
45 pub m: u64,
47}
48
49impl Default for BtOptions {
50 fn default() -> BtOptions {
51 BtOptions {
52 layout: BtLayout::Footer,
53 m: DEFAULT_BTREE_M,
54 }
55 }
56}
57
58pub fn build_bt(
60 kv_path: impl AsRef<Path>,
61 bt_path: impl AsRef<Path>,
62 opts: BtOptions,
63) -> Result<()> {
64 let seg = Seg::open(kv_path)?;
65 build_bt_from_seg(&seg, bt_path, opts)
66}
67
68pub fn build_bt_from_seg(seg: &Seg, bt_path: impl AsRef<Path>, opts: BtOptions) -> Result<()> {
70 let bt_path = bt_path.as_ref();
71 let key_count = seg.words_count() / 2;
72
73 if key_count == 0 {
75 File::create(bt_path).map_err(|e| Error::io(bt_path, e))?;
76 return Ok(());
77 }
78
79 let max_offset = seg.len() as u64;
80 let m = opts.m.max(1);
81 let mut ef = EfBuilder::new(key_count, max_offset);
82
83 let mut nodes: Vec<u8> = Vec::new();
85 let footer = opts.layout == BtLayout::Footer;
86 if footer {
87 nodes.push(FIRST_BYTE_FOOTER);
88 }
89
90 let mut g = seg.getter();
91 for di in 0..key_count {
92 let off = g.offset();
93 if footer && di % m == 0 {
94 let key = g.next(); let klen = u16::try_from(key.len()).map_err(|_| {
96 Error::format("key longer than 65535 bytes (unsupported in .bt node)")
97 })?;
98 nodes.extend_from_slice(&klen.to_be_bytes());
99 nodes.extend_from_slice(&key);
100 } else {
101 g.skip(); }
103 g.skip(); ef.add_offset(off);
105 }
106 ef.build();
107
108 match opts.layout {
109 BtLayout::Legacy => {
110 let mut out = Vec::with_capacity(ef.serialized_len());
111 ef.write_to(&mut out);
112 write_all(bt_path, &out)
113 }
114 BtLayout::Footer => {
115 let mut out = nodes;
116 pad_to(&mut out, BT_EF_ALIGN);
117 let ef_offset = out.len() as u64;
118 ef.write_to(&mut out);
119 pad_to(&mut out, BT_FOOTER_ALIGN);
120 out.extend_from_slice(&key_count.to_be_bytes());
122 out.extend_from_slice(&m.to_be_bytes());
123 out.extend_from_slice(&ef_offset.to_be_bytes());
124 out.extend_from_slice(&BT_METADATA_LEN.to_be_bytes());
126 out.extend_from_slice(&0u16.to_be_bytes()); out.extend_from_slice(&BT_VERSION.to_be_bytes());
128 out.extend_from_slice(&FOOTER_MAGIC);
129 write_all(bt_path, &out)
130 }
131 }
132}
133
134fn pad_to(out: &mut Vec<u8>, align: usize) {
135 let rem = out.len() % align;
136 if rem != 0 {
137 out.resize(out.len() + (align - rem), 0);
138 }
139}
140
141fn write_all(path: &Path, bytes: &[u8]) -> Result<()> {
142 let mut f = File::create(path).map_err(|e| Error::io(path, e))?;
143 f.write_all(bytes).map_err(|e| Error::io(path, e))?;
144 f.flush().map_err(|e| Error::io(path, e))
145}