1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum QuadClustering {
23 Auto,
25 GraphMajor,
27 OrderMajor,
29}
30
31impl QuadClustering {
32 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#[derive(Debug, Clone, Copy)]
44pub struct WriteOpts {
45 pub internal_codec: u8,
47 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#[derive(Debug, Clone)]
62pub struct Written {
63 pub header: Header06,
65 pub root: Root06,
67}
68
69#[derive(Debug, Clone, Copy)]
71struct Seg {
72 input: u64,
73 len: u64,
74 raw: bool,
76}
77
78struct DataSec {
80 kind: u8,
81 segs: Vec<Seg>,
82 measure: Option<u8>,
85 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#[derive(Clone)]
98struct Targets {
99 dict: [usize; 4],
100 index: [usize; 6],
101 quad: [usize; 6],
102 text: usize,
103}
104
105struct 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 leaf_max: [u64; 4],
115 counts: Vec<(usize, usize)>,
117}
118
119const DICT_ORDER: [usize; 4] = [3, 0, 1, 2];
121
122#[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 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 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 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 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 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
287fn 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
305fn 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 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 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 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 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 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 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#[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 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 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
586fn 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
609pub 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
631pub struct Hooks<'a> {
633 pub copied: &'a mut dyn FnMut(u8),
637 pub progress: &'a mut dyn FnMut(u64, u64),
639 pub needed: &'a mut dyn FnMut(u64, u64),
642 pub consumed: &'a mut dyn FnMut(u64, u64),
646 pub settle: &'a mut dyn FnMut(&mut dyn Write) -> std::io::Result<()>,
652}
653
654pub 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 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 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 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 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 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 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 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 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(®ion)?;
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
934pub const MIN_FILE_LEN: u64 = 1024;
936
937pub 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
976pub 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
1001pub 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
1010pub 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
1019pub 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 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
1046struct 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
1076pub 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
1107pub 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#[derive(Debug, Clone, Copy)]
1132pub enum NewCard<'a> {
1133 Keep,
1136 Set(Option<&'a [u8]>),
1138}
1139
1140fn 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(®ion);
1187 out.extend_from_slice(sections);
1188 Ok((nh, out))
1189}
1190
1191pub 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 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
1236pub 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 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
1326pub 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}