Skip to main content

lumen/
writer.rs

1//! Writing a file: [`Encoder`], which assembles a conforming `.lumen` container.
2//!
3//! The encoder owns the choices the specification leaves to it (section 5.7): it
4//! picks a tag per layer, picks how many layers a `LAYR` chunk spans, and decides
5//! whether to train a dictionary. Layer masks are encoded as they are pushed, so
6//! the encoder never holds the pixel data of the whole print - only its run-end
7//! encoded form.
8//!
9//! The layer axis is cut into groups of [`DEFAULT_LAYERS_PER_CHUNK`] layers (or
10//! whatever [`Encoder::set_layers_per_chunk`] asks for), and each (sector,
11//! layer-group) becomes one `LAYR` chunk holding that sector's masks for the
12//! group as a single zstd frame. `LTBL` then records, per (layer, sector), which
13//! chunk holds the layer's run, where in that chunk's decompressed output the
14//! run sits, and which `LROV` chunk carries the (layer, sector)'s overrides.
15//!
16//! Two of `finish`'s passes are a function of one item each: one zstd frame per
17//! (sector, layer-group) chunk, and one `LHAS` leaf per layer. Neither reads
18//! another's result, so both run across worker threads
19//! ([`Encoder::set_worker_threads`]) - which is where the time of a long print
20//! sits once its masks are encoded. The file does not depend on how many: an
21//! item is placed by the index it came from, not by the order it finished in.
22//!
23//! `finish` produces the container in the order section 3 recommends: `HEAD`,
24//! `META`, the optional chunks, `LTBL`, `LAYR`, and the directory at the end.
25//! Payloads are 8-byte aligned, as section 3.2 prefers, and the trailer's
26//! CRC-32C covers every preceding byte.
27//!
28//! Encryption follows section 9.1 exactly: `HEAD`, `AUTH`, `LTBL`, `LHAS` and the
29//! `LAYR` version fields stay plaintext; `META`, `PROF`, `LROV`, `VOXL` and
30//! `ZDIC` are sealed as single units; each `LAYR` frame is its own unit, bound by
31//! its associated data to the chunk's directory index. `PREV` and `EXTD` are
32//! written in the clear, which section 9.1 permits and which keeps a thumbnail
33//! usable without a key.
34
35use crate::check::Check;
36use crate::chunks::extd::Extension;
37use crate::chunks::head::Head;
38use crate::chunks::lhas::{self, LayerHashes};
39use crate::chunks::ltbl::{LayerEntry, LayerTable};
40use crate::chunks::preview::PreviewRole;
41use crate::chunks::zdic::ZstdDictionary;
42use crate::chunks::{self, json_chunks, layr};
43use crate::container::{
44    self, ChunkDescriptor, ChunkType, Directory, FileHeader, CHUNK_FLAG_ENCRYPTED, FLAG_ENCRYPTED,
45    FLAG_MULTI_SECTOR,
46};
47use crate::crypto::{self, Auth, Cipher, RecipientEntry, SessionKey};
48use crate::error::{Error, Result};
49use crate::io::Writer;
50use crate::json::{Meta, Profile, Sector, Timing};
51use crate::ree::{self, EncodeMode, Run};
52use std::collections::HashMap;
53use std::sync::atomic::{AtomicUsize, Ordering};
54
55/// Argon2id parameters for password-mode encryption.
56///
57/// The defaults sit well inside the ceilings a reader is required to enforce
58/// (ยง11.4: iterations <= 10, memory <= 4 GiB, lanes <= 16), so a file this crate
59/// writes is one every conforming reader must be willing to open. The encoder
60/// does not police the values it is handed: a caller who deliberately asks for
61/// more is writing a file that readers MUST refuse rather than attempt, and the
62/// conformance corpus contains exactly such a file on purpose.
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub struct Argon2Params {
65    /// Time cost.
66    pub iterations: u32,
67    /// Memory cost in KiB.
68    pub memory_kib: u32,
69    /// Lanes.
70    pub parallelism: u8,
71}
72
73impl Default for Argon2Params {
74    fn default() -> Self {
75        // Print-file decryption happens on printer hardware, so the cost has to
76        // land between "a weak password is cheap to attack" and "the printer
77        // takes a minute to start". A caller who wants more can say so.
78        Argon2Params {
79            iterations: 3,
80            memory_kib: 65_536,
81            parallelism: 1,
82        }
83    }
84}
85
86/// How to encrypt the file being written.
87#[derive(Debug, Clone)]
88pub struct EncryptOptions {
89    /// The AEAD cipher.
90    pub cipher: Cipher,
91    /// Wrap the session key under this password.
92    pub password: Option<String>,
93    /// Wrap the session key for each of these X25519 public keys.
94    pub recipients: Vec<[u8; 32]>,
95    /// Argon2id parameters, used when a password is set.
96    pub argon2: Argon2Params,
97}
98
99impl EncryptOptions {
100    /// Password-mode encryption with default Argon2id parameters.
101    pub fn password(password: impl Into<String>) -> EncryptOptions {
102        EncryptOptions {
103            cipher: Cipher::Aes256Gcm,
104            password: Some(password.into()),
105            recipients: Vec::new(),
106            argon2: Argon2Params::default(),
107        }
108    }
109}
110
111/// The layers one `LAYR` chunk spans, which a fresh encoder starts with
112/// (section 4.9 recommends 32-64).
113pub const DEFAULT_LAYERS_PER_CHUNK: u32 = 64;
114/// The zstd level a fresh encoder starts with.
115pub const DEFAULT_ZSTD_LEVEL: i32 = 6;
116/// The zstd level for the small JSON chunks (section 6.3).
117pub const JSON_ZSTD_LEVEL: i32 = 3;
118/// How many 1 KiB windows the dictionary may be sampled from (section 6.2).
119///
120/// The unit is a window, not a layer: this constant was read as a layer count and
121/// applied to the windows of the print's *first* chunks, which made the training set
122/// the print's bottom layers - 256 KiB, a couple of layers at 12K.
123pub const DICTIONARY_SAMPLE_WINDOWS: usize = 256;
124
125/// The largest dictionary the specification allows, in bytes.
126const DICTIONARY_MAX_BYTES: usize = 112_640;
127
128/// The most bytes the pays-probe compresses in each direction.
129const DICTIONARY_PROBE_BYTES: usize = 4 * 1024 * 1024;
130
131/// The plaintext size below which the probe is not trusted, in bytes.
132///
133/// On a small print the probe's noise and the dictionary's own bytes are the same
134/// order of magnitude, so the writer keeps the specification's behavior and writes
135/// the dictionary it trained. Above this, the data decides.
136const DICTIONARY_PROBE_MIN_PLAINTEXT: usize = 8 * 1024 * 1024;
137
138/// The margin a dictionary must beat by, as a fraction of the probe, to be written.
139const DICTIONARY_PROBE_MARGIN: f64 = 0.002;
140
141/// One `(layer, sector)`'s timing delta, written as its own `LROV` chunk.
142///
143/// A delta replaces only the fields it carries; every field it leaves out keeps
144/// the value the pipeline resolved for that (layer, sector) from `META`, the
145/// sector's `META.sectors` entry and the bottom/transition blend.
146#[derive(Debug, Clone, PartialEq)]
147pub struct Override {
148    /// The layer the delta applies to.
149    pub layer: u32,
150    /// The sector it applies to. `0` is the primary sector.
151    pub sector_id: u32,
152    /// The fields the delta replaces, in META's names and units.
153    pub timing: Timing,
154}
155
156/// A layer's encoded form, ready to store: what `push_layer` builds internally.
157#[derive(Debug, Clone)]
158pub enum EncodedLayer {
159    /// The empty layer: no sector carries bytes.
160    Empty,
161    /// `(sector_id, tag plus REE stream)`, ascending by `sector_id`, with the
162    /// empty sectors omitted.
163    Sectors(Vec<(u32, Vec<u8>)>),
164}
165
166impl EncodedLayer {
167    /// One sector's stream, as sector 0 of a layer.
168    ///
169    /// `stream` is exactly what [`ree::encode`] or [`ree::encode_runs`] returned:
170    /// the REE tag, then the stream. A stream with no tag cannot be stored - an
171    /// all-black sector is the empty-layer form - so it is rejected here rather
172    /// than reaching the layer table.
173    pub fn single(stream: Vec<u8>) -> Result<Self> {
174        check_stream(&stream)?;
175        Ok(EncodedLayer::Sectors(vec![(0, stream)]))
176    }
177
178    /// The pushed stream for one sector, when this layer has data for it.
179    fn stream(&self, sector_id: u32) -> Option<&[u8]> {
180        match self {
181            EncodedLayer::Empty => None,
182            EncodedLayer::Sectors(list) => list
183                .iter()
184                .find(|(id, _)| *id == sector_id)
185                .map(|(_, stream)| stream.as_slice()),
186        }
187    }
188}
189
190/// Reject a stream that does not begin with a layer encoding tag.
191///
192/// An empty stream is not degenerate but different: it is the empty-layer form,
193/// which [`EncodedLayer::Empty`] carries and no sector's data does. The messages
194/// match the reader's, so a caller sees the same words either side of the file.
195/// A single-sector layer record for a maybe-empty encoding.
196///
197/// `None` is the empty-layer form, and the returned mask data carries its tag,
198/// which is what a layer stores.
199fn single_sector(encoded: Option<(u8, Vec<u8>)>) -> EncodedLayer {
200    match encoded {
201        None => EncodedLayer::Empty,
202        Some((_tag, mask)) => EncodedLayer::Sectors(vec![(0, mask)]),
203    }
204}
205
206fn check_stream(stream: &[u8]) -> Result<()> {
207    match stream.first() {
208        None => Err(Error::new(
209            Check::ReeTag,
210            "layer mask data is empty: the empty-layer form carries no bytes",
211        )),
212        Some(&(ree::TAG_BINARY | ree::TAG_GRAYSCALE | ree::TAG_SPLIT | ree::TAG_ATTACHED)) => {
213            Ok(())
214        }
215        Some(&tag) => Err(Error::new(
216            Check::ReeTag,
217            format!("unknown layer encoding tag 0x{tag:02X}"),
218        )),
219    }
220}
221
222/// A layer group's plaintext under one candidate encoding, with the entries that
223/// place each layer's slice inside it.
224struct ChunkCandidate {
225    plaintext: Vec<u8>,
226    placement: Vec<(u32, u64, u32)>,
227}
228
229/// A chunk ready to be laid out.
230#[derive(Debug)]
231struct Pending {
232    chunk_type: ChunkType,
233    /// The bytes as they will sit on disk.
234    stored: Vec<u8>,
235    /// What the descriptor's `size_uncompressed` field reports.
236    size_uncompressed: u64,
237    /// The stored length when it differs from the raw payload, else zero.
238    size_compressed: u64,
239    flags: u32,
240}
241
242/// Builds a `.lumen` file.
243#[derive(Debug)]
244pub struct Encoder {
245    head: Head,
246    meta: Meta,
247    profile: Option<Profile>,
248    overrides: Vec<Override>,
249    previews: Vec<(PreviewRole, Vec<u8>)>,
250    extensions: Vec<Extension>,
251    voxl: Option<Vec<u8>>,
252    layers_per_chunk: u32,
253    zstd_level: i32,
254    dictionary: bool,
255    layer_hashes: bool,
256    worker_threads: usize,
257    encryption: Option<EncryptOptions>,
258    session: Option<SessionKey>,
259    /// Whether each layer group is encoded both ways and the smaller kept
260    /// ([`Encoder::set_tag_probe`]).
261    tag_probe: bool,
262    /// One entry per layer: its mask data as [`EncodeMode::Auto`] chose it, and
263    /// the alternative tag 0x03 encoding the probe compares it against, for a
264    /// layer that was pushed under `Auto` while the probe was on.
265    layers: Vec<(EncodedLayer, Option<EncodedLayer>)>,
266}
267
268impl Encoder {
269    /// Start a file with the given `HEAD` and `META`.
270    ///
271    /// `head.total_layers` must equal the number of layers pushed.
272    pub fn new(head: Head, meta: Meta) -> Encoder {
273        Encoder {
274            head,
275            meta,
276            profile: None,
277            overrides: Vec::new(),
278            previews: Vec::new(),
279            extensions: Vec::new(),
280            voxl: None,
281            layers_per_chunk: DEFAULT_LAYERS_PER_CHUNK,
282            zstd_level: DEFAULT_ZSTD_LEVEL,
283            dictionary: true,
284            layer_hashes: true,
285            worker_threads: 0,
286            encryption: None,
287            session: None,
288            tag_probe: true,
289            layers: Vec::new(),
290        }
291    }
292
293    /// Attach a reusable print profile.
294    pub fn set_profile(&mut self, profile: Profile) {
295        self.profile = Some(profile);
296    }
297
298    /// Declare the non-zero sectors this print uses, as `META.sectors`.
299    ///
300    /// Sector 0 is implicit and must not appear; a sector carries no material
301    /// until its entry names one.
302    pub fn set_sectors(&mut self, sectors: Vec<Sector>) {
303        self.meta.sectors = Some(sectors);
304    }
305
306    /// Attach the per-(layer, sector) timing deltas.
307    ///
308    /// Each pair's matching `LTBL` entry points at an `LROV` chunk, and pairs
309    /// whose deltas are equal share one: a delta set on a range - or on any set of
310    /// pairs - costs one chunk, and every entry that names it applies it. A layer
311    /// emitted here also gets a table entry for that sector even when the layer
312    /// holds no data for it, which is how an override reaches a sector a layer does
313    /// not print.
314    pub fn set_overrides(&mut self, overrides: Vec<Override>) -> Result<()> {
315        let mut overrides = overrides;
316        overrides.sort_by_key(|over| (over.layer, over.sector_id));
317        for pair in overrides.windows(2) {
318            if pair[0].layer == pair[1].layer && pair[0].sector_id == pair[1].sector_id {
319                return Err(Error::new(
320                    Check::LrovOrphan,
321                    format!(
322                        "layer {} sector {} carries two override sets; a (layer, sector) has one \
323                         or none",
324                        pair[0].layer, pair[0].sector_id
325                    ),
326                ));
327            }
328        }
329        self.overrides = overrides;
330        Ok(())
331    }
332
333    /// Add a PNG preview.
334    pub fn add_preview(&mut self, role: PreviewRole, png: Vec<u8>) {
335        self.previews.push((role, png));
336    }
337
338    /// Add a vendor or standard extension.
339    pub fn add_extension(&mut self, extension: Extension) {
340        self.extensions.push(extension);
341    }
342
343    /// Embed the scene that produced this print, verbatim.
344    pub fn set_voxl(&mut self, voxl: Vec<u8>) {
345        self.voxl = Some(voxl);
346    }
347
348    /// How many layers one `LAYR` chunk spans.
349    pub fn set_layers_per_chunk(&mut self, layers: u32) {
350        self.layers_per_chunk = layers.max(1);
351    }
352
353    /// Whether to encode each layer group both ways and keep the smaller.
354    ///
355    /// On by default. A group's layers are encoded with the tag
356    /// [`EncodeMode::Auto`] chooses and with [`EncodeMode::Attached`], both are
357    /// compressed exactly as the file will frame them, and the smaller result is
358    /// kept. The comparison is what makes the attached form a gain rather than a
359    /// gamble: its raw size is not a guide to its compressed size, so an encoder
360    /// that picked it by bytes alone would lose on a print whose layers barely
361    /// change.
362    ///
363    /// The probe costs one extra encoding and one extra compression of every
364    /// layer group, and it measures without the `ZDIC` dictionary, which is only
365    /// trained once the whole print has been pushed. Turning it off means every
366    /// layer is stored with the tag [`EncodeMode::Auto`] picks.
367    pub fn set_tag_probe(&mut self, enabled: bool) {
368        self.tag_probe = enabled;
369    }
370
371    /// zstd compression level for the layer chunks.
372    pub fn set_zstd_level(&mut self, level: i32) {
373        self.zstd_level = level;
374    }
375
376    /// Whether to train and emit a `ZDIC` dictionary (section 6.2).
377    ///
378    /// Enabled, the writer samples the print, trains a dictionary and writes it *if
379    /// the data says it pays* - see [`dictionary_pays`]. Disabled, it compresses every
380    /// frame without one, which is a smaller file than a dictionary that does not fit
381    /// the print.
382    pub fn set_dictionary(&mut self, enabled: bool) {
383        self.dictionary = enabled;
384    }
385
386    /// Whether to emit an `LHAS` integrity tree.
387    pub fn set_layer_hashes(&mut self, enabled: bool) {
388        self.layer_hashes = enabled;
389    }
390
391    /// How many threads `finish` may frame chunks and hash layers on.
392    ///
393    /// `0`, what a fresh encoder starts with, means one per available core, and
394    /// a request above that is capped by it: this is CPU-bound work on data
395    /// already in memory, so a thread beyond the cores available would take time
396    /// from another rather than spend a wait. `1` is the serial path, and it
397    /// writes the same bytes any other count does - the output of this encoder
398    /// is a function of its input and settings alone, as section 5.7 intends.
399    pub fn set_worker_threads(&mut self, threads: usize) {
400        self.worker_threads = threads;
401    }
402
403    /// The worker count `finish` will use, never zero.
404    fn workers(&self) -> usize {
405        match self.worker_threads {
406            0 => available_parallelism(),
407            threads => threads.min(available_parallelism()),
408        }
409    }
410
411    /// Encrypt the file.
412    pub fn set_encryption(&mut self, options: EncryptOptions) {
413        self.session = Some(SessionKey::generate());
414        self.encryption = Some(options);
415    }
416
417    /// The session key this encoder uses, for a caller that wants to keep it.
418    pub fn session_key(&self) -> Option<SessionKey> {
419        self.session
420    }
421
422    /// Push one layer's mask, choosing its encoding per [`EncodeMode::Auto`].
423    ///
424    /// The mask is sector 0's: it is what a single-material reader prints.
425    pub fn push_layer(&mut self, pixels: &[u8]) -> Result<()> {
426        self.push_layer_with_mode(pixels, EncodeMode::Auto)
427    }
428
429    /// Push one layer's mask with an explicit encoding.
430    pub fn push_layer_with_mode(&mut self, pixels: &[u8], mode: EncodeMode) -> Result<()> {
431        let total_pixels = self.total_pixels();
432        check_mask(pixels, total_pixels, "the layer")?;
433        let record = single_sector(ree::encode(pixels, total_pixels, mode)?);
434        let alternative = match (mode, self.tag_probe) {
435            (EncodeMode::Auto, true) => Some(single_sector(ree::encode(
436                pixels,
437                total_pixels,
438                EncodeMode::Attached,
439            )?)),
440            _ => None,
441        };
442        self.layers.push((record, alternative));
443        Ok(())
444    }
445
446    /// Push one layer as a set of sectors, each with its own mask.
447    ///
448    /// Sector 0 is implicit, so a layer whose only content is sector 0 is pushed
449    /// with [`Encoder::push_layer`] instead. A sector with no exposed pixel is
450    /// omitted, and a layer whose every sector is empty is stored as the
451    /// empty-layer form. The pairs are sorted by sector, which is the order the
452    /// layer table requires.
453    pub fn push_layer_sectors(&mut self, sectors: &[(u32, Vec<u8>)]) -> Result<()> {
454        let total_pixels = self.total_pixels();
455        let mut masks: Vec<(u32, Vec<u8>)> = Vec::new();
456        let mut alternatives: Vec<(u32, Vec<u8>)> = Vec::new();
457        for (sector_id, pixels) in sectors {
458            check_mask(
459                pixels,
460                total_pixels,
461                &format!("sector {sector_id} of the layer"),
462            )?;
463            if sectors.iter().filter(|(id, _)| id == sector_id).count() > 1 {
464                return Err(Error::new(
465                    Check::LtblSectorIdUnique,
466                    format!("sector {sector_id} appears twice in one layer"),
467                ));
468            }
469            if let Some((_tag, mask)) = ree::encode(pixels, total_pixels, EncodeMode::Auto)? {
470                masks.push((*sector_id, mask));
471                if self.tag_probe {
472                    if let Some((_tag, other)) =
473                        ree::encode(pixels, total_pixels, EncodeMode::Attached)?
474                    {
475                        alternatives.push((*sector_id, other));
476                    }
477                }
478            }
479        }
480        if masks.is_empty() {
481            self.layers.push((EncodedLayer::Empty, None));
482        } else {
483            masks.sort_by_key(|(sector_id, _)| *sector_id);
484            alternatives.sort_by_key(|(sector_id, _)| *sector_id);
485            let alternative =
486                (!alternatives.is_empty()).then_some(EncodedLayer::Sectors(alternatives));
487            self.layers
488                .push((EncodedLayer::Sectors(masks), alternative));
489        }
490        Ok(())
491    }
492
493    /// Push one layer given as runs instead of pixels, as sector 0.
494    ///
495    /// Equivalent to [`Encoder::push_layer`] on the mask the runs describe,
496    /// except that the mask is never built: a slicer's run-length rasterizer
497    /// output goes straight to the encoder. The runs must be canonical and cover
498    /// the layer, as [`ree::encode_runs`] requires.
499    pub fn push_layer_runs(&mut self, runs: &[Run], mode: EncodeMode) -> Result<()> {
500        let total_pixels = self.total_pixels();
501        let record = single_sector(ree::encode_runs(runs, total_pixels, mode)?);
502        let alternative = match (mode, self.tag_probe) {
503            (EncodeMode::Auto, true) => Some(single_sector(ree::encode_runs(
504                runs,
505                total_pixels,
506                EncodeMode::Attached,
507            )?)),
508            _ => None,
509        };
510        self.layers.push((record, alternative));
511        Ok(())
512    }
513
514    /// Push a layer that is already encoded.
515    ///
516    /// This is the parallel path: workers turn one layer's runs into bytes with
517    /// [`ree::encode_runs`], and the writer stores each result in layer order
518    /// without ever seeing pixel data. The streams are checked to carry a layer
519    /// encoding tag, the sectors are sorted by id as the layer table requires,
520    /// and a layer with no sector data becomes the empty-layer form.
521    pub fn push_encoded_layer(&mut self, layer: EncodedLayer) -> Result<()> {
522        self.push_encoded_layer_pair(layer, None)
523    }
524
525    /// Push a layer that is already encoded, with the other encoding the tag probe
526    /// should weigh it against.
527    ///
528    /// [`Encoder::push_encoded_layer`] is this with no alternative. A caller that
529    /// encodes layers on its own threads - the slicer adapter does, one worker per
530    /// layer - can hand over the tag `0x03` stream of the same layer as well, and
531    /// the probe then has two candidates to compress when the group closes. Both
532    /// must describe the same pixels: the probe's choice is a choice of encoding,
533    /// and nothing checks that they agree.
534    ///
535    /// The alternative is stored as it is when the probe is off, which costs the
536    /// caller nothing but the encoding it computed.
537    pub fn push_encoded_layer_pair(
538        &mut self,
539        layer: EncodedLayer,
540        alternative: Option<EncodedLayer>,
541    ) -> Result<()> {
542        if let Some(other) = alternative.as_ref() {
543            match other {
544                EncodedLayer::Empty => {}
545                EncodedLayer::Sectors(list) => {
546                    for (_, stream) in list {
547                        check_stream(stream)?;
548                    }
549                }
550            }
551        }
552        let record = match layer {
553            EncodedLayer::Empty => EncodedLayer::Empty,
554            EncodedLayer::Sectors(mut list) => {
555                for (_, stream) in &list {
556                    check_stream(stream)?;
557                }
558                list.sort_by_key(|(sector_id, _)| *sector_id);
559                for pair in list.windows(2) {
560                    if pair[0].0 == pair[1].0 {
561                        return Err(Error::new(
562                            Check::LtblSectorIdUnique,
563                            format!("sector {} appears twice in one layer", pair[0].0),
564                        ));
565                    }
566                }
567                if list.is_empty() {
568                    EncodedLayer::Empty
569                } else {
570                    EncodedLayer::Sectors(list)
571                }
572            }
573        };
574        self.layers.push((record, alternative));
575        Ok(())
576    }
577
578    fn total_pixels(&self) -> u32 {
579        self.head.total_pixels()
580    }
581
582    /// Assemble the container.
583    ///
584    /// Fails if the number of pushed layers disagrees with `head.total_layers`.
585    pub fn finish(self) -> Result<Vec<u8>> {
586        let pushed = self.layers.len() as u32;
587        if pushed != self.head.total_layers {
588            return Err(Error::new(
589                Check::HeadTotalLayers,
590                format!(
591                    "HEAD declares {} layers but {pushed} were pushed",
592                    self.head.total_layers
593                ),
594            ));
595        }
596
597        // 1. META's own shape rules, so the encoder cannot write a file its own
598        //    validator would reject.
599        if let Some(sectors) = self.meta.sectors.as_deref() {
600            let mut seen: Vec<u32> = Vec::with_capacity(sectors.len());
601            for (i, sector) in sectors.iter().enumerate() {
602                if sector.sector_id == 0 {
603                    return Err(Error::new(
604                        Check::MetaSectorsShape,
605                        format!(
606                            "META.sectors[{i}] uses sector_id 0, which names the implicit primary \
607                             sector"
608                        ),
609                    ));
610                }
611                if seen.contains(&sector.sector_id) {
612                    return Err(Error::new(
613                        Check::MetaSectorsShape,
614                        format!("sector_id {} appears twice", sector.sector_id),
615                    ));
616                }
617                seen.push(sector.sector_id);
618            }
619        }
620
621        // 2. Which sectors each layer's table entry set covers: sector 0, every
622        //    sector it holds data for, and every sector an override targets.
623        let sector_sets = self.sector_sets()?;
624
625        // 3. One LAYR chunk per (sector, layer-group): the group's masks for that
626        //    sector, concatenated in ascending layer order. The placement of each
627        //    layer's run inside its chunk is what LTBL records.
628        let group_size = self.layers_per_chunk.max(1) as usize;
629        // Per (group, sector): the plaintext the file frames, the placement of
630        // each layer's slice inside it, and - when the probe computed one - the
631        // same pair for the tag 0x03 encoding of the very same layers. The choice
632        // between them is made in step 5, with the dictionary the frames will
633        // actually use: an encoding's raw size is no guide, and neither is its
634        // compressed size without the dictionary that will compress it.
635        let mut chunk_plaintexts: Vec<Vec<u8>> = Vec::new();
636        let mut chunk_placements: Vec<Vec<(u32, u64, u32)>> = Vec::new();
637        let mut chunk_alternatives: Vec<Option<ChunkCandidate>> = Vec::new();
638        let mut chunk_sectors: Vec<u32> = Vec::new();
639        for group_start in (0..pushed as usize).step_by(group_size) {
640            let group_end = (group_start + group_size).min(pushed as usize);
641            let mut sectors: Vec<u32> = sector_sets[group_start..group_end]
642                .iter()
643                .flatten()
644                .copied()
645                .collect();
646            sectors.sort_unstable();
647            sectors.dedup();
648            let probed = self.layers[group_start..group_end]
649                .iter()
650                .any(|(_, alternative)| alternative.is_some());
651            for sector in sectors {
652                let mut plaintext = Vec::new();
653                let mut placed: Vec<(u32, u64, u32)> = Vec::new();
654                let mut alternative_text = Vec::new();
655                let mut alternative_placed: Vec<(u32, u64, u32)> = Vec::new();
656                for (index, (record, alternative)) in self
657                    .layers
658                    .iter()
659                    .enumerate()
660                    .take(group_end)
661                    .skip(group_start)
662                {
663                    if let Some(mask) = record.stream(sector) {
664                        placed.push((index as u32, plaintext.len() as u64, mask.len() as u32));
665                        plaintext.extend_from_slice(mask);
666                    }
667                    // A layer without an alternative is its own alternative, so
668                    // the two candidates always describe the same layers.
669                    let chosen = alternative.as_ref().unwrap_or(record);
670                    if let Some(mask) = chosen.stream(sector) {
671                        alternative_placed.push((
672                            index as u32,
673                            alternative_text.len() as u64,
674                            mask.len() as u32,
675                        ));
676                        alternative_text.extend_from_slice(mask);
677                    }
678                }
679                // A sector enters a layer's entries for other reasons too - an
680                // override with no data behind it - so a group may name a sector
681                // that holds nothing here. It gets no chunk.
682                if plaintext.is_empty() {
683                    continue;
684                }
685                chunk_plaintexts.push(plaintext);
686                chunk_placements.push(placed);
687                chunk_alternatives.push(
688                    (probed && alternative_text != chunk_plaintexts[chunk_plaintexts.len() - 1])
689                        .then_some(ChunkCandidate {
690                            plaintext: alternative_text,
691                            placement: alternative_placed,
692                        }),
693                );
694                chunk_sectors.push(sector);
695            }
696        }
697        // A file carries at least one LAYR chunk even when every layer is empty
698        // (section 3's presence rule).
699        if chunk_plaintexts.is_empty() {
700            chunk_plaintexts.push(Vec::new());
701            chunk_placements.push(Vec::new());
702            chunk_alternatives.push(None);
703            chunk_sectors.push(0);
704        }
705
706        // 4. The dictionary: trained on samples spread across the print, and written
707        //    only when it earns its own bytes back (section 6.2).
708        let dictionary = if self.dictionary {
709            let samples = dictionary_samples(&chunk_plaintexts);
710            match chunks::train_dictionary(&samples, DICTIONARY_MAX_BYTES) {
711                Ok(bytes) if !bytes.is_empty() && chunks::dictionary_id(&bytes) != 0 => {
712                    if dictionary_pays(&chunk_plaintexts, &bytes, self.zstd_level)? {
713                        Some(ZstdDictionary {
714                            zdic_version: 1,
715                            dict_id: chunks::dictionary_id(&bytes),
716                            dict_bytes: bytes,
717                        })
718                    } else {
719                        // A dictionary that makes the frames larger is bytes nobody
720                        // asked for, and the specification permits omitting it.
721                        None
722                    }
723                }
724                // Training fails on a degenerate print - too little or too
725                // uniform data. The specification's answer is to omit ZDIC and
726                // compress without a dictionary.
727                _ => None,
728            }
729        } else {
730            None
731        };
732        let dict_bytes = dictionary.as_ref().map(|d| d.dict_bytes.as_slice());
733
734        // 5. Compress each chunk's plaintext into its frame, on the worker
735        //    threads: one frame is a function of one plaintext and the
736        //    dictionary, and no frame is a function of another.
737        let workers = self.workers();
738        let level = self.zstd_level;
739        let alternatives = &chunk_alternatives;
740        let choices: Vec<(Vec<u8>, bool)> =
741            map_ordered(&chunk_plaintexts, workers, |index, plaintext| {
742                let frame = |text: &[u8]| -> Result<Vec<u8>> {
743                    let frame = chunks::compress(text, level, dict_bytes)?;
744                    // A reader sizes the frame's output from the frame's own
745                    // header, so a writer must set its content size; refusing here
746                    // keeps the guarantee in the writer rather than in a comment.
747                    chunks::frame_content_size(&frame)?;
748                    Ok(frame)
749                };
750                let primary = frame(plaintext)?;
751                match &alternatives[index] {
752                    // The smaller frame wins, and a tie keeps the encoding
753                    // `EncodeMode::Auto` chose.
754                    Some(candidate) => match frame(&candidate.plaintext) {
755                        Ok(other) if other.len() < primary.len() => Ok((other, true)),
756                        Ok(_) | Err(_) => Ok((primary, false)),
757                    },
758                    None => Ok((primary, false)),
759                }
760            })?;
761
762        // The placement a reader needs is the chosen encoding's, and LTBL is the
763        // only thing that says where a slice sits.
764        let mut placement: HashMap<(u32, u32), (u32, u64, u32)> = HashMap::new();
765        let mut frames: Vec<Vec<u8>> = Vec::with_capacity(choices.len());
766        for (index, (frame, chose_alternative)) in choices.into_iter().enumerate() {
767            let sector = chunk_sectors[index];
768            let entries = if chose_alternative {
769                let candidate = chunk_alternatives[index]
770                    .take()
771                    .expect("an alternative that was compressed");
772                chunk_plaintexts[index] = candidate.plaintext;
773                candidate.placement
774            } else {
775                std::mem::take(&mut chunk_placements[index])
776            };
777            for (layer, offset, size) in entries {
778                placement.insert((layer, sector), (index as u32, offset, size));
779            }
780            frames.push(frame);
781        }
782
783        // 6. LHAS, over each layer's slices concatenated in ascending sector_id -
784        //    read back out of the chunk plaintexts, so the leaves are over the
785        //    very bytes a reader will reconstruct - and on the worker threads,
786        //    a leaf being a function of one layer and nothing else.
787        let hashes = if self.layer_hashes {
788            let leaves = map_ordered(&sector_sets, workers, |index, sectors| {
789                let mut data = Vec::new();
790                for sector in sectors {
791                    let key = (index as u32, *sector);
792                    if let Some(&(ordinal, offset, size)) = placement.get(&key) {
793                        let start = offset as usize;
794                        data.extend_from_slice(
795                            &chunk_plaintexts[ordinal as usize][start..start + size as usize],
796                        );
797                    }
798                }
799                Ok(lhas::leaf_hash(&data))
800            })?;
801            Some(LayerHashes {
802                hash_algorithm: lhas::HASH_ALGORITHM_SHA256,
803                hash_size: lhas::HASH_SIZE_SHA256,
804                layer_count: leaves.len() as u32,
805                merkle_root: lhas::merkle_root(&leaves),
806                layer_hashes: leaves,
807            })
808        } else {
809            None
810        };
811
812        // 7. Everything before LTBL, then LTBL, then everything after: the LAYR
813        //    chunks need their directory indices, and LTBL needs them too, so the
814        //    one slot LTBL occupies and the optional LHAS slot are what the LAYR
815        //    indices are computed from.
816        let mut flags = 0u32;
817        if self
818            .layers
819            .iter()
820            .any(|(record, _)| matches!(record, EncodedLayer::Sectors(list) if list.len() > 1))
821        {
822            flags |= FLAG_MULTI_SECTOR;
823        }
824        let sealed = self.session.is_some() && self.encryption.is_some();
825        if self.encryption.is_some() {
826            flags |= FLAG_ENCRYPTED;
827        }
828
829        let mut pending = Vec::new();
830        pending.push(Pending::raw(ChunkType::HEAD, self.head.to_bytes()));
831        pending.push(self.seal_if_needed(
832            ChunkType::META,
833            json_chunks::meta_to_bytes(&self.meta)?,
834            true,
835        )?);
836        if let Some(profile) = self.profile.as_ref() {
837            pending.push(self.seal_if_needed(
838                ChunkType::PROF,
839                json_chunks::profile_to_bytes(profile)?,
840                true,
841            )?);
842        }
843        if let Some(auth) = self.build_auth()? {
844            pending.push(Pending::raw(ChunkType::AUTH, crypto::encode_auth(&auth)));
845        }
846        // The overrides are already sorted by (layer, sector), so their chunks are
847        // too, and each (layer, sector) remembers which index it landed at. Two pairs
848        // whose deltas are identical share one chunk: a delta that covers a range, or
849        // any set of pairs, is written once and every entry that names it applies it
850        // (section 4.5). The map is keyed by the serialized payload, so a slicer that
851        // builds each pair's delta the same way gets the sharing for free.
852        let mut lrov_index: HashMap<(u32, u32), u32> = HashMap::new();
853        let mut shared: HashMap<Vec<u8>, u32> = HashMap::new();
854        for over in &self.overrides {
855            let payload = json_chunks::lrov_to_bytes(&over.timing)?;
856            let index = match shared.get(&payload) {
857                Some(index) => *index,
858                None => {
859                    let index = pending.len() as u32;
860                    pending.push(self.seal_if_needed(ChunkType::LROV, payload.clone(), true)?);
861                    shared.insert(payload, index);
862                    index
863                }
864            };
865            lrov_index.insert((over.layer, over.sector_id), index);
866        }
867        if let Some(dictionary) = dictionary.as_ref() {
868            pending.push(self.seal_if_needed(ChunkType::ZDIC, dictionary.to_bytes(), false)?);
869        }
870        for (role, png) in &self.previews {
871            pending.push(Pending::raw_with_flags(
872                ChunkType::PREV,
873                png.clone(),
874                role.to_flags(),
875            ));
876        }
877        if let Some(voxl) = self.voxl.as_ref() {
878            pending.push(self.seal_if_needed(ChunkType::VOXL, voxl.clone(), true)?);
879        }
880
881        let layr_base = pending.len() as u32 + 1 + u32::from(hashes.is_some());
882        let entries =
883            self.table_entries(pushed, &sector_sets, &placement, layr_base, &lrov_index)?;
884        let table = LayerTable::new(entries, pushed)?;
885        pending.push(Pending::raw(ChunkType::LTBL, table.to_bytes()));
886        if let Some(hashes) = hashes.as_ref() {
887            pending.push(Pending::raw(ChunkType::LHAS, hashes.to_bytes()));
888        }
889        debug_assert_eq!(
890            pending.len() as u32,
891            layr_base,
892            "the LAYR chunks start where their indices say they do"
893        );
894
895        for (ordinal, frame) in frames.iter().enumerate() {
896            let directory_index = layr_base + ordinal as u32;
897            let container = if sealed {
898                let sealed_frame = crypto::seal(
899                    self.cipher()?,
900                    self.session.as_ref().expect("set with encryption"),
901                    ChunkType::LAYR,
902                    directory_index,
903                    frame,
904                )?;
905                layr::to_bytes(&sealed_frame)
906            } else {
907                layr::to_bytes(frame)
908            };
909            pending.push(Pending {
910                chunk_type: ChunkType::LAYR,
911                stored: container,
912                size_uncompressed: layr::container_len(frame.len()),
913                size_compressed: if sealed {
914                    // The container's stored length: the version field, the
915                    // frame, and the one sealed unit's nonce and tag.
916                    layr::container_len(frame.len()) + crypto::UNIT_OVERHEAD as u64
917                } else {
918                    0
919                },
920                flags: if sealed { CHUNK_FLAG_ENCRYPTED } else { 0 },
921            });
922        }
923        for extension in &self.extensions {
924            pending.push(Pending::raw_with_flags(
925                ChunkType::EXTD,
926                extension.to_bytes(),
927                extension.to_flags(),
928            ));
929        }
930
931        // 8. Lay out: header, payloads at 8-byte alignment, directory, trailer.
932        let mut out = Writer::with_capacity(
933            pending.iter().map(|p| p.stored.len() + 8).sum::<usize>() + 64 + pending.len() * 32,
934        );
935        out.bytes(&[0u8; 32]);
936        let mut directory = Directory::new();
937        let mut total_uncompressed = 0u64;
938        for chunk in &pending {
939            out.pad_to(8);
940            let offset = out.len() as u64;
941            out.bytes(&chunk.stored);
942            total_uncompressed += chunk.size_uncompressed;
943            directory.push(ChunkDescriptor {
944                chunk_type: chunk.chunk_type,
945                offset,
946                size_uncompressed: chunk.size_uncompressed,
947                size_compressed: chunk.size_compressed,
948                flags: chunk.flags,
949            });
950        }
951        out.pad_to(8);
952        let dir_offset = out.len() as u64;
953        out.bytes(&directory.to_bytes());
954        let header = FileHeader {
955            version: 1,
956            dir_offset,
957            chunk_count: directory.descriptors.len() as u32,
958            flags,
959            total_uncompressed_size: total_uncompressed,
960        };
961        out.patch_bytes(0, &header.to_bytes());
962        let crc = container::crc32c(out.as_slice());
963        out.bytes(&container::TRAILER_MAGIC);
964        out.u32(crc);
965        Ok(out.into_vec())
966    }
967
968    /// The sectors each layer's entries cover, ascending, sector 0 first.
969    fn sector_sets(&self) -> Result<Vec<Vec<u32>>> {
970        let mut sets: Vec<Vec<u32>> = self
971            .layers
972            .iter()
973            .map(|(record, _)| match record {
974                EncodedLayer::Empty => Vec::new(),
975                EncodedLayer::Sectors(list) => list.iter().map(|(id, _)| *id).collect(),
976            })
977            .collect();
978        for over in &self.overrides {
979            let index = over.layer as usize;
980            let Some(ids) = sets.get_mut(index) else {
981                return Err(Error::new(
982                    Check::LtblLayerIndexRange,
983                    format!(
984                        "an override names layer {} of {}",
985                        over.layer,
986                        self.layers.len()
987                    ),
988                ));
989            };
990            if !ids.contains(&over.sector_id) {
991                ids.push(over.sector_id);
992            }
993        }
994        for ids in &mut sets {
995            if !ids.contains(&0) {
996                ids.push(0);
997            }
998            ids.sort_unstable();
999            ids.dedup();
1000        }
1001        Ok(sets)
1002    }
1003
1004    /// The layer table's entries, one per (layer, sector), in table order.
1005    fn table_entries(
1006        &self,
1007        pushed: u32,
1008        sector_sets: &[Vec<u32>],
1009        placement: &HashMap<(u32, u32), (u32, u64, u32)>,
1010        layr_base: u32,
1011        lrov_index: &HashMap<(u32, u32), u32>,
1012    ) -> Result<Vec<LayerEntry>> {
1013        let mut entries = Vec::new();
1014        for index in 0..pushed {
1015            let ids = &sector_sets[index as usize];
1016            let additional = (ids.len() - 1) as u32;
1017            // Every entry names a `LAYR` chunk, as section 4.9 requires, even
1018            // when the (layer, sector) holds no bytes: an entry with no run of
1019            // its own takes the chunk of the layer's first run, or the file's
1020            // first chunk when the layer has none at all.
1021            let home = ids
1022                .iter()
1023                .find_map(|sector| placement.get(&(index, *sector)))
1024                .map_or(layr_base, |(ordinal, _, _)| layr_base + *ordinal);
1025            for (position, sector) in ids.iter().enumerate() {
1026                let run = placement.get(&(index, *sector));
1027                entries.push(LayerEntry {
1028                    data_size: run.map_or(0, |(_, _, size)| *size),
1029                    first_lrov: lrov_index.get(&(index, *sector)).copied().unwrap_or(0),
1030                    first_layr: run.map_or(home, |(ordinal, _, _)| layr_base + *ordinal),
1031                    additional_sector_count: if position == 0 { additional } else { 0 },
1032                    data_offset: run.map_or(0, |(_, offset, _)| *offset),
1033                    sector_id: *sector,
1034                });
1035            }
1036        }
1037        Ok(entries)
1038    }
1039
1040    fn cipher(&self) -> Result<Cipher> {
1041        Ok(self
1042            .encryption
1043            .as_ref()
1044            .expect("cipher asked for without encryption")
1045            .cipher)
1046    }
1047
1048    /// Compress and seal a content chunk when the file is encrypted.
1049    ///
1050    /// These are one-unit chunks, so their associated data names slot zero; a
1051    /// `LAYR` frame is sealed by the layout step instead, bound to the chunk's
1052    /// directory index.
1053    fn seal_if_needed(
1054        &self,
1055        chunk_type: ChunkType,
1056        plaintext: Vec<u8>,
1057        compress: bool,
1058    ) -> Result<Pending> {
1059        let payload = if compress {
1060            chunks::compress(&plaintext, JSON_ZSTD_LEVEL, None)?
1061        } else {
1062            plaintext.clone()
1063        };
1064        let compressed_len = if compress { payload.len() as u64 } else { 0 };
1065        match self.session.as_ref() {
1066            None => Ok(Pending {
1067                chunk_type,
1068                stored: payload,
1069                size_uncompressed: plaintext.len() as u64,
1070                size_compressed: compressed_len,
1071                flags: 0,
1072            }),
1073            Some(key) => {
1074                let sealed = crypto::seal(self.cipher()?, key, chunk_type, 0, &payload)?;
1075                Ok(Pending {
1076                    chunk_type,
1077                    size_uncompressed: plaintext.len() as u64,
1078                    stored: sealed,
1079                    size_compressed: payload.len() as u64 + crypto::UNIT_OVERHEAD as u64,
1080                    flags: CHUNK_FLAG_ENCRYPTED,
1081                })
1082            }
1083        }
1084    }
1085
1086    /// Build the `AUTH` chunk, when the file is encrypted.
1087    fn build_auth(&self) -> Result<Option<Auth>> {
1088        let Some(options) = self.encryption.as_ref() else {
1089            return Ok(None);
1090        };
1091        let session = self.session.as_ref().expect("set with encryption");
1092        let mut mode = 0u32;
1093        let mut password = None;
1094        let mut recipients: Vec<RecipientEntry> = Vec::new();
1095        if let Some(secret) = options.password.as_deref() {
1096            mode |= crypto::MODE_PASSWORD;
1097            let params = options.argon2;
1098            password = Some(crypto::wrap_password(
1099                session,
1100                secret,
1101                params.iterations,
1102                params.memory_kib,
1103                params.parallelism,
1104            )?);
1105        }
1106        for public in &options.recipients {
1107            mode |= crypto::MODE_MACHINE;
1108            recipients.push(crypto::wrap_machine(session, public)?);
1109        }
1110        if mode == 0 {
1111            return Err(Error::new(
1112                Check::CryptModeEmpty,
1113                "encryption was requested with neither a password nor a recipient",
1114            ));
1115        }
1116        Ok(Some(Auth {
1117            cipher: options.cipher,
1118            auth_version: 1,
1119            mode,
1120            password: password.map(Option::Some).unwrap_or(None),
1121            recipients,
1122        }))
1123    }
1124}
1125
1126/// One worker per core the platform reports, and one where it reports none.
1127///
1128/// `wasm32-unknown-unknown` is a target this crate is built for - DragonFruit
1129/// renders its UI with it - and it can neither report parallelism nor spawn a
1130/// thread, so an encoder running there takes the serial path however many workers
1131/// it is asked for.
1132fn available_parallelism() -> usize {
1133    #[cfg(target_arch = "wasm32")]
1134    {
1135        1
1136    }
1137    #[cfg(not(target_arch = "wasm32"))]
1138    {
1139        std::thread::available_parallelism().map_or(1, |threads| threads.get())
1140    }
1141}
1142
1143/// Map `items` through `f` on up to `workers` threads, keeping the input order.
1144///
1145/// This is what `finish` runs its per-item passes on. Work is handed out one item
1146/// at a time rather than as a block per thread: the items are not equally
1147/// expensive - a dense layer's zstd frame costs many times an empty one's - so a
1148/// block-sized slice of the print would leave every other worker waiting behind
1149/// whichever one of them drew the hard chunk.
1150///
1151/// A result is placed by the index it came from and never by the order it
1152/// completed in, so the returned vector is the one a single thread would have
1153/// built. `f` sees the item's index because an item may need to say where it
1154/// sits: an `LHAS` leaf names the layer it hashes.
1155fn map_ordered<T, R>(
1156    items: &[T],
1157    workers: usize,
1158    f: impl Fn(usize, &T) -> Result<R> + Sync,
1159) -> Result<Vec<R>>
1160where
1161    T: Sync,
1162    R: Send,
1163{
1164    if workers <= 1 || items.len() <= 1 {
1165        return items
1166            .iter()
1167            .enumerate()
1168            .map(|(index, item)| f(index, item))
1169            .collect();
1170    }
1171    let next = AtomicUsize::new(0);
1172    let mut ordered: Vec<Option<R>> = (0..items.len()).map(|_| None).collect();
1173    std::thread::scope(|scope| -> Result<()> {
1174        let handles: Vec<_> = (0..workers.min(items.len()))
1175            .map(|_| {
1176                scope.spawn(|| -> Result<Vec<(usize, R)>> {
1177                    let mut finished = Vec::new();
1178                    loop {
1179                        // Relaxed: the counter orders nothing but itself, and
1180                        // every value it hands out is filled by whoever took it.
1181                        let index = next.fetch_add(1, Ordering::Relaxed);
1182                        let Some(item) = items.get(index) else {
1183                            break;
1184                        };
1185                        finished.push((index, f(index, item)?));
1186                    }
1187                    Ok(finished)
1188                })
1189            })
1190            .collect();
1191        for handle in handles {
1192            for (index, value) in handle.join().expect("a worker does not panic")? {
1193                ordered[index] = Some(value);
1194            }
1195        }
1196        Ok(())
1197    })?;
1198    Ok(ordered
1199        .into_iter()
1200        .map(|value| value.expect("every index is taken by one worker"))
1201        .collect())
1202}
1203
1204/// Reject a mask whose pixel count disagrees with the display.
1205fn check_mask(pixels: &[u8], total_pixels: u32, what: &str) -> Result<()> {
1206    if pixels.len() != total_pixels as usize {
1207        return Err(Error::new(
1208            Check::ReeDataSize,
1209            format!(
1210                "{what} has {} pixels but the display holds {total_pixels}",
1211                pixels.len()
1212            ),
1213        ));
1214    }
1215    Ok(())
1216}
1217
1218impl Pending {
1219    fn raw(chunk_type: ChunkType, payload: Vec<u8>) -> Pending {
1220        let len = payload.len() as u64;
1221        Pending {
1222            chunk_type,
1223            stored: payload,
1224            size_uncompressed: len,
1225            size_compressed: 0,
1226            flags: 0,
1227        }
1228    }
1229
1230    fn raw_with_flags(chunk_type: ChunkType, payload: Vec<u8>, flags: u32) -> Pending {
1231        let mut pending = Pending::raw(chunk_type, payload);
1232        pending.flags = flags;
1233        pending
1234    }
1235}
1236
1237/// Split a chunk's plaintext into samples for dictionary training.
1238///
1239/// The boundaries between layers are not recoverable from the concatenation
1240/// alone, so this yields the whole chunk as one sample per fixed stride, which is
1241/// what `ZDICT_trainFromBuffer` wants: many samples of similar data. Layer data is
1242/// already highly redundant, so a sample per 1 KiB window trains well and costs
1243/// nothing to compute.
1244fn split_samples(plaintext: &[u8]) -> Vec<&[u8]> {
1245    const WINDOW: usize = 1024;
1246    if plaintext.is_empty() {
1247        return Vec::new();
1248    }
1249    plaintext.chunks(WINDOW).collect()
1250}
1251
1252/// Training samples spread across the whole print (section 6.2).
1253///
1254/// The samples span the print - evenly over the chunks, and evenly within each - rather
1255/// than coming from its start, because the start of a print is its raft and its bottom
1256/// layers. Measured on a 12K print: samples taken from the start compressed 3% worse
1257/// than *no* dictionary on a binary print and 5% worse on an anti-aliased one, while
1258/// samples spread across the print were the best set of the three. A dictionary is
1259/// worth its own 112 KiB only when it describes the data it is used on.
1260fn dictionary_samples(plaintexts: &[Vec<u8>]) -> Vec<&[u8]> {
1261    let windows: Vec<Vec<&[u8]>> = plaintexts.iter().map(|p| split_samples(p)).collect();
1262    let mut samples = Vec::with_capacity(DICTIONARY_SAMPLE_WINDOWS);
1263    if windows.is_empty() {
1264        return samples;
1265    }
1266    // With more chunks than the budget allows, visit a stride of them and take one
1267    // window each; with fewer, take an equal share from every chunk.
1268    let chunk_stride = windows.len().div_ceil(DICTIONARY_SAMPLE_WINDOWS).max(1);
1269    let per_chunk = (DICTIONARY_SAMPLE_WINDOWS / (windows.len() / chunk_stride).max(1)).max(1);
1270    for (index, chunk) in windows.iter().enumerate() {
1271        if index % chunk_stride != 0 || chunk.is_empty() {
1272            continue;
1273        }
1274        let step = (chunk.len() / per_chunk).max(1);
1275        for window in (0..chunk.len()).step_by(step).take(per_chunk) {
1276            if samples.len() == DICTIONARY_SAMPLE_WINDOWS {
1277                return samples;
1278            }
1279            samples.push(chunk[window]);
1280        }
1281    }
1282    samples
1283}
1284
1285/// Whether a trained dictionary earns its own bytes back on this print (section 6.2).
1286///
1287/// The answer is a property of the data rather than of the format. An anti-aliased
1288/// print's greyscale planes are high-entropy, and a dictionary built from the print's
1289/// own statistics made its frames 4.7% *larger* on a measured 800-layer 12K print,
1290/// while a binary print's run-length planes are matched by one. So the writer
1291/// compresses a bounded probe of the print both ways - a share of the first, middle and
1292/// last chunks, so the probe spans the print - and keeps the dictionary only when the
1293/// probe says it wins by more than noise, with the dictionary's own bytes charged
1294/// against it at the share of the print the probe covers. Below
1295/// [`DICTIONARY_PROBE_MIN_PLAINTEXT`] the probe is not trusted and the dictionary stands.
1296fn dictionary_pays(plaintexts: &[Vec<u8>], dictionary: &[u8], level: i32) -> Result<bool> {
1297    let total: usize = plaintexts.iter().map(Vec::len).sum();
1298    if total < DICTIONARY_PROBE_MIN_PLAINTEXT {
1299        return Ok(true);
1300    }
1301    let share = DICTIONARY_PROBE_BYTES / 3;
1302    let last = plaintexts.len() - 1;
1303    let mut probe = Vec::with_capacity(DICTIONARY_PROBE_BYTES);
1304    for index in [0, last / 2, last] {
1305        let plaintext = &plaintexts[index];
1306        probe.extend_from_slice(&plaintext[..plaintext.len().min(share)]);
1307    }
1308    if probe.is_empty() {
1309        return Ok(true);
1310    }
1311    let with = chunks::compress(&probe, level, Some(dictionary))?.len();
1312    let without = chunks::compress(&probe, level, None)?.len();
1313    // The dictionary is stored once and the probe is one slice of the print, so its
1314    // own bytes are charged at the share of the print the probe stands for.
1315    let charge = dictionary.len() as f64 * probe.len() as f64 / total as f64;
1316    Ok((with as f64 + charge) < without as f64 * (1.0 - DICTIONARY_PROBE_MARGIN))
1317}
1318
1319#[cfg(test)]
1320mod tests {
1321    use super::*;
1322    use crate::json::{Meta, Timing};
1323
1324    const WIDTH: u32 = 8;
1325    const HEIGHT: u32 = 4;
1326    const PIXELS: usize = (WIDTH * HEIGHT) as usize;
1327
1328    fn head(layers: u32) -> Head {
1329        Head {
1330            head_version: 1,
1331            encoder_name: "lumen writer tests".to_string(),
1332            created_unix_sec: 1_750_000_000,
1333            display_width_px: WIDTH,
1334            display_height_px: HEIGHT,
1335            build_width_um: 40_000,
1336            build_depth_um: 20_000,
1337            build_height_um: 20_000,
1338            layer_height_um: 50,
1339            total_layers: layers,
1340        }
1341    }
1342
1343    fn meta() -> Meta {
1344        let mut timing = Timing {
1345            layer_height_um: Some(50),
1346            ..Timing::default()
1347        };
1348        timing.extra.clear();
1349        Meta {
1350            meta_version: Some(1),
1351            timing,
1352            ..Meta::default()
1353        }
1354    }
1355
1356    /// The canonical runs of a mask.
1357    fn runs_of(pixels: &[u8]) -> Vec<Run> {
1358        let mut runs: Vec<Run> = Vec::new();
1359        for &pixel in pixels {
1360            match runs.last_mut() {
1361                Some(last) if last.value == pixel => last.length += 1,
1362                _ => runs.push(Run::new(1, pixel)),
1363            }
1364        }
1365        runs
1366    }
1367
1368    /// One mask per encoding, plus the empty layer: binary stripes, an
1369    /// anti-aliased edge, and a mask neither tag can shrink.
1370    fn masks() -> Vec<Vec<u8>> {
1371        let mut stripes = vec![0u8; PIXELS];
1372        for (index, pixel) in stripes.iter_mut().enumerate() {
1373            if index % 3 == 0 {
1374                *pixel = 255;
1375            }
1376        }
1377        let mut edge = vec![0u8; PIXELS];
1378        edge[PIXELS / 2 - 1] = 90;
1379        edge[PIXELS / 2] = 200;
1380        let gradient: Vec<u8> = (0..PIXELS).map(|index| (index * 8) as u8).collect();
1381        vec![vec![0u8; PIXELS], stripes, edge, gradient]
1382    }
1383
1384    /// Runs pushed with `push_layer_runs` must produce the same file plain
1385    /// pixels do, and an already-encoded layer must store those same bytes.
1386    #[test]
1387    fn the_run_path_writes_the_same_file_as_the_pixel_path() {
1388        let masks = masks();
1389        let layers = masks.len() as u32;
1390
1391        let mut from_masks = Encoder::new(head(layers), meta());
1392        let mut from_runs = Encoder::new(head(layers), meta());
1393        let mut encoded = Encoder::new(head(layers), meta());
1394        for pixels in &masks {
1395            from_masks.push_layer(pixels).expect("a pushable mask");
1396            from_runs
1397                .push_layer_runs(&runs_of(pixels), EncodeMode::Auto)
1398                .expect("a pushable run list");
1399            let record = match ree::encode(pixels, PIXELS as u32, EncodeMode::Auto).unwrap() {
1400                None => EncodedLayer::Empty,
1401                Some((_tag, stream)) => EncodedLayer::single(stream).expect("a tagged stream"),
1402            };
1403            encoded
1404                .push_encoded_layer(record)
1405                .expect("a storable layer");
1406        }
1407
1408        let expected = from_masks.finish().expect("a writable file");
1409        assert_eq!(from_runs.finish().expect("a writable file"), expected);
1410        assert_eq!(encoded.finish().expect("a writable file"), expected);
1411    }
1412
1413    /// The worker threads `finish` uses are a wall-clock knob, not a format one:
1414    /// a print written on a machine asked for two threads is the file it would
1415    /// have been on one asked for eight, `LHAS` and all.
1416    #[test]
1417    fn the_worker_thread_count_does_not_change_the_file() {
1418        let masks = masks();
1419        let layers = masks.len() as u32;
1420
1421        let write = |threads: usize| {
1422            let mut encoder = Encoder::new(head(layers), meta());
1423            // One layer per chunk, so `finish` has one item per layer to hand to
1424            // its workers and the parallel path is the one that runs.
1425            encoder.set_layers_per_chunk(1);
1426            encoder.set_worker_threads(threads);
1427            for pixels in &masks {
1428                encoder.push_layer(pixels).expect("a pushable mask");
1429            }
1430            encoder.finish().expect("a writable file")
1431        };
1432
1433        let serial = write(1);
1434        for threads in [2, 4, 8] {
1435            assert_eq!(write(threads), serial, "with {threads} workers");
1436        }
1437    }
1438
1439    /// A layer's sectors are sorted by id, and a layer with no sector data is the
1440    /// empty-layer form rather than a table entry per sector.
1441    #[test]
1442    fn an_encoded_layer_is_sorted_and_may_be_empty() {
1443        let masks = masks();
1444        let mut sectors = Encoder::new(head(2), meta());
1445        sectors
1446            .push_encoded_layer(EncodedLayer::Sectors(vec![
1447                (
1448                    1,
1449                    ree::encode(&masks[1], PIXELS as u32, EncodeMode::Auto)
1450                        .unwrap()
1451                        .unwrap()
1452                        .1,
1453                ),
1454                (
1455                    0,
1456                    ree::encode(&masks[2], PIXELS as u32, EncodeMode::Auto)
1457                        .unwrap()
1458                        .unwrap()
1459                        .1,
1460                ),
1461            ]))
1462            .expect("a storable layer");
1463        sectors
1464            .push_encoded_layer(EncodedLayer::Sectors(Vec::new()))
1465            .expect("a storable empty layer");
1466
1467        let mut from_masks = Encoder::new(head(2), meta());
1468        from_masks
1469            .push_layer_sectors(&[(0, masks[2].clone()), (1, masks[1].clone())])
1470            .expect("a pushable layer");
1471        from_masks
1472            .push_layer_sectors(&[])
1473            .expect("a pushable empty layer");
1474
1475        assert_eq!(
1476            sectors.finish().expect("a writable file"),
1477            from_masks.finish().expect("a writable file")
1478        );
1479    }
1480
1481    /// The META minima section 11.2 requires, so a test can validate a file it
1482    /// wrote rather than only read its structure back.
1483    fn conforming_meta() -> Meta {
1484        Meta {
1485            meta_version: Some(1),
1486            timing: Timing {
1487                normal_exposure_ms: Some(2500),
1488                bottom_exposure_ms: Some(30000),
1489                bottom_layer_count: Some(2),
1490                transition_layer_count: Some(2),
1491                layer_height_um: Some(50),
1492                lift_slow_distance_um: Some(5000),
1493                lift_slow_speed_um_min: Some(65000),
1494                retract_fast_distance_um: Some(5000),
1495                retract_fast_speed_um_min: Some(150000),
1496                ..Timing::default()
1497            },
1498            ..Meta::default()
1499        }
1500    }
1501
1502    /// Overrides whose deltas are equal share one `LROV` chunk, so a delta set on a
1503    /// range costs one chunk while a delta that differs by one field is its own; and
1504    /// each pair still resolves from the chunk its own entry names (section 4.5).
1505    #[test]
1506    fn equal_override_deltas_share_one_chunk() {
1507        use crate::reader::LumenFile;
1508        use crate::validate::Level;
1509
1510        let exposure = |ms: u32| Timing {
1511            normal_exposure_ms: Some(ms),
1512            ..Timing::default()
1513        };
1514        let mut encoder = Encoder::new(head(4), conforming_meta());
1515        for _ in 0..4 {
1516            encoder.push_layer(&[0u8; PIXELS]).expect("a layer");
1517        }
1518        encoder
1519            .set_overrides(vec![
1520                Override {
1521                    layer: 0,
1522                    sector_id: 0,
1523                    timing: exposure(2800),
1524                },
1525                Override {
1526                    layer: 1,
1527                    sector_id: 0,
1528                    timing: exposure(2800),
1529                },
1530                Override {
1531                    layer: 2,
1532                    sector_id: 0,
1533                    timing: exposure(2800),
1534                },
1535                Override {
1536                    layer: 3,
1537                    sector_id: 0,
1538                    timing: exposure(3000),
1539                },
1540            ])
1541            .expect("one override set per (layer, sector)");
1542
1543        let bytes = encoder.finish().expect("a writable file");
1544        let file =
1545            LumenFile::open(&bytes, Level::Strict).expect("a shared chunk is a conforming file");
1546        let first_lrov = |layer: u32| {
1547            file.layer_table()
1548                .entry(layer, 0)
1549                .expect("an entry for every layer")
1550                .first_lrov
1551        };
1552        assert_eq!(first_lrov(0), first_lrov(1), "equal deltas share a chunk");
1553        assert_eq!(first_lrov(1), first_lrov(2), "equal deltas share a chunk");
1554        assert_ne!(
1555            first_lrov(2),
1556            first_lrov(3),
1557            "a delta that differs is its own chunk"
1558        );
1559        let chunks = file
1560            .directory()
1561            .entries()
1562            .filter(|d| d.chunk_type == crate::container::ChunkType::LROV)
1563            .count();
1564        assert_eq!(chunks, 2, "four overrides with two distinct deltas");
1565
1566        // The sharing is an encoding detail: every pair still resolves from the
1567        // chunk its entry names.
1568        for layer in 0..3 {
1569            assert_eq!(
1570                file.timing_for(layer, 0).expect("resolvable").exposure_ms,
1571                2800
1572            );
1573        }
1574        assert_eq!(file.timing_for(3, 0).expect("resolvable").exposure_ms, 3000);
1575    }
1576
1577    /// The dictionary's training samples come from across the print rather than from
1578    /// its first chunks: the start of a print is its raft and its bottom layers, and a
1579    /// dictionary trained on those describes the rest of the file badly.
1580    #[test]
1581    fn dictionary_samples_span_the_print() {
1582        // Sixteen chunks, each filled with its own byte value, so every sample says by
1583        // its content which chunk it came from. Each chunk is four windows long.
1584        let plaintexts: Vec<Vec<u8>> = (1u8..=16).map(|value| vec![value; 4096]).collect();
1585
1586        let samples = dictionary_samples(&plaintexts);
1587
1588        let chunks: std::collections::BTreeSet<u8> = samples.iter().map(|s| s[0]).collect();
1589        assert_eq!(chunks.len(), 16, "every chunk of the print contributes");
1590        assert!(
1591            samples.len() <= DICTIONARY_SAMPLE_WINDOWS,
1592            "the budget holds"
1593        );
1594        assert!(
1595            samples.iter().all(|sample| sample.len() == 1024),
1596            "a sample is one window"
1597        );
1598        assert_eq!(
1599            samples,
1600            dictionary_samples(&plaintexts),
1601            "sampling is deterministic"
1602        );
1603    }
1604
1605    /// A print a dictionary cannot help is written without one: high-entropy
1606    /// greyscale, which is what an anti-aliased print's planes look like.
1607    #[test]
1608    fn a_dictionary_that_does_not_pay_is_not_written() {
1609        let layers = 8u32;
1610        let mut head = head(layers);
1611        head.display_width_px = 1024;
1612        head.display_height_px = 1024;
1613        let pixels = (head.display_width_px * head.display_height_px) as usize;
1614
1615        let mut state = 0x2545_F491_4F6C_DD1Du64;
1616        let mut random = || {
1617            state ^= state << 13;
1618            state ^= state >> 7;
1619            state ^= state << 17;
1620            state >> 24
1621        };
1622        let mut encoder = Encoder::new(head, conforming_meta());
1623        encoder.set_dictionary(true);
1624        for _ in 0..layers {
1625            let mask: Vec<u8> = (0..pixels).map(|_| random() as u8).collect();
1626            encoder.push_layer(&mask).expect("a layer");
1627        }
1628
1629        let bytes = encoder.finish().expect("a writable file");
1630        let file = crate::reader::LumenFile::open(&bytes, crate::validate::Level::Strict)
1631            .expect("a conforming file");
1632        assert!(
1633            file.dictionary().is_none(),
1634            "a dictionary that makes the frames larger is not written"
1635        );
1636    }
1637
1638    /// A stream that carries no tag is not a layer, and the encoder says so
1639    /// before it reaches the layer table.
1640    /// A print whose layers barely move: a white band per row whose left edge
1641    /// advances one pixel per layer, so every layer's run lengths are close to its
1642    /// neighbours' without being equal. This is the shape the attached encoding is
1643    /// for, and the shape a byte-counting encoder would get wrong the other way -
1644    /// on a print this smooth the plain form's bytes are the ones that compress.
1645    fn moving_band_layers(count: u32) -> Vec<Vec<u8>> {
1646        const SIDE: usize = 64;
1647        let mut layers = Vec::new();
1648        for index in 0..count as usize {
1649            let mut pixels = vec![0u8; SIDE * SIDE];
1650            for (row, line) in pixels.chunks_mut(SIDE).enumerate() {
1651                // A lens that walks one pixel to the right per layer: its chord
1652                // grows and shrinks smoothly down the rows, and its boundaries
1653                // move between layers without the run lengths repeating.
1654                let half = if row < 32 { row / 2 } else { (63 - row) / 2 };
1655                let centre = 24 + index;
1656                let start = centre.saturating_sub(half);
1657                let end = (centre + half + 1).min(SIDE);
1658                if start < end {
1659                    line[start..end].fill(255);
1660                }
1661            }
1662            layers.push(pixels);
1663        }
1664        layers
1665    }
1666
1667    /// The band fixture's pixels per layer.
1668    const PIXELS_BAND: u32 = 64 * 64;
1669
1670    fn band_head(layers: u32) -> Head {
1671        Head {
1672            display_width_px: 64,
1673            display_height_px: 64,
1674            ..head(layers)
1675        }
1676    }
1677
1678    fn written_with_probe(probe: bool) -> (Vec<u8>, Vec<Vec<u8>>) {
1679        let layers = moving_band_layers(8);
1680        let mut encoder = Encoder::new(band_head(layers.len() as u32), conforming_meta());
1681        encoder.set_tag_probe(probe);
1682        for layer in &layers {
1683            encoder.push_layer(layer).expect("a pushable layer");
1684        }
1685        (encoder.finish().expect("a writable file"), layers)
1686    }
1687
1688    #[test]
1689    fn the_tag_probe_keeps_the_smaller_encoding() {
1690        let (probed, layers) = written_with_probe(true);
1691        let (plain, _) = written_with_probe(false);
1692
1693        // The probe changes bytes, never content: both files decode to the layers
1694        // that were pushed.
1695        for bytes in [&probed, &plain] {
1696            let file = crate::LumenFile::open(bytes, crate::validate::Level::Strict)
1697                .expect("a conforming file");
1698            for (index, mask) in layers.iter().enumerate() {
1699                assert_eq!(&file.layer(index as u32).unwrap().pixels, mask);
1700            }
1701        }
1702
1703        // Whichever way it decided, it did not write a larger file, and on this
1704        // print it decided the other way from the plain encoding.
1705        assert!(
1706            probed.len() <= plain.len(),
1707            "the probe must not write a larger file: {} against {}",
1708            probed.len(),
1709            plain.len()
1710        );
1711        assert!(
1712            probed.len() < plain.len(),
1713            "a print whose layers barely move is where the attached form wins: \
1714             {} against {}",
1715            probed.len(),
1716            plain.len()
1717        );
1718    }
1719
1720    #[test]
1721    fn a_pre_encoded_layer_can_carry_its_alternative() {
1722        // The parallel path: a caller encodes a layer on its own thread and hands
1723        // over both encodings, so the probe still has something to weigh when the
1724        // group closes. Both files decode to the same layers, and the pair is never
1725        // larger than the plain one.
1726        let layers = moving_band_layers(8);
1727        let mut paired = Encoder::new(band_head(layers.len() as u32), conforming_meta());
1728        let mut plain = Encoder::new(band_head(layers.len() as u32), conforming_meta());
1729        for layer in &layers {
1730            let (tag, primary) = ree::encode(layer, PIXELS_BAND, EncodeMode::Auto)
1731                .unwrap()
1732                .expect("a non-empty layer");
1733            assert_eq!(tag, ree::TAG_BINARY);
1734            let (other_tag, alternative) = ree::encode(layer, PIXELS_BAND, EncodeMode::Attached)
1735                .unwrap()
1736                .expect("a non-empty layer");
1737            assert_eq!(other_tag, ree::TAG_ATTACHED);
1738            paired
1739                .push_encoded_layer_pair(
1740                    EncodedLayer::single(primary).unwrap(),
1741                    Some(EncodedLayer::single(alternative).unwrap()),
1742                )
1743                .expect("a pair of encodings of one layer");
1744            plain
1745                .push_encoded_layer_pair(
1746                    EncodedLayer::single(
1747                        ree::encode(layer, PIXELS_BAND, EncodeMode::Auto)
1748                            .unwrap()
1749                            .expect("a non-empty layer")
1750                            .1,
1751                    )
1752                    .unwrap(),
1753                    None,
1754                )
1755                .expect("one encoding");
1756        }
1757        let paired = paired.finish().expect("a writable file");
1758        let plain = plain.finish().expect("a writable file");
1759        for bytes in [&paired, &plain] {
1760            let file = crate::LumenFile::open(bytes, crate::validate::Level::Strict)
1761                .expect("a conforming file");
1762            for (index, mask) in layers.iter().enumerate() {
1763                assert_eq!(&file.layer(index as u32).unwrap().pixels, mask);
1764            }
1765        }
1766        assert!(
1767            paired.len() < plain.len(),
1768            "the pair decides what the probe would have decided: {} against {}",
1769            paired.len(),
1770            plain.len()
1771        );
1772    }
1773
1774    #[test]
1775    fn the_tag_probe_can_be_turned_off() {
1776        let (probed, _) = written_with_probe(true);
1777        let (plain, _) = written_with_probe(false);
1778        assert_eq!(
1779            plain.len(),
1780            written_with_probe(false).0.len(),
1781            "the same input writes the same file"
1782        );
1783        assert_ne!(probed.len(), plain.len());
1784    }
1785
1786    #[test]
1787    fn an_encoded_layer_is_checked_before_it_is_stored() {
1788        let mut encoder = Encoder::new(head(1), meta());
1789        for empty_or_unknown in [Vec::new(), vec![0x04, 0x00], vec![0xFF]] {
1790            assert_eq!(
1791                EncodedLayer::single(empty_or_unknown.clone())
1792                    .unwrap_err()
1793                    .check(),
1794                Check::ReeTag
1795            );
1796            assert_eq!(
1797                encoder
1798                    .push_encoded_layer(EncodedLayer::Sectors(vec![(0, empty_or_unknown)]))
1799                    .unwrap_err()
1800                    .check(),
1801                Check::ReeTag
1802            );
1803        }
1804
1805        // One sector id twice on a layer is not a sector list.
1806        let stream = vec![ree::TAG_GRAYSCALE, 0x01, 0x00, PIXELS as u8];
1807        assert_eq!(
1808            encoder
1809                .push_encoded_layer(EncodedLayer::Sectors(vec![
1810                    (2, stream.clone()),
1811                    (2, stream)
1812                ]))
1813                .unwrap_err()
1814                .check(),
1815            Check::LtblSectorIdUnique
1816        );
1817
1818        // Nothing rejected above was pushed.
1819        assert!(encoder.layers.is_empty());
1820    }
1821}