Skip to main content

rete_core/v06/
write.rs

1//! The `0x06` assembler: [`super::parts::Parts`] in, a `0x06` file out. Every
2//! writer goes through here — the in-RAM build (via `0x05` section images),
3//! `rete convert`, and the memory-bounded build and `merge` (which hand it their
4//! spill as an address space, see `extbuild_v06`). `recard` only re-lays the
5//! front.
6//!
7//! The front (root, card region, stats, leaf directories) is built in memory;
8//! the data region is streamed from the input reader, hashed on the way. The
9//! data is written first at its final offset, then the front at offset 0, so the
10//! output needs `Write + Seek` and the input is read once.
11
12use std::io::{Seek, SeekFrom, Write};
13
14use super::parts::Parts;
15use super::*;
16use crate::file::{compress, FileError};
17use crate::header::CURRENT_FORMAT_VERSION;
18use crate::reader::RangeReader;
19
20/// How quad tiles are clustered in the data region.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum QuadClustering {
23    /// Graph-major when the mean graph is smaller than a tile, else order-major.
24    Auto,
25    /// The tiles of every order for one block of graphs side by side.
26    GraphMajor,
27    /// Each order's tiles contiguous.
28    OrderMajor,
29}
30
31impl QuadClustering {
32    /// `RETE_QUAD_CLUSTERING=graph|order|auto` (default auto).
33    pub fn from_env() -> Self {
34        match std::env::var("RETE_QUAD_CLUSTERING").ok().as_deref() {
35            Some("graph") => QuadClustering::GraphMajor,
36            Some("order") => QuadClustering::OrderMajor,
37            _ => QuadClustering::Auto,
38        }
39    }
40}
41
42/// Assembly options.
43#[derive(Debug, Clone, Copy)]
44pub struct WriteOpts {
45    /// Codec of the root, leaves, card and stats (`writer_codec()`).
46    pub internal_codec: u8,
47    /// Quad tile clustering.
48    pub quad: QuadClustering,
49}
50
51impl Default for WriteOpts {
52    fn default() -> Self {
53        WriteOpts {
54            internal_codec: crate::file::writer_codec(),
55            quad: QuadClustering::from_env(),
56        }
57    }
58}
59
60/// What the assembler wrote.
61#[derive(Debug, Clone)]
62pub struct Written {
63    /// The header as written.
64    pub header: Header06,
65    /// The root as written.
66    pub root: Root06,
67}
68
69/// One segment of the data region: input bytes to copy.
70#[derive(Debug, Clone, Copy)]
71struct Seg {
72    input: u64,
73    len: u64,
74    /// The block is stored uncompressed (a raw quad tile).
75    raw: bool,
76}
77
78/// A data section: its kind and the input segments, in output order.
79struct DataSec {
80    kind: u8,
81    segs: Vec<Seg>,
82    /// Each segment is one compressed block of this codec, whose decompressed
83    /// size the writer records ([`Sec06::max_block`]); `None` = not measured.
84    measure: Option<u8>,
85    /// A largest block known without measuring (the dictionary's, from the
86    /// rekey), or 0.
87    known_max: u64,
88}
89
90impl DataSec {
91    fn len(&self) -> u64 {
92        self.segs.iter().map(|s| s.len).sum()
93    }
94}
95
96/// Leaf groups whose leaf size the root budget may grow.
97#[derive(Clone)]
98struct Targets {
99    dict: [usize; 4],
100    index: [usize; 6],
101    quad: [usize; 6],
102    text: usize,
103}
104
105/// The front built for a set of leaf targets.
106struct Front {
107    units: Vec<(u8, Vec<u8>)>,
108    dict_leaves: Vec<u8>,
109    index_leaves: Vec<u8>,
110    quad_leaves: Vec<u8>,
111    text_leaves: Vec<u8>,
112    /// The largest leaf of each leaves section, decompressed: dict, index,
113    /// quad, text.
114    leaf_max: [u64; 4],
115    /// Leaves per group, to pick which group to grow: `(group id, count)`.
116    counts: Vec<(usize, usize)>,
117}
118
119/// Order in which dictionary sections are laid out (predicates first).
120const DICT_ORDER: [usize; 4] = [3, 0, 1, 2];
121
122/// Lay out the data region: the input segments of every data section and, for
123/// the quad index, each tile's new offset within QUAD_DATA.
124#[allow(clippy::type_complexity)]
125fn layout_data(
126    parts: &Parts,
127    graph_major: bool,
128) -> (Vec<DataSec>, [u64; 4], Vec<u64>, Vec<Vec<u64>>) {
129    let mut secs = Vec::new();
130    // Dictionary: sections in DICT_ORDER, each a run of chunks.
131    let mut dict_off = [0u64; 4];
132    let mut segs = Vec::new();
133    let mut at = 0u64;
134    for &s in &DICT_ORDER {
135        dict_off[s] = at;
136        for (c, inp) in &parts.dict[s].chunks {
137            segs.push(Seg {
138                input: *inp,
139                len: c.clen,
140                raw: false,
141            });
142            at += c.clen;
143        }
144    }
145    if at > 0 {
146        secs.push(DataSec {
147            kind: kind::DICT_DATA,
148            segs,
149            measure: None,
150            known_max: parts.dict.iter().map(|d| d.max_raw).max().unwrap_or(0),
151        });
152    }
153    // Index: permutations in section-index order.
154    let mut perm_off = Vec::new();
155    let mut segs = Vec::new();
156    let mut at = 0u64;
157    for p in &parts.index {
158        perm_off.push(at);
159        for (t, inp) in &p.tiles {
160            segs.push(Seg {
161                input: *inp,
162                len: t.clen,
163                raw: false,
164            });
165            at += t.clen;
166        }
167    }
168    if at > 0 {
169        secs.push(DataSec {
170            kind: kind::INDEX_DATA,
171            segs,
172            measure: Some(parts.block_codec),
173            known_max: 0,
174        });
175    }
176    // Quad: order-major, or graph-major by directory block.
177    let mut quad_off: Vec<Vec<u64>> = Vec::new();
178    if let Some(q) = &parts.quad {
179        let mut segs = Vec::new();
180        quad_off = q.orders.iter().map(|(_, t)| vec![0u64; t.len()]).collect();
181        let mut at = 0u64;
182        if graph_major {
183            let firsts = block_first_graphs(&q.p0);
184            // The block a tile belongs to: the last block whose first graph is
185            // <= the tile's first graph.
186            let block_of = |g: u32| firsts.partition_point(|&f| f <= g).saturating_sub(1);
187            let mut cursor = vec![0usize; q.orders.len()];
188            for b in 0..firsts.len().max(1) {
189                for (k, (_, tiles)) in q.orders.iter().enumerate() {
190                    while cursor[k] < tiles.len()
191                        && (firsts.is_empty() || block_of(tiles[cursor[k]].first.0) <= b)
192                    {
193                        let t = &tiles[cursor[k]];
194                        quad_off[k][cursor[k]] = at;
195                        segs.push(Seg {
196                            input: t.off,
197                            len: t.clen,
198                            raw: t.raw,
199                        });
200                        at += t.clen;
201                        cursor[k] += 1;
202                    }
203                }
204            }
205            // Anything left (a tile past the last block's first graph is already
206            // covered above; this only guards a malformed P0).
207            for (k, (_, tiles)) in q.orders.iter().enumerate() {
208                while cursor[k] < tiles.len() {
209                    let t = &tiles[cursor[k]];
210                    quad_off[k][cursor[k]] = at;
211                    segs.push(Seg {
212                        input: t.off,
213                        len: t.clen,
214                        raw: t.raw,
215                    });
216                    at += t.clen;
217                    cursor[k] += 1;
218                }
219            }
220        } else {
221            for (k, (_, tiles)) in q.orders.iter().enumerate() {
222                for (i, t) in tiles.iter().enumerate() {
223                    quad_off[k][i] = at;
224                    segs.push(Seg {
225                        input: t.off,
226                        len: t.clen,
227                        raw: t.raw,
228                    });
229                    at += t.clen;
230                }
231            }
232        }
233        if at > 0 {
234            secs.push(DataSec {
235                kind: kind::QUAD_DATA,
236                segs,
237                measure: Some(parts.block_codec),
238                known_max: 0,
239            });
240        }
241    }
242    if let Some(l) = &parts.labels {
243        if l.data_len > 0 {
244            secs.push(DataSec {
245                kind: kind::LABEL_DATA,
246                segs: vec![Seg {
247                    input: l.body.offset,
248                    len: l.data_len,
249                    raw: false,
250                }],
251                measure: None,
252                known_max: 0,
253            });
254        }
255    }
256    if let Some(t) = &parts.text {
257        if t.postings.len > 0 {
258            secs.push(DataSec {
259                kind: kind::TEXT_DATA,
260                segs: vec![Seg {
261                    input: t.postings.offset,
262                    len: t.postings.len,
263                    raw: false,
264                }],
265                measure: None,
266                known_max: 0,
267            });
268        }
269    }
270    if let Some(py) = parts.pyramid {
271        if py.len > 0 {
272            secs.push(DataSec {
273                kind: kind::PYRAMID,
274                segs: vec![Seg {
275                    input: py.offset,
276                    len: py.len,
277                    raw: false,
278                }],
279                measure: None,
280                known_max: 0,
281            });
282        }
283    }
284    (secs, dict_off, perm_off, quad_off)
285}
286
287/// First graph id of each graph-directory block, from a `P0` block index.
288fn block_first_graphs(p0: &[u8]) -> Vec<u32> {
289    let mut c = Cur::new(p0);
290    let mut out = Vec::new();
291    let Ok(n) = c.count(u32::MAX as u64) else {
292        return out;
293    };
294    let mut g = 0u32;
295    for _ in 0..n {
296        let (Ok(d), Ok(_)) = (c.u32v(), c.uv()) else {
297            break;
298        };
299        g = g.saturating_add(d);
300        out.push(g);
301    }
302    out
303}
304
305/// Compress each leaf of `entries` (cut at `target` encoded bytes) and append
306/// it to `section`. Returns `(leaf ref, first entry index)` per leaf.
307fn leaves_of<E>(
308    entries: &[E],
309    target: usize,
310    codec: u8,
311    section: &mut Vec<u8>,
312    max_raw: &mut u64,
313    encode: impl Fn(&[E], usize) -> Vec<u8>,
314) -> Vec<LeafRef> {
315    let sizes: Vec<usize> = entries
316        .chunks(1)
317        .enumerate()
318        .map(|(i, e)| encode(e, i).len())
319        .collect();
320    let mut out = Vec::new();
321    for (a, b) in cut(&sizes, target) {
322        let raw = encode(&entries[a..b], a);
323        *max_raw = (*max_raw).max(raw.len() as u64);
324        let stored = compress(codec, &raw);
325        out.push(LeafRef {
326            first: a,
327            n: b - a,
328            off: section.len() as u64,
329            len: stored.len() as u64,
330        });
331        section.extend_from_slice(&stored);
332    }
333    out
334}
335
336fn build_front(
337    parts: &Parts,
338    t: &Targets,
339    codec: u8,
340    dict_off: &[u64; 4],
341    perm_off: &[u64],
342    quad_off: &[Vec<u64>],
343    quad_gdir: &[u8],
344) -> Front {
345    let mut f = Front {
346        units: Vec::new(),
347        dict_leaves: Vec::new(),
348        index_leaves: Vec::new(),
349        quad_leaves: Vec::new(),
350        text_leaves: Vec::new(),
351        leaf_max: [0; 4],
352        counts: Vec::new(),
353    };
354    // Dictionary.
355    let mut droots: [DictSecRoot; 4] = Default::default();
356    for &s in &DICT_ORDER {
357        let d = &parts.dict[s];
358        let ents: Vec<ChunkEnt> = d.chunks.iter().map(|(c, _)| c.clone()).collect();
359        let leaves = leaves_of(
360            &ents,
361            t.dict[s],
362            codec,
363            &mut f.dict_leaves,
364            &mut f.leaf_max[0],
365            |e, _| encode_chunk_leaf(e),
366        );
367        // Offset of each chunk within the section's data.
368        let mut coff = Vec::with_capacity(ents.len());
369        let mut at = 0u64;
370        for e in &ents {
371            coff.push(at);
372            at += e.clen;
373        }
374        f.counts.push((s, leaves.len()));
375        droots[s] = DictSecRoot {
376            term_count: d.term_count,
377            restart_interval: d.restart_interval,
378            n_chunks: ents.len(),
379            data_off: dict_off[s],
380            data_len: at,
381            leaves: leaves
382                .into_iter()
383                .map(|l| {
384                    let e = &ents[l.first];
385                    let run = e.first_run;
386                    let key = e.key.clone();
387                    let co = coff[l.first];
388                    (l, run, co, key)
389                })
390                .collect(),
391        };
392    }
393    f.units.push((tag::DICT, encode_dict_root(&droots)));
394    // Index.
395    let mut proots = Vec::new();
396    for (pi, p) in parts.index.iter().enumerate() {
397        let ents: Vec<TileEnt> = p.tiles.iter().map(|(t, _)| *t).collect();
398        let leaves = leaves_of(
399            &ents,
400            t.index[p.si],
401            codec,
402            &mut f.index_leaves,
403            &mut f.leaf_max[1],
404            |e, _| encode_tile_leaf(e),
405        );
406        let mut toff = Vec::with_capacity(ents.len());
407        let mut at = 0u64;
408        for e in &ents {
409            toff.push(at);
410            at += e.clen;
411        }
412        f.counts.push((4 + p.si, leaves.len()));
413        proots.push(PermRoot {
414            si: p.si,
415            n_tiles: ents.len(),
416            data_off: perm_off[pi],
417            data_len: at,
418            leaves: leaves
419                .into_iter()
420                .map(|l| {
421                    let mn = ents[l.first].min_a;
422                    let mx = ents[l.first + l.n - 1].max_a;
423                    let o = toff[l.first];
424                    (l, mn, mx, o)
425                })
426                .collect(),
427        });
428    }
429    f.units.push((tag::INDEX, encode_index_root(&proots)));
430    // Quad.
431    if let Some(q) = &parts.quad {
432        f.quad_leaves.extend_from_slice(quad_gdir);
433        let mut orders = Vec::new();
434        for (k, (si, tiles)) in q.orders.iter().enumerate() {
435            let ents: Vec<QTileEnt> = tiles
436                .iter()
437                .enumerate()
438                .map(|(i, e)| QTileEnt {
439                    off: quad_off[k][i],
440                    ..*e
441                })
442                .collect();
443            let leaves = leaves_of(
444                &ents,
445                t.quad[*si],
446                codec,
447                &mut f.quad_leaves,
448                &mut f.leaf_max[2],
449                |e, i| {
450                    let _ = i;
451                    encode_qtile_leaf(e, e[0].off)
452                },
453            );
454            f.counts.push((10 + si, leaves.len()));
455            orders.push(QuadOrderRoot {
456                si: *si,
457                n_tiles: ents.len(),
458                leaves: leaves
459                    .into_iter()
460                    .map(|l| {
461                        let a = &ents[l.first];
462                        let b = &ents[l.first + l.n - 1];
463                        let o = a.off;
464                        (l, a.first, b.last, o)
465                    })
466                    .collect(),
467            });
468        }
469        f.units.push((
470            tag::QUAD,
471            encode_quad_root(&QuadRoot {
472                graph_count: q.graph_count,
473                quad_count: q.quad_count,
474                block_graphs: q.block_graphs,
475                p0: q.p0.clone(),
476                gdir_len: quad_gdir.len() as u64,
477                orders,
478            }),
479        ));
480    }
481    // Labels.
482    if let Some(l) = &parts.labels {
483        f.units.push((
484            tag::LABELS,
485            encode_label_root(&LabelRoot {
486                data_len: l.data_len,
487                head: l.head.clone(),
488            }),
489        ));
490    }
491    // Text.
492    if let Some(tp) = &parts.text {
493        let leaves = leaves_of(
494            &tp.tokens,
495            t.text,
496            codec,
497            &mut f.text_leaves,
498            &mut f.leaf_max[3],
499            |e, _| encode_token_leaf(e),
500        );
501        f.counts.push((16, leaves.len()));
502        f.units.push((
503            tag::TEXT,
504            encode_text_root(&TextRoot {
505                scope: tp.scope.clone(),
506                n_tokens: tp.tokens.len(),
507                leaves: leaves
508                    .into_iter()
509                    .map(|l| {
510                        let tok = tp.tokens[l.first].token.clone();
511                        (l, tok)
512                    })
513                    .collect(),
514            }),
515        ));
516    }
517    f
518}
519
520/// Copy `segs` from `input` to `out`, coalescing input-adjacent segments into
521/// reads of at most 8 MiB, hashing everything written into `h`.
522#[allow(clippy::too_many_arguments)]
523fn copy_segs(
524    input: &dyn RangeReader,
525    segs: &[Seg],
526    measure: Option<u8>,
527    out: &mut dyn Write,
528    h: &mut blake3::Hasher,
529    progress: &mut dyn FnMut(u64),
530    consumed: &mut dyn FnMut(u64, u64),
531    settle: &mut dyn FnMut(&mut dyn Write) -> std::io::Result<()>,
532) -> Result<u64, FileError> {
533    const SPAN: u64 = 8 << 20;
534    // The ceiling for measuring a block whose frame does not state its size:
535    // generous, since this only bounds the writer's own check.
536    const MEASURE_CEILING: u64 = 1 << 30;
537    let mut max_raw = 0u64;
538    let mut done = 0u64;
539    let mut i = 0;
540    while i < segs.len() {
541        let start = segs[i].input;
542        let mut end = start + segs[i].len;
543        let mut j = i + 1;
544        while j < segs.len() && segs[j].input == end && end - start < SPAN {
545            end += segs[j].len;
546            j += 1;
547        }
548        if let Some(codec) = measure {
549            // Blocks are measured whole, so the run is read in one piece.
550            let b = input.read_at(start, end - start)?;
551            for sg in &segs[i..j] {
552                let lo = (sg.input - start) as usize;
553                let blk = &b[lo..lo + sg.len as usize];
554                let n = if sg.raw {
555                    sg.len
556                } else {
557                    crate::file::block_raw_len(codec, blk, MEASURE_CEILING)?
558                };
559                max_raw = max_raw.max(n);
560            }
561            h.update(&b);
562            out.write_all(&b)?;
563            done += b.len() as u64;
564            progress(done);
565        } else {
566            let mut at = start;
567            while at < end {
568                let n = (end - at).min(SPAN);
569                let b = input.read_at(at, n)?;
570                h.update(&b);
571                out.write_all(&b)?;
572                at += n;
573                done += n;
574                progress(done);
575            }
576        }
577        for sg in &segs[i..j] {
578            consumed(sg.input, sg.len);
579        }
580        settle(out)?;
581        i = j;
582    }
583    Ok(max_raw)
584}
585
586/// The largest graph-directory block of a quad index, decompressed.
587fn gdir_max(p0: &[u8], gdir: &[u8], codec: u8) -> Result<u64, FileError> {
588    let mut c = Cur::new(p0);
589    let n = c.count(u32::MAX as u64)?;
590    let (mut at, mut max) = (0usize, 0u64);
591    for _ in 0..n {
592        let _ = c.u32v()?;
593        let packed = c.uv()?;
594        let len = (packed >> 1) as usize;
595        let blk = gdir
596            .get(at..at + len)
597            .ok_or(FileError::Container("quad graph directory overruns"))?;
598        let raw = if packed & 1 == 1 {
599            len as u64
600        } else {
601            crate::file::block_raw_len(codec, blk, 1 << 30)?
602        };
603        max = max.max(raw);
604        at += len;
605    }
606    Ok(max)
607}
608
609/// Assemble a `0x06` file from `parts`, whose payload ranges point into `input`.
610pub fn assemble<W: Write + Seek>(
611    parts: &Parts,
612    input: &dyn RangeReader,
613    out: &mut W,
614    opts: WriteOpts,
615) -> Result<Written, FileError> {
616    assemble_with(
617        parts,
618        input,
619        out,
620        opts,
621        &mut Hooks {
622            copied: &mut |_| {},
623            progress: &mut |_, _| {},
624            needed: &mut |_, _| {},
625            consumed: &mut |_, _| {},
626            settle: &mut |_| Ok(()),
627        },
628    )
629}
630
631/// What [`assemble_with`] reports while it writes.
632pub struct Hooks<'a> {
633    /// Called once each data section (a [`kind`] of `DICT_DATA`, `INDEX_DATA`,
634    /// `QUAD_DATA`, `LABEL_DATA`, `TEXT_DATA` or `PYRAMID`) is in the output:
635    /// the input bytes it came from are not read again.
636    pub copied: &'a mut dyn FnMut(u8),
637    /// `(data bytes written so far, data bytes in all)`, after each write.
638    pub progress: &'a mut dyn FnMut(u64, u64),
639    /// `(input offset, length)` of every payload range the copy will read,
640    /// announced once before the first byte is written.
641    pub needed: &'a mut dyn FnMut(u64, u64),
642    /// `(input offset, length)` of a payload range now in the output, whatever
643    /// order the ranges come in. Once every announced range of a spill file
644    /// has been consumed, nothing reads that file again.
645    pub consumed: &'a mut dyn FnMut(u64, u64),
646    /// Called with the output after every copied run and every finished data
647    /// section, once `consumed` / `copied` may have made spill files ready to
648    /// go. The implementation flushes `out`, makes the bytes durable, and only
649    /// then deletes: a file must never be removed while its bytes are still in
650    /// a buffer or an unsynced page.
651    pub settle: &'a mut dyn FnMut(&mut dyn Write) -> std::io::Result<()>,
652}
653
654/// [`assemble`], reporting through `hooks`. The memory-bounded build deletes
655/// each spill file as soon as every range it feeds is in the output, so the
656/// output grows as the spill shrinks, and drives `--progress` from the byte
657/// count.
658pub fn assemble_with<W: Write + Seek>(
659    parts: &Parts,
660    input: &dyn RangeReader,
661    out: &mut W,
662    opts: WriteOpts,
663    hooks: &mut Hooks<'_>,
664) -> Result<Written, FileError> {
665    let codec = opts.internal_codec;
666    if parts.pyramid_levels > u8::MAX as u16 {
667        return Err(FileError::Container(
668            "0x06 stores at most 255 pyramid levels",
669        ));
670    }
671    // Ids are dense and 1-based, so no position holds an id above the term
672    // count (graphs: the graph count). Checked, never truncated.
673    let id_widths = super::id_widths_for([
674        parts.term_count,
675        parts.term_count,
676        parts.term_count,
677        parts.quad.as_ref().map_or(0, |q| q.graph_count),
678    ])?;
679    // Quad clustering: graph-major when the mean graph is under a tile.
680    let graph_major = match (&parts.quad, opts.quad) {
681        (None, _) => false,
682        (Some(_), QuadClustering::GraphMajor) => true,
683        (Some(_), QuadClustering::OrderMajor) => false,
684        (Some(q), QuadClustering::Auto) => {
685            let bytes: u64 = q.orders.iter().flat_map(|(_, t)| t).map(|t| t.clen).sum();
686            q.graph_count > 0 && bytes / q.graph_count.max(1) < 64 * 1024
687        }
688    };
689    let (data, dict_off, perm_off, quad_off) = layout_data(parts, graph_major);
690    for d in &data {
691        for sg in &d.segs {
692            (hooks.needed)(sg.input, sg.len);
693        }
694    }
695    let quad_gdir = match &parts.quad {
696        Some(q) if q.gdir.len > 0 => input.read_at(q.gdir.offset, q.gdir.len)?,
697        _ => Vec::new(),
698    };
699    let label_nodes = match &parts.labels {
700        Some(l) if l.body.len > l.data_len => {
701            input.read_at(l.body.offset + l.data_len, l.body.len - l.data_len)?
702        }
703        _ => Vec::new(),
704    };
705
706    // Grow leaves until the stored root fits its budget.
707    let mut t = Targets {
708        dict: [LEAF_START; 4],
709        index: [LEAF_START; 6],
710        quad: [LEAF_START; 6],
711        text: LEAF_START,
712    };
713    let mut front;
714    loop {
715        front = build_front(
716            parts, &t, codec, &dict_off, &perm_off, &quad_off, &quad_gdir,
717        );
718        let probe = Root06 {
719            sections: vec![
720                Sec06 {
721                    kind: 0,
722                    max_block: 0,
723                    offset: 0,
724                    len: 0,
725                    digest: [0; 16]
726                };
727                kind::ORDER.len()
728            ],
729            units: front.units.clone(),
730        };
731        if probe.encode(codec).len() <= ROOT_BUDGET {
732            break;
733        }
734        let Some(&(g, n)) = front.counts.iter().max_by_key(|(_, n)| *n) else {
735            break;
736        };
737        if n <= 1 {
738            break;
739        }
740        match g {
741            0..=3 => t.dict[g] *= 2,
742            4..=9 => t.index[g - 4] *= 2,
743            10..=15 => t.quad[g - 10] *= 2,
744            _ => t.text *= 2,
745        }
746    }
747
748    // Front sections, in file order, with their largest block.
749    let schema = parts
750        .schema
751        .as_ref()
752        .map(|s| (compress(codec, s), s.len() as u64));
753    let quad_gdir_max = match &parts.quad {
754        Some(q) if !quad_gdir.is_empty() => gdir_max(&q.p0, &quad_gdir, parts.block_codec)?,
755        _ => 0,
756    };
757    let mut fronts: Vec<(u8, Vec<u8>, u64)> = Vec::new();
758    if let Some((s, raw)) = schema {
759        fronts.push((kind::SCHEMA, s, raw));
760    }
761    if !front.dict_leaves.is_empty() {
762        fronts.push((
763            kind::DICT_LEAVES,
764            std::mem::take(&mut front.dict_leaves),
765            front.leaf_max[0],
766        ));
767    }
768    if !front.index_leaves.is_empty() {
769        fronts.push((
770            kind::INDEX_LEAVES,
771            std::mem::take(&mut front.index_leaves),
772            front.leaf_max[1],
773        ));
774    }
775    if !front.quad_leaves.is_empty() {
776        fronts.push((
777            kind::QUAD_LEAVES,
778            std::mem::take(&mut front.quad_leaves),
779            front.leaf_max[2].max(quad_gdir_max),
780        ));
781    }
782    if !label_nodes.is_empty() {
783        fronts.push((kind::LABEL_LEAVES, label_nodes, 0));
784    }
785    if !front.text_leaves.is_empty() {
786        fronts.push((
787            kind::TEXT_LEAVES,
788            std::mem::take(&mut front.text_leaves),
789            front.leaf_max[3],
790        ));
791    }
792    let clamp = |n: u64| n.min(u32::MAX as u64) as u32;
793
794    // Card region: card | record | slack.
795    let card = parts
796        .card
797        .as_ref()
798        .filter(|c| !c.is_empty())
799        .map(|c| compress(codec, &stamp_card_version(c)));
800    let build = parts.build.clone().filter(|b| !b.is_empty());
801    let card_len = card.as_ref().map_or(0, |c| c.len()) as u32;
802    let build_len = build.as_ref().map_or(0, |b| b.len()) as u32;
803    let mut region_len = card_region_len(card_len, build_len);
804
805    // Section table (digests filled once the data is hashed).
806    let mut sections = Vec::new();
807    let mut at = 0u64;
808    for (k, b, raw) in &fronts {
809        sections.push(Sec06 {
810            kind: *k,
811            max_block: clamp(*raw),
812            offset: at,
813            len: b.len() as u64,
814            digest: digest16(b),
815        });
816        at += b.len() as u64;
817    }
818    let front_sections_len = at;
819    let mut at = 0u64;
820    for d in &data {
821        sections.push(Sec06 {
822            kind: d.kind,
823            max_block: clamp(d.known_max),
824            offset: at,
825            len: d.len(),
826            digest: [0; 16],
827        });
828        at += d.len();
829    }
830    let data_len = at;
831    let mut root = Root06 {
832        sections,
833        units: front.units,
834    };
835    let root_len = root.encode(codec).len() as u32;
836    // A file is never shorter than the 1 KiB a 0x05-era reader fetches first:
837    // a tiny one gets the difference as card-region slack.
838    let total = HEADER_LEN as u64
839        + root_len as u64
840        + region_len as u64
841        + front_sections_len
842        + data_len
843        + FOOTER.len() as u64;
844    if total < MIN_FILE_LEN {
845        region_len += (MIN_FILE_LEN - total) as u32;
846    }
847    let data_offset = HEADER_LEN as u64 + root_len as u64 + region_len as u64 + front_sections_len;
848
849    // Data first, at its final offset, hashed per section.
850    out.seek(SeekFrom::Start(data_offset))?;
851    let n_front = fronts.len();
852    let mut written_before = 0u64;
853    for (i, d) in data.iter().enumerate() {
854        let mut h = blake3::Hasher::new();
855        let progress = &mut hooks.progress;
856        let consumed = &mut hooks.consumed;
857        let settle = &mut hooks.settle;
858        let measured = copy_segs(
859            input,
860            &d.segs,
861            d.measure,
862            out,
863            &mut h,
864            &mut |n| progress(written_before + n, data_len),
865            &mut **consumed,
866            &mut **settle,
867        )?;
868        written_before += d.len();
869        (hooks.copied)(d.kind);
870        (hooks.settle)(out)?;
871        let sec = &mut root.sections[n_front + i];
872        sec.digest.copy_from_slice(&h.finalize().as_bytes()[..16]);
873        if d.measure.is_some() {
874            sec.max_block = clamp(measured);
875        }
876    }
877    out.write_all(&FOOTER)?;
878    let root_bytes = root.encode(codec);
879    debug_assert_eq!(root_bytes.len() as u32, root_len);
880
881    let mut header = Header06 {
882        flags: (if parts.named { F_NAMED } else { 0 })
883            | (if parts.quoted { F_QUOTED } else { 0 })
884            | (if parts.text.as_ref().is_some_and(|t| t.named) {
885                F_TEXT_NAMED
886            } else {
887                0
888            })
889            | (if graph_major { F_QUAD_GRAPH_MAJOR } else { 0 }),
890        content_hash: [0; 16],
891        file_len: data_offset + data_len + FOOTER.len() as u64,
892        root_len,
893        card_region_len: region_len,
894        data_offset,
895        quad_count: parts.quad_count,
896        term_count: parts.term_count,
897        graph_count: parts.quad.as_ref().map_or(0, |q| q.graph_count),
898        perm_mask: parts.perm_mask,
899        dict_codec: parts.dict_codec,
900        block_codec: parts.block_codec,
901        internal_codec: codec,
902        pyramid_levels: parts.pyramid_levels,
903        id_widths,
904        presence: presence_of(parts),
905        card_off: 0,
906        card_len,
907        card_digest: card.as_ref().map_or([0; 16], |c| digest16(c)),
908        build_off: card_len,
909        build_len,
910        root_digest: digest16(&root_bytes),
911    };
912    header.content_hash = content_hash(&header);
913
914    out.seek(SeekFrom::Start(0))?;
915    out.write_all(&header.to_bytes())?;
916    out.write_all(&root_bytes)?;
917    let mut region = Vec::with_capacity(region_len as usize);
918    if let Some(c) = &card {
919        region.extend_from_slice(c);
920    }
921    if let Some(b) = &build {
922        region.extend_from_slice(b);
923    }
924    region.resize(region_len as usize, 0);
925    out.write_all(&region)?;
926    for (_, b, _) in &fronts {
927        out.write_all(b)?;
928    }
929    out.seek(SeekFrom::Start(header.file_len))?;
930    out.flush()?;
931    Ok(Written { header, root })
932}
933
934/// The smallest `0x06` file: 1 KiB, what every reader reads first anyway.
935pub const MIN_FILE_LEN: u64 = 1024;
936
937/// The card region's length for a card and record of these stored lengths.
938///
939/// With a card: both, twice over, plus [`CARD_SLACK_MIN`] — so a re-card can
940/// write the new pair beside the old one and switch with one header write
941/// (SPEC.md §4.4). With **no card** the region is the record alone: that space
942/// is in the 64 KiB the open reads first, so slack there would only push the
943/// leaf directories out of it, and a file nobody carded yet is re-carded by a
944/// copy. No card and no record: no region.
945pub fn card_region_len(card_len: u32, build_len: u32) -> u32 {
946    match (card_len, build_len) {
947        (0, 0) => 0,
948        (0, b) => b,
949        (c, b) => 2 * (c + b) + CARD_SLACK_MIN,
950    }
951}
952
953fn presence_of(p: &Parts) -> u16 {
954    let mut x = 0u16;
955    if let Some(t) = &p.text {
956        x |= P_TEXT;
957        if t.partial {
958            x |= P_TEXT_PARTIAL;
959        }
960    }
961    if p.labels.is_some() {
962        x |= P_LABELS;
963    }
964    if p.schema.is_some() {
965        x |= P_SCHEMA;
966    }
967    if p.pyramid.is_some_and(|r| r.len > 0) {
968        x |= P_PYRAMID;
969    }
970    if p.quad.as_ref().is_some_and(|q| q.graph_count > 0) {
971        x |= P_QUAD;
972    }
973    x
974}
975
976/// A card written into a `0x06` file says `"format_version":6`. The cards the
977/// core builders derive stamp the `0x05` layout they were produced for
978/// (`LEGACY_FORMAT_VERSION`); every card entering a `0x06` file passes here, so
979/// a build, a converted build and a re-card agree on the stamp. Only the first
980/// compact `"format_version":N` with N below 6 is rewritten; anything else is
981/// returned unchanged.
982pub fn stamp_card_version(card: &[u8]) -> std::borrow::Cow<'_, [u8]> {
983    const KEY: &[u8] = b"\"format_version\":";
984    let Some(at) = card.windows(KEY.len()).position(|w| w == KEY) else {
985        return std::borrow::Cow::Borrowed(card);
986    };
987    let s = at + KEY.len();
988    let digits = card[s..].iter().take_while(|b| b.is_ascii_digit()).count();
989    let old = std::str::from_utf8(&card[s..s + digits]).unwrap_or("");
990    match old.parse::<u8>() {
991        Ok(v) if v < CURRENT_FORMAT_VERSION => {
992            let mut out = card[..s].to_vec();
993            out.extend_from_slice(CURRENT_FORMAT_VERSION.to_string().as_bytes());
994            out.extend_from_slice(&card[s + digits..]);
995            std::borrow::Cow::Owned(out)
996        }
997        _ => std::borrow::Cow::Borrowed(card),
998    }
999}
1000
1001/// Convert a whole `0x05` image held in memory to a `0x06` image.
1002pub fn convert_image(image: &[u8], opts: WriteOpts) -> Result<Vec<u8>, FileError> {
1003    let r = crate::reader::SliceReader::new(image);
1004    let parts = parts::from_v05(&r)?;
1005    let mut out = std::io::Cursor::new(Vec::with_capacity(image.len()));
1006    assemble(&parts, &r, &mut out, opts)?;
1007    Ok(out.into_inner())
1008}
1009
1010/// Convert a `0x05` file read through `input` into `out`.
1011pub fn convert<W: Write + Seek>(
1012    input: &dyn RangeReader,
1013    out: &mut W,
1014    opts: WriteOpts,
1015) -> Result<Written, FileError> {
1016    convert_with_labels(input, None, out, opts)
1017}
1018
1019/// [`convert`], giving the file the LABEL_INDEX section `labels` (its `0x05`
1020/// encoding, from [`label_section`]) when it has none.
1021pub fn convert_with_labels<W: Write + Seek>(
1022    input: &dyn RangeReader,
1023    labels: Option<&[u8]>,
1024    out: &mut W,
1025    opts: WriteOpts,
1026) -> Result<Written, FileError> {
1027    let mut parts = parts::from_v05(input)?;
1028    let Some(sec) = labels.filter(|s| !s.is_empty() && parts.labels.is_none()) else {
1029        return assemble(&parts, input, out, opts);
1030    };
1031    // The new section is read through an overlay that places it right after
1032    // the input, so the assembler copies it like any other.
1033    let overlay = Overlay {
1034        base: input,
1035        extra: sec,
1036    };
1037    parts.labels = Some(parts::label_part(
1038        &overlay,
1039        input.len(),
1040        sec.len() as u64,
1041        parts.block_codec,
1042    )?);
1043    assemble(&parts, &overlay, out, opts)
1044}
1045
1046/// `base`, then `extra` past its end.
1047struct Overlay<'a> {
1048    base: &'a dyn RangeReader,
1049    extra: &'a [u8],
1050}
1051
1052impl RangeReader for Overlay<'_> {
1053    fn len(&self) -> u64 {
1054        self.base.len() + self.extra.len() as u64
1055    }
1056
1057    fn read_at(&self, offset: u64, len: u64) -> std::io::Result<Vec<u8>> {
1058        let b = self.base.len();
1059        if offset + len <= b {
1060            return self.base.read_at(offset, len);
1061        }
1062        if offset >= b {
1063            let s = (offset - b) as usize;
1064            return self
1065                .extra
1066                .get(s..s + len as usize)
1067                .map(<[u8]>::to_vec)
1068                .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "oob"));
1069        }
1070        let mut v = self.base.read_at(offset, b - offset)?;
1071        v.extend_from_slice(self.read_at(b, offset + len - b)?.as_slice());
1072        Ok(v)
1073    }
1074}
1075
1076/// The LABEL_INDEX section (`0x05` encoding) of the default graph of `rete`,
1077/// computed as `rete build --label-index` computes it — from the triples of
1078/// the label predicates only, read through the index. Empty when the graph
1079/// has no label predicate.
1080pub fn label_section(rete: &crate::file::Rete) -> Vec<u8> {
1081    use crate::label_index::{
1082        DIRECT_LABEL_PREDICATES, FAMILY_NAME_PREDICATES, GIVEN_NAME_PREDICATES,
1083    };
1084    let dict = rete.dictionary();
1085    let index = rete.default_index();
1086    let mut triples: Vec<(u32, u32, u32)> = Vec::new();
1087    for list in [
1088        DIRECT_LABEL_PREDICATES,
1089        FAMILY_NAME_PREDICATES,
1090        GIVEN_NAME_PREDICATES,
1091    ] {
1092        for p in list {
1093            if let Some(pid) = dict.predicate_id(p) {
1094                triples.extend(index.scan_iter((None, Some(pid), None)));
1095            }
1096        }
1097    }
1098    triples.sort_unstable();
1099    triples.dedup();
1100    let entries = crate::label_index::entries_from_sorted_triples(dict, &triples);
1101    if entries.is_empty() {
1102        return Vec::new();
1103    }
1104    crate::label_index::encode_section(&entries, crate::file::writer_codec())
1105}
1106
1107/// Replace the card (and the build record) of a `0x06` image held in memory:
1108/// the front is re-laid out with the new card region; every section is copied
1109/// byte for byte. The content hash follows the card.
1110pub fn replace_card(
1111    image: &[u8],
1112    card: Option<&[u8]>,
1113    build: Option<&[u8]>,
1114) -> Result<Vec<u8>, FileError> {
1115    let h = Header06::from_bytes(image)?;
1116    let front = image
1117        .get(..h.data_offset as usize)
1118        .ok_or(FileError::Container("0x06 front overruns"))?;
1119    let data = image
1120        .get(h.data_offset as usize..)
1121        .ok_or(FileError::Container("0x06 data overruns"))?;
1122    let stored = card
1123        .filter(|c| !c.is_empty())
1124        .map(|c| compress(h.internal_codec, &stamp_card_version(c)));
1125    let (_, mut out) = recard_front(&h, front, stored.as_deref(), build)?;
1126    out.extend_from_slice(data);
1127    Ok(out)
1128}
1129
1130/// The card a re-card writes.
1131#[derive(Debug, Clone, Copy)]
1132pub enum NewCard<'a> {
1133    /// The file's own card, byte for byte (only the build record changes, so
1134    /// the content hash does not).
1135    Keep,
1136    /// This card (JSON), or none.
1137    Set(Option<&'a [u8]>),
1138}
1139
1140/// The header and front (everything before the data region) of `h`'s file
1141/// with a new card and build record: the card region re-laid out (no slack
1142/// reuse — a fresh region sized for the new pair), the root and the front
1143/// sections copied, the data region shifted by the region's change.
1144fn recard_front(
1145    h: &Header06,
1146    front: &[u8],
1147    card_stored: Option<&[u8]>,
1148    build: Option<&[u8]>,
1149) -> Result<(Header06, Vec<u8>), FileError> {
1150    let root_bytes = front
1151        .get(HEADER_LEN..HEADER_LEN + h.root_len as usize)
1152        .ok_or(FileError::Container("0x06 root overruns"))?;
1153    let sections = front
1154        .get(h.front_base() as usize..h.data_offset as usize)
1155        .ok_or(FileError::Container("0x06 front overruns"))?;
1156    let card_len = card_stored.map_or(0, |c| c.len()) as u32;
1157    let build_len = build.map_or(0, |b| b.len()) as u32;
1158    let mut region_len = card_region_len(card_len, build_len);
1159    let rest_len = h.file_len - h.front_base();
1160    let total = HEADER_LEN as u64 + h.root_len as u64 + region_len as u64 + rest_len;
1161    if total < MIN_FILE_LEN {
1162        region_len += (MIN_FILE_LEN - total) as u32;
1163    }
1164    let mut nh = h.clone();
1165    nh.card_region_len = region_len;
1166    nh.card_off = 0;
1167    nh.card_len = card_len;
1168    nh.card_digest = card_stored.map_or([0; 16], digest16);
1169    nh.build_off = card_len;
1170    nh.build_len = build_len;
1171    let shift = region_len as i64 - h.card_region_len as i64;
1172    nh.data_offset = (h.data_offset as i64 + shift) as u64;
1173    nh.file_len = (h.file_len as i64 + shift) as u64;
1174    nh.content_hash = content_hash(&nh);
1175    let mut out = Vec::with_capacity(nh.data_offset as usize);
1176    out.extend_from_slice(&nh.to_bytes());
1177    out.extend_from_slice(root_bytes);
1178    let mut region = Vec::with_capacity(region_len as usize);
1179    if let Some(c) = card_stored {
1180        region.extend_from_slice(c);
1181    }
1182    if let Some(b) = build {
1183        region.extend_from_slice(b);
1184    }
1185    region.resize(region_len as usize, 0);
1186    out.extend_from_slice(&region);
1187    out.extend_from_slice(sections);
1188    Ok((nh, out))
1189}
1190
1191/// Rewrite a `0x06` file read through `src` into `out` with a new card and
1192/// build record — [`replace_card`] without holding the file: the front is
1193/// re-laid out around the new card region, the data region copied byte for
1194/// byte (its digests, and so the root, unchanged). `progress` receives the
1195/// data bytes copied so far.
1196pub fn rewrite_card_stream<W: Write>(
1197    src: &dyn RangeReader,
1198    out: &mut W,
1199    card: NewCard<'_>,
1200    build: Option<&[u8]>,
1201    progress: &mut dyn FnMut(u64),
1202) -> Result<Header06, FileError> {
1203    const SPAN: u64 = 8 << 20;
1204    let head = src.read_at(0, HEADER_LEN as u64)?;
1205    let h = Header06::from_bytes(&head)?;
1206    if h.file_len != src.len() {
1207        return Err(FileError::Container(
1208            "0x06 file length does not match its header",
1209        ));
1210    }
1211    // Everything before the data region is small: the header, the root, the
1212    // card region and the front sections.
1213    let front = src.read_at(0, h.data_offset)?;
1214    let stored = match card {
1215        NewCard::Keep => {
1216            let r = h.card_range();
1217            (r.len > 0).then(|| front[r.offset as usize..(r.offset + r.len) as usize].to_vec())
1218        }
1219        NewCard::Set(c) => c
1220            .filter(|c| !c.is_empty())
1221            .map(|c| compress(h.internal_codec, &stamp_card_version(c))),
1222    };
1223    let (nh, new_front) = recard_front(&h, &front, stored.as_deref(), build)?;
1224    out.write_all(&new_front)?;
1225    let mut at = h.data_offset;
1226    while at < h.file_len {
1227        let n = (h.file_len - at).min(SPAN);
1228        out.write_all(&src.read_at(at, n)?)?;
1229        at += n;
1230        progress(at - h.data_offset);
1231    }
1232    out.flush()?;
1233    Ok(nh)
1234}
1235
1236/// Re-card a `0x06` file **in place**, when the new card and record fit in a
1237/// free gap of its card region (beside the old pair, which stays readable until
1238/// the switch): they are written there and the new header is returned, *not*
1239/// written. The caller syncs the file, then [`commit_header`] — the 128-byte
1240/// header is the one commit point, so a crash before it leaves the old file
1241/// whole. A card equal to the current one stays where it is (only the record
1242/// moves; the content hash does not change). `Ok(None)` when the pair does not
1243/// fit: the caller copies instead ([`rewrite_card_stream`]).
1244pub fn recard_in_place<F: std::io::Read + Write + Seek>(
1245    file: &mut F,
1246    card: NewCard<'_>,
1247    build: Option<&[u8]>,
1248) -> Result<Option<Header06>, FileError> {
1249    let mut head = vec![0u8; HEADER_LEN];
1250    file.seek(SeekFrom::Start(0))?;
1251    file.read_exact(&mut head)?;
1252    let h = Header06::from_bytes(&head)?;
1253    let card_stored = match card {
1254        NewCard::Keep => None,
1255        NewCard::Set(c) => c
1256            .filter(|c| !c.is_empty())
1257            .map(|c| compress(h.internal_codec, &stamp_card_version(c))),
1258    };
1259    let keep_card = match (card, &card_stored) {
1260        (NewCard::Keep, _) => true,
1261        (_, Some(c)) => c.len() as u32 == h.card_len && digest16(c) == h.card_digest,
1262        (_, None) => h.card_len == 0,
1263    };
1264    let card_need = if keep_card {
1265        0
1266    } else {
1267        card_stored.as_ref().map_or(0, |c| c.len()) as u64
1268    };
1269    let build_need = build.map_or(0, |b| b.len()) as u64;
1270    let need = card_need + build_need;
1271    // The free gaps of the region around the pair in use.
1272    let mut used: Vec<(u64, u64)> = [
1273        (h.card_off as u64, h.card_len as u64),
1274        (h.build_off as u64, h.build_len as u64),
1275    ]
1276    .into_iter()
1277    .filter(|&(_, l)| l > 0)
1278    .collect();
1279    used.sort_unstable();
1280    let region = h.card_region_len as u64;
1281    let mut gaps = Vec::new();
1282    let mut at = 0u64;
1283    for (o, l) in used {
1284        if o > at {
1285            gaps.push((at, o - at));
1286        }
1287        at = at.max(o + l);
1288    }
1289    if region > at {
1290        gaps.push((at, region - at));
1291    }
1292    let start = if need == 0 {
1293        0
1294    } else {
1295        match gaps.iter().find(|&&(_, l)| l >= need) {
1296            Some(&(o, _)) => o,
1297            None => return Ok(None),
1298        }
1299    };
1300    file.seek(SeekFrom::Start(h.card_region_offset() + start))?;
1301    if !keep_card {
1302        if let Some(c) = &card_stored {
1303            file.write_all(c)?;
1304        }
1305    }
1306    if let Some(b) = build {
1307        file.write_all(b)?;
1308    }
1309    file.flush()?;
1310    let mut nh = h.clone();
1311    if !keep_card {
1312        nh.card_off = if card_need > 0 { start as u32 } else { 0 };
1313        nh.card_len = card_need as u32;
1314        nh.card_digest = card_stored.as_ref().map_or([0; 16], |c| digest16(c));
1315    }
1316    nh.build_off = if build_need > 0 {
1317        (start + card_need) as u32
1318    } else {
1319        0
1320    };
1321    nh.build_len = build_need as u32;
1322    nh.content_hash = content_hash(&nh);
1323    Ok(Some(nh))
1324}
1325
1326/// Write `h` over the first 128 bytes (the commit of [`recard_in_place`]).
1327pub fn commit_header<F: Write + Seek>(file: &mut F, h: &Header06) -> Result<(), FileError> {
1328    file.seek(SeekFrom::Start(0))?;
1329    file.write_all(&h.to_bytes())?;
1330    file.flush()?;
1331    Ok(())
1332}
1333
1334#[cfg(test)]
1335mod stamp_tests {
1336    use super::stamp_card_version;
1337
1338    #[test]
1339    fn stamps_older_versions_only() {
1340        let s = |b: &[u8]| stamp_card_version(b).into_owned();
1341        assert_eq!(
1342            s(br#"{"a":1,"format_version":5}"#),
1343            br#"{"a":1,"format_version":6}"#
1344        );
1345        assert_eq!(
1346            s(br#"{"format_version":5,"z":2}"#),
1347            br#"{"format_version":6,"z":2}"#
1348        );
1349        assert_eq!(s(br#"{"format_version":6}"#), br#"{"format_version":6}"#);
1350        assert_eq!(s(br#"{"format_version":7}"#), br#"{"format_version":7}"#);
1351        assert_eq!(s(br#"{"title":"x"}"#), br#"{"title":"x"}"#);
1352    }
1353}