use crate::{enc_varint_into, local_payload_len, Value};
#[derive(Debug, Clone, PartialEq)]
pub struct RebuildRow {
pub page: u32,
pub offset: usize,
pub rowid: Option<i64>,
pub source: String,
pub confidence: f32,
pub cells: Vec<Value>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct FragmentRow {
pub page: u32,
pub offset: usize,
pub missing: usize,
pub confidence: f32,
pub surviving: Vec<(usize, Value)>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct RecoveredTable {
pub name: String,
pub columns: Vec<String>,
pub rows: Vec<Vec<Value>>,
}
#[must_use]
pub fn build_recovered_db_tables(tables: &[RecoveredTable]) -> Vec<u8> {
let specs: Vec<TableSpec> = tables
.iter()
.map(|t| TableSpec {
name: t.name.clone(),
create_sql: create_table_sql(&t.name, &t.columns),
rows: t.rows.clone(),
})
.collect();
build_from_specs(&specs)
}
const PAGE_SIZE: usize = 4096;
const USABLE: usize = PAGE_SIZE;
const OVERFLOW_PAYLOAD: usize = USABLE - 4;
const LEAD_COLS: usize = 5;
struct TableSpec {
name: String,
create_sql: String,
rows: Vec<Vec<Value>>,
}
#[must_use]
pub fn build_recovered_db(rows: &[RebuildRow]) -> Vec<u8> {
build_recovered_db_with_fragments(rows, None)
}
#[must_use]
pub fn build_recovered_db_with_fragments(
rows: &[RebuildRow],
fragments: Option<&[FragmentRow]>,
) -> Vec<u8> {
let mut specs = vec![records_spec(rows)];
if let Some(frags) = fragments {
specs.push(fragments_spec(frags));
}
build_from_specs(&specs)
}
fn build_from_specs(specs: &[TableSpec]) -> Vec<u8> {
let mut builder = Builder::new();
let _schema_page = builder.alloc_page(); let placed: Vec<(&str, u32, &str)> = specs
.iter()
.map(|spec| {
let root = builder.build_table(&spec.rows);
(spec.name.as_str(), root, spec.create_sql.as_str())
})
.collect();
let schema_leaf = build_schema_leaf(&placed, builder.page_count());
builder.set_page(1, schema_leaf);
builder.finish()
}
fn quote_ident(ident: &str) -> String {
let mut out = String::with_capacity(ident.len() + 2);
out.push('"');
for ch in ident.chars() {
if ch == '"' {
out.push('"');
}
out.push(ch);
}
out.push('"');
out
}
fn create_table_sql(name: &str, columns: &[String]) -> String {
let mut sql = String::from("CREATE TABLE ");
sql.push_str("e_ident(name));
sql.push_str(" (");
for (i, col) in columns.iter().enumerate() {
if i > 0 {
sql.push_str(", ");
}
sql.push_str("e_ident(col));
}
sql.push(')');
sql
}
fn records_spec(rows: &[RebuildRow]) -> TableSpec {
let max_cells = rows.iter().map(|r| r.cells.len()).max().unwrap_or(0);
let values = rows.iter().map(|r| record_values(r, max_cells)).collect();
TableSpec {
name: "recovered_records".to_string(),
create_sql: create_records_sql(max_cells),
rows: values,
}
}
fn fragments_spec(frags: &[FragmentRow]) -> TableSpec {
let max_col = frags
.iter()
.flat_map(|f| f.surviving.iter().map(|(idx, _)| *idx))
.max();
let cell_cols = max_col.map_or(0, |m| m + 1);
let values = frags
.iter()
.map(|f| fragment_values(f, cell_cols))
.collect();
TableSpec {
name: "recovered_fragments".to_string(),
create_sql: create_fragments_sql(cell_cols),
rows: values,
}
}
fn create_records_sql(max_cells: usize) -> String {
let mut sql = String::from(
"CREATE TABLE recovered_records (\n _page INTEGER, _offset INTEGER, \
_rowid INTEGER, _source TEXT, _confidence REAL",
);
append_cell_columns(&mut sql, max_cells);
sql.push(')');
sql
}
fn create_fragments_sql(cell_cols: usize) -> String {
let mut sql = String::from(
"CREATE TABLE recovered_fragments (\n _page INTEGER, _offset INTEGER, \
_missing INTEGER, _confidence REAL",
);
append_cell_columns(&mut sql, cell_cols);
sql.push(')');
sql
}
fn append_cell_columns(sql: &mut String, n: usize) {
for i in 0..n {
use std::fmt::Write as _;
let _ = write!(sql, ", c{i}");
}
}
fn record_values(row: &RebuildRow, max_cells: usize) -> Vec<Value> {
let mut values: Vec<Value> = Vec::with_capacity(LEAD_COLS + max_cells);
values.push(Value::Integer(i64::from(row.page)));
values.push(Value::Integer(row.offset as i64));
values.push(match row.rowid {
Some(id) => Value::Integer(id),
None => Value::Null,
});
values.push(Value::Text(row.source.clone()));
values.push(Value::Real(f64::from(row.confidence)));
for i in 0..max_cells {
values.push(row.cells.get(i).cloned().unwrap_or(Value::Null));
}
values
}
fn fragment_values(frag: &FragmentRow, cell_cols: usize) -> Vec<Value> {
let mut cells = vec![Value::Null; cell_cols];
for (idx, v) in &frag.surviving {
if let Some(slot) = cells.get_mut(*idx) {
*slot = v.clone();
}
}
let mut values: Vec<Value> = Vec::with_capacity(FRAG_LEAD_COLS + cell_cols);
values.push(Value::Integer(i64::from(frag.page)));
values.push(Value::Integer(frag.offset as i64));
values.push(Value::Integer(frag.missing as i64));
values.push(Value::Real(f64::from(frag.confidence)));
values.extend(cells);
values
}
const FRAG_LEAD_COLS: usize = 4;
struct Builder {
pages: Vec<Vec<u8>>, }
impl Builder {
fn new() -> Self {
Self { pages: Vec::new() }
}
fn alloc_page(&mut self) -> u32 {
self.pages.push(vec![0u8; PAGE_SIZE]);
self.pages.len() as u32
}
fn set_page(&mut self, page: u32, mut content: Vec<u8>) {
content.resize(PAGE_SIZE, 0);
let idx = page as usize - 1;
self.pages[idx] = content;
}
fn push_page(&mut self, mut content: Vec<u8>) -> u32 {
content.resize(PAGE_SIZE, 0);
self.pages.push(content);
self.pages.len() as u32
}
fn page_count(&self) -> u32 {
self.pages.len() as u32
}
fn build_table(&mut self, rows: &[Vec<Value>]) -> u32 {
let payloads: Vec<(i64, Vec<u8>)> = rows
.iter()
.enumerate()
.map(|(i, values)| (i as i64 + 1, encode_record(values)))
.collect();
let leaves = self.pack_leaves(&payloads);
debug_assert!(!leaves.is_empty());
if leaves.len() == 1 {
leaves[0].0
} else {
self.build_interiors(&leaves)
}
}
fn pack_leaves(&mut self, payloads: &[(i64, Vec<u8>)]) -> Vec<(u32, i64)> {
let mut leaves: Vec<(u32, i64)> = Vec::new();
let mut cur: Vec<LeafCell> = Vec::new();
let mut cur_max_rowid = 0i64;
for (rowid, payload) in payloads {
let cell = self.make_leaf_cell(*rowid, payload);
let used: usize = 8 + cur.iter().map(|c| 2 + c.on_page.len()).sum::<usize>();
let need = 2 + cell.on_page.len();
if !cur.is_empty() && used + need > PAGE_SIZE {
let page = self.flush_leaf(&cur);
leaves.push((page, cur_max_rowid));
cur.clear();
}
cur_max_rowid = *rowid;
cur.push(cell);
}
let page = self.flush_leaf(&cur);
leaves.push((page, cur_max_rowid));
leaves
}
fn make_leaf_cell(&mut self, rowid: i64, payload: &[u8]) -> LeafCell {
let total = payload.len();
let local = local_payload_len(total, USABLE);
let mut on_page = Vec::new();
on_page.extend_from_slice(&enc_varint_into(total));
on_page.extend_from_slice(&enc_varint_into(rowid_as_usize(rowid)));
if local >= total {
on_page.extend_from_slice(payload);
} else {
on_page.extend_from_slice(&payload[..local]);
let first_overflow = self.write_overflow_chain(&payload[local..]);
on_page.extend_from_slice(&first_overflow.to_be_bytes());
}
LeafCell { on_page }
}
fn write_overflow_chain(&mut self, rest: &[u8]) -> u32 {
let n_pages = rest.len().div_ceil(OVERFLOW_PAYLOAD);
let first = self.page_count() + 1;
let mut chunks = rest.chunks(OVERFLOW_PAYLOAD);
for i in 0..n_pages {
let chunk = chunks.next().unwrap_or(&[]);
let next = if i + 1 < n_pages {
first + i as u32 + 1
} else {
0
};
let mut page = Vec::with_capacity(PAGE_SIZE);
page.extend_from_slice(&next.to_be_bytes());
page.extend_from_slice(chunk);
self.push_page(page);
}
first
}
fn flush_leaf(&mut self, cells: &[LeafCell]) -> u32 {
let mut page = vec![0u8; PAGE_SIZE];
page[0] = 0x0d; write_u16(&mut page, 1, 0); write_u16(&mut page, 3, cells.len() as u16); let mut content_start = PAGE_SIZE;
let header_end = 8 + cells.len() * 2;
for (i, cell) in cells.iter().enumerate() {
content_start -= cell.on_page.len();
page[content_start..content_start + cell.on_page.len()].copy_from_slice(&cell.on_page);
write_u16(&mut page, 8 + i * 2, content_start as u16);
}
write_u16(&mut page, 5, content_start as u16);
debug_assert!(header_end <= content_start);
page[7] = 0; self.push_page(page)
}
fn build_interiors(&mut self, leaves: &[(u32, i64)]) -> u32 {
let mut level: Vec<(u32, i64)> = leaves.to_vec();
while level.len() > 1 {
let mut next_level: Vec<(u32, i64)> = Vec::new();
let mut group: Vec<(u32, i64)> = Vec::new();
for child in &level {
let used: usize = 12
+ group
.iter()
.map(|c| 2 + 4 + enc_varint_into(rowid_as_usize(c.1)).len())
.sum::<usize>();
let need = 2 + 4 + enc_varint_into(rowid_as_usize(child.1)).len();
if !group.is_empty() && used + need > PAGE_SIZE {
next_level.push(self.flush_interior(&group));
group.clear();
}
group.push(*child);
}
if !group.is_empty() {
next_level.push(self.flush_interior(&group));
}
level = next_level;
}
level[0].0
}
fn flush_interior(&mut self, children: &[(u32, i64)]) -> (u32, i64) {
let mut page = vec![0u8; PAGE_SIZE];
page[0] = 0x05; write_u16(&mut page, 1, 0); let Some((rightmost, keyed)) = children.split_last() else {
return (self.push_page(page), 0);
};
write_u32(&mut page, 8, rightmost.0); write_u16(&mut page, 3, keyed.len() as u16);
let mut content_start = PAGE_SIZE;
for (i, (child, key)) in keyed.iter().enumerate() {
let mut cell = Vec::with_capacity(4 + 9);
cell.extend_from_slice(&child.to_be_bytes());
cell.extend_from_slice(&enc_varint_into(rowid_as_usize(*key)));
content_start -= cell.len();
page[content_start..content_start + cell.len()].copy_from_slice(&cell);
write_u16(&mut page, 12 + i * 2, content_start as u16);
}
write_u16(&mut page, 5, content_start as u16);
page[7] = 0;
let page_no = self.push_page(page);
(page_no, rightmost.1)
}
fn finish(mut self) -> Vec<u8> {
let count = self.page_count();
let _ = count;
let mut out = Vec::with_capacity(self.pages.len() * PAGE_SIZE);
for page in self.pages.drain(..) {
out.extend_from_slice(&page);
}
out
}
}
struct LeafCell {
on_page: Vec<u8>,
}
fn encode_record(values: &[Value]) -> Vec<u8> {
let mut serials: Vec<i64> = Vec::with_capacity(values.len());
let mut body: Vec<u8> = Vec::new();
for v in values {
let (serial, bytes) = encode_value(v);
serials.push(serial);
body.extend_from_slice(&bytes);
}
let mut serial_bytes: Vec<u8> = Vec::new();
for &s in &serials {
serial_bytes.extend_from_slice(&enc_varint_into(s as usize));
}
let header_len = resolve_header_len(serial_bytes.len());
let mut out = Vec::with_capacity(header_len + body.len());
out.extend_from_slice(&enc_varint_into(header_len));
out.extend_from_slice(&serial_bytes);
out.extend_from_slice(&body);
out
}
fn resolve_header_len(serial_len: usize) -> usize {
let mut header_len = serial_len + 1;
loop {
let with_varint = serial_len + enc_varint_into(header_len).len();
if with_varint == header_len {
return header_len;
}
header_len = with_varint;
}
}
fn encode_value(v: &Value) -> (i64, Vec<u8>) {
match v {
Value::Null => (0, Vec::new()),
Value::Integer(i) => encode_int(*i),
Value::Real(r) => (7, r.to_bits().to_be_bytes().to_vec()),
Value::Text(t) => (13 + 2 * t.len() as i64, t.as_bytes().to_vec()),
Value::Blob(b) => (12 + 2 * b.len() as i64, b.clone()),
}
}
fn encode_int(i: i64) -> (i64, Vec<u8>) {
match i {
0 => (8, Vec::new()),
1 => (9, Vec::new()),
_ if (i64::from(i8::MIN)..=i64::from(i8::MAX)).contains(&i) => (1, vec![i as u8]),
_ if (i64::from(i16::MIN)..=i64::from(i16::MAX)).contains(&i) => {
(2, (i as i16).to_be_bytes().to_vec())
}
_ if (-(1 << 23)..(1 << 23)).contains(&i) => (3, i.to_be_bytes()[5..].to_vec()),
_ if (i64::from(i32::MIN)..=i64::from(i32::MAX)).contains(&i) => {
(4, (i as i32).to_be_bytes().to_vec())
}
_ if (-(1 << 47)..(1 << 47)).contains(&i) => (5, i.to_be_bytes()[2..].to_vec()),
_ => (6, i.to_be_bytes().to_vec()),
}
}
fn rowid_as_usize(rowid: i64) -> usize {
usize::try_from(rowid).unwrap_or(0)
}
fn build_schema_leaf(tables: &[(&str, u32, &str)], page_count: u32) -> Vec<u8> {
let mut page = vec![0u8; PAGE_SIZE];
write_file_header(&mut page, page_count);
let h = 100;
page[h] = 0x0d; write_u16(&mut page, h + 1, 0); write_u16(&mut page, h + 3, tables.len() as u16); page[h + 7] = 0;
let mut content_start = PAGE_SIZE;
for (i, (name, root, sql)) in tables.iter().enumerate() {
let cell = schema_cell(name, *root, sql, i as i64 + 1);
content_start -= cell.len();
page[content_start..content_start + cell.len()].copy_from_slice(&cell);
write_u16(&mut page, h + 8 + i * 2, content_start as u16); }
debug_assert!(h + 8 + tables.len() * 2 <= content_start);
write_u16(&mut page, h + 5, content_start as u16); page
}
fn schema_cell(name: &str, root: u32, create_sql: &str, rowid: i64) -> Vec<u8> {
let schema_record = encode_record(&[
Value::Text("table".into()),
Value::Text(name.into()),
Value::Text(name.into()),
Value::Integer(i64::from(root)),
Value::Text(create_sql.into()),
]);
let mut cell = Vec::new();
cell.extend_from_slice(&enc_varint_into(schema_record.len()));
cell.extend_from_slice(&enc_varint_into(rowid_as_usize(rowid)));
cell.extend_from_slice(&schema_record);
cell
}
fn write_file_header(page: &mut [u8], page_count: u32) {
page[..16].copy_from_slice(b"SQLite format 3\0");
write_u16(page, 16, PAGE_SIZE as u16);
page[18] = 1; page[19] = 1; page[20] = 0; page[21] = 64; page[22] = 32; page[23] = 32; write_u32(page, 24, 1); write_u32(page, 28, page_count); write_u32(page, 32, 0); write_u32(page, 36, 0); write_u32(page, 40, 1); write_u32(page, 44, 4); write_u32(page, 48, 0); write_u32(page, 52, 0); write_u32(page, 56, 1); write_u32(page, 60, 0); write_u32(page, 64, 0); write_u32(page, 68, 0); write_u32(page, 92, 1); write_u32(page, 96, 3_045_000); }
fn write_u16(buf: &mut [u8], off: usize, val: u16) {
buf[off..off + 2].copy_from_slice(&val.to_be_bytes());
}
fn write_u32(buf: &mut [u8], off: usize, val: u32) {
buf[off..off + 4].copy_from_slice(&val.to_be_bytes());
}