use crate::check::Check;
use crate::chunks::extd::Extension;
use crate::chunks::head::Head;
use crate::chunks::lhas::{self, LayerHashes};
use crate::chunks::ltbl::{LayerEntry, LayerTable};
use crate::chunks::preview::PreviewRole;
use crate::chunks::zdic::ZstdDictionary;
use crate::chunks::{self, json_chunks, layr};
use crate::container::{
self, ChunkDescriptor, ChunkType, Directory, FileHeader, CHUNK_FLAG_ENCRYPTED, FLAG_ENCRYPTED,
FLAG_MULTI_SECTOR,
};
use crate::crypto::{self, Auth, Cipher, RecipientEntry, SessionKey};
use crate::error::{Error, Result};
use crate::io::Writer;
use crate::json::{Meta, Profile, Sector, Timing};
use crate::ree::{self, EncodeMode, Run};
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Argon2Params {
pub iterations: u32,
pub memory_kib: u32,
pub parallelism: u8,
}
impl Default for Argon2Params {
fn default() -> Self {
Argon2Params {
iterations: 3,
memory_kib: 65_536,
parallelism: 1,
}
}
}
#[derive(Debug, Clone)]
pub struct EncryptOptions {
pub cipher: Cipher,
pub password: Option<String>,
pub recipients: Vec<[u8; 32]>,
pub argon2: Argon2Params,
}
impl EncryptOptions {
pub fn password(password: impl Into<String>) -> EncryptOptions {
EncryptOptions {
cipher: Cipher::Aes256Gcm,
password: Some(password.into()),
recipients: Vec::new(),
argon2: Argon2Params::default(),
}
}
}
pub const DEFAULT_LAYERS_PER_CHUNK: u32 = 64;
pub const DEFAULT_ZSTD_LEVEL: i32 = 6;
pub const JSON_ZSTD_LEVEL: i32 = 3;
pub const DICTIONARY_SAMPLE_WINDOWS: usize = 256;
const DICTIONARY_MAX_BYTES: usize = 112_640;
const DICTIONARY_PROBE_BYTES: usize = 4 * 1024 * 1024;
const DICTIONARY_PROBE_MIN_PLAINTEXT: usize = 8 * 1024 * 1024;
const DICTIONARY_PROBE_MARGIN: f64 = 0.002;
#[derive(Debug, Clone, PartialEq)]
pub struct Override {
pub layer: u32,
pub sector_id: u32,
pub timing: Timing,
}
#[derive(Debug, Clone)]
pub enum EncodedLayer {
Empty,
Sectors(Vec<(u32, Vec<u8>)>),
}
impl EncodedLayer {
pub fn single(stream: Vec<u8>) -> Result<Self> {
check_stream(&stream)?;
Ok(EncodedLayer::Sectors(vec![(0, stream)]))
}
fn stream(&self, sector_id: u32) -> Option<&[u8]> {
match self {
EncodedLayer::Empty => None,
EncodedLayer::Sectors(list) => list
.iter()
.find(|(id, _)| *id == sector_id)
.map(|(_, stream)| stream.as_slice()),
}
}
}
fn single_sector(encoded: Option<(u8, Vec<u8>)>) -> EncodedLayer {
match encoded {
None => EncodedLayer::Empty,
Some((_tag, mask)) => EncodedLayer::Sectors(vec![(0, mask)]),
}
}
fn check_stream(stream: &[u8]) -> Result<()> {
match stream.first() {
None => Err(Error::new(
Check::ReeTag,
"layer mask data is empty: the empty-layer form carries no bytes",
)),
Some(&(ree::TAG_BINARY | ree::TAG_GRAYSCALE | ree::TAG_SPLIT | ree::TAG_ATTACHED)) => {
Ok(())
}
Some(&tag) => Err(Error::new(
Check::ReeTag,
format!("unknown layer encoding tag 0x{tag:02X}"),
)),
}
}
struct ChunkCandidate {
plaintext: Vec<u8>,
placement: Vec<(u32, u64, u32)>,
}
#[derive(Debug)]
struct Pending {
chunk_type: ChunkType,
stored: Vec<u8>,
size_uncompressed: u64,
size_compressed: u64,
flags: u32,
}
#[derive(Debug)]
pub struct Encoder {
head: Head,
meta: Meta,
profile: Option<Profile>,
overrides: Vec<Override>,
previews: Vec<(PreviewRole, Vec<u8>)>,
extensions: Vec<Extension>,
voxl: Option<Vec<u8>>,
layers_per_chunk: u32,
zstd_level: i32,
dictionary: bool,
layer_hashes: bool,
worker_threads: usize,
encryption: Option<EncryptOptions>,
session: Option<SessionKey>,
tag_probe: bool,
layers: Vec<(EncodedLayer, Option<EncodedLayer>)>,
}
impl Encoder {
pub fn new(head: Head, meta: Meta) -> Encoder {
Encoder {
head,
meta,
profile: None,
overrides: Vec::new(),
previews: Vec::new(),
extensions: Vec::new(),
voxl: None,
layers_per_chunk: DEFAULT_LAYERS_PER_CHUNK,
zstd_level: DEFAULT_ZSTD_LEVEL,
dictionary: true,
layer_hashes: true,
worker_threads: 0,
encryption: None,
session: None,
tag_probe: true,
layers: Vec::new(),
}
}
pub fn set_profile(&mut self, profile: Profile) {
self.profile = Some(profile);
}
pub fn set_sectors(&mut self, sectors: Vec<Sector>) {
self.meta.sectors = Some(sectors);
}
pub fn set_overrides(&mut self, overrides: Vec<Override>) -> Result<()> {
let mut overrides = overrides;
overrides.sort_by_key(|over| (over.layer, over.sector_id));
for pair in overrides.windows(2) {
if pair[0].layer == pair[1].layer && pair[0].sector_id == pair[1].sector_id {
return Err(Error::new(
Check::LrovOrphan,
format!(
"layer {} sector {} carries two override sets; a (layer, sector) has one \
or none",
pair[0].layer, pair[0].sector_id
),
));
}
}
self.overrides = overrides;
Ok(())
}
pub fn add_preview(&mut self, role: PreviewRole, png: Vec<u8>) {
self.previews.push((role, png));
}
pub fn add_extension(&mut self, extension: Extension) {
self.extensions.push(extension);
}
pub fn set_voxl(&mut self, voxl: Vec<u8>) {
self.voxl = Some(voxl);
}
pub fn set_layers_per_chunk(&mut self, layers: u32) {
self.layers_per_chunk = layers.max(1);
}
pub fn set_tag_probe(&mut self, enabled: bool) {
self.tag_probe = enabled;
}
pub fn set_zstd_level(&mut self, level: i32) {
self.zstd_level = level;
}
pub fn set_dictionary(&mut self, enabled: bool) {
self.dictionary = enabled;
}
pub fn set_layer_hashes(&mut self, enabled: bool) {
self.layer_hashes = enabled;
}
pub fn set_worker_threads(&mut self, threads: usize) {
self.worker_threads = threads;
}
fn workers(&self) -> usize {
match self.worker_threads {
0 => available_parallelism(),
threads => threads.min(available_parallelism()),
}
}
pub fn set_encryption(&mut self, options: EncryptOptions) {
self.session = Some(SessionKey::generate());
self.encryption = Some(options);
}
pub fn session_key(&self) -> Option<SessionKey> {
self.session
}
pub fn push_layer(&mut self, pixels: &[u8]) -> Result<()> {
self.push_layer_with_mode(pixels, EncodeMode::Auto)
}
pub fn push_layer_with_mode(&mut self, pixels: &[u8], mode: EncodeMode) -> Result<()> {
let total_pixels = self.total_pixels();
check_mask(pixels, total_pixels, "the layer")?;
let record = single_sector(ree::encode(pixels, total_pixels, mode)?);
let alternative = match (mode, self.tag_probe) {
(EncodeMode::Auto, true) => Some(single_sector(ree::encode(
pixels,
total_pixels,
EncodeMode::Attached,
)?)),
_ => None,
};
self.layers.push((record, alternative));
Ok(())
}
pub fn push_layer_sectors(&mut self, sectors: &[(u32, Vec<u8>)]) -> Result<()> {
let total_pixels = self.total_pixels();
let mut masks: Vec<(u32, Vec<u8>)> = Vec::new();
let mut alternatives: Vec<(u32, Vec<u8>)> = Vec::new();
for (sector_id, pixels) in sectors {
check_mask(
pixels,
total_pixels,
&format!("sector {sector_id} of the layer"),
)?;
if sectors.iter().filter(|(id, _)| id == sector_id).count() > 1 {
return Err(Error::new(
Check::LtblSectorIdUnique,
format!("sector {sector_id} appears twice in one layer"),
));
}
if let Some((_tag, mask)) = ree::encode(pixels, total_pixels, EncodeMode::Auto)? {
masks.push((*sector_id, mask));
if self.tag_probe {
if let Some((_tag, other)) =
ree::encode(pixels, total_pixels, EncodeMode::Attached)?
{
alternatives.push((*sector_id, other));
}
}
}
}
if masks.is_empty() {
self.layers.push((EncodedLayer::Empty, None));
} else {
masks.sort_by_key(|(sector_id, _)| *sector_id);
alternatives.sort_by_key(|(sector_id, _)| *sector_id);
let alternative =
(!alternatives.is_empty()).then_some(EncodedLayer::Sectors(alternatives));
self.layers
.push((EncodedLayer::Sectors(masks), alternative));
}
Ok(())
}
pub fn push_layer_runs(&mut self, runs: &[Run], mode: EncodeMode) -> Result<()> {
let total_pixels = self.total_pixels();
let record = single_sector(ree::encode_runs(runs, total_pixels, mode)?);
let alternative = match (mode, self.tag_probe) {
(EncodeMode::Auto, true) => Some(single_sector(ree::encode_runs(
runs,
total_pixels,
EncodeMode::Attached,
)?)),
_ => None,
};
self.layers.push((record, alternative));
Ok(())
}
pub fn push_encoded_layer(&mut self, layer: EncodedLayer) -> Result<()> {
self.push_encoded_layer_pair(layer, None)
}
pub fn push_encoded_layer_pair(
&mut self,
layer: EncodedLayer,
alternative: Option<EncodedLayer>,
) -> Result<()> {
if let Some(other) = alternative.as_ref() {
match other {
EncodedLayer::Empty => {}
EncodedLayer::Sectors(list) => {
for (_, stream) in list {
check_stream(stream)?;
}
}
}
}
let record = match layer {
EncodedLayer::Empty => EncodedLayer::Empty,
EncodedLayer::Sectors(mut list) => {
for (_, stream) in &list {
check_stream(stream)?;
}
list.sort_by_key(|(sector_id, _)| *sector_id);
for pair in list.windows(2) {
if pair[0].0 == pair[1].0 {
return Err(Error::new(
Check::LtblSectorIdUnique,
format!("sector {} appears twice in one layer", pair[0].0),
));
}
}
if list.is_empty() {
EncodedLayer::Empty
} else {
EncodedLayer::Sectors(list)
}
}
};
self.layers.push((record, alternative));
Ok(())
}
fn total_pixels(&self) -> u32 {
self.head.total_pixels()
}
pub fn finish(self) -> Result<Vec<u8>> {
let pushed = self.layers.len() as u32;
if pushed != self.head.total_layers {
return Err(Error::new(
Check::HeadTotalLayers,
format!(
"HEAD declares {} layers but {pushed} were pushed",
self.head.total_layers
),
));
}
if let Some(sectors) = self.meta.sectors.as_deref() {
let mut seen: Vec<u32> = Vec::with_capacity(sectors.len());
for (i, sector) in sectors.iter().enumerate() {
if sector.sector_id == 0 {
return Err(Error::new(
Check::MetaSectorsShape,
format!(
"META.sectors[{i}] uses sector_id 0, which names the implicit primary \
sector"
),
));
}
if seen.contains(§or.sector_id) {
return Err(Error::new(
Check::MetaSectorsShape,
format!("sector_id {} appears twice", sector.sector_id),
));
}
seen.push(sector.sector_id);
}
}
let sector_sets = self.sector_sets()?;
let group_size = self.layers_per_chunk.max(1) as usize;
let mut chunk_plaintexts: Vec<Vec<u8>> = Vec::new();
let mut chunk_placements: Vec<Vec<(u32, u64, u32)>> = Vec::new();
let mut chunk_alternatives: Vec<Option<ChunkCandidate>> = Vec::new();
let mut chunk_sectors: Vec<u32> = Vec::new();
for group_start in (0..pushed as usize).step_by(group_size) {
let group_end = (group_start + group_size).min(pushed as usize);
let mut sectors: Vec<u32> = sector_sets[group_start..group_end]
.iter()
.flatten()
.copied()
.collect();
sectors.sort_unstable();
sectors.dedup();
let probed = self.layers[group_start..group_end]
.iter()
.any(|(_, alternative)| alternative.is_some());
for sector in sectors {
let mut plaintext = Vec::new();
let mut placed: Vec<(u32, u64, u32)> = Vec::new();
let mut alternative_text = Vec::new();
let mut alternative_placed: Vec<(u32, u64, u32)> = Vec::new();
for (index, (record, alternative)) in self
.layers
.iter()
.enumerate()
.take(group_end)
.skip(group_start)
{
if let Some(mask) = record.stream(sector) {
placed.push((index as u32, plaintext.len() as u64, mask.len() as u32));
plaintext.extend_from_slice(mask);
}
let chosen = alternative.as_ref().unwrap_or(record);
if let Some(mask) = chosen.stream(sector) {
alternative_placed.push((
index as u32,
alternative_text.len() as u64,
mask.len() as u32,
));
alternative_text.extend_from_slice(mask);
}
}
if plaintext.is_empty() {
continue;
}
chunk_plaintexts.push(plaintext);
chunk_placements.push(placed);
chunk_alternatives.push(
(probed && alternative_text != chunk_plaintexts[chunk_plaintexts.len() - 1])
.then_some(ChunkCandidate {
plaintext: alternative_text,
placement: alternative_placed,
}),
);
chunk_sectors.push(sector);
}
}
if chunk_plaintexts.is_empty() {
chunk_plaintexts.push(Vec::new());
chunk_placements.push(Vec::new());
chunk_alternatives.push(None);
chunk_sectors.push(0);
}
let dictionary = if self.dictionary {
let samples = dictionary_samples(&chunk_plaintexts);
match chunks::train_dictionary(&samples, DICTIONARY_MAX_BYTES) {
Ok(bytes) if !bytes.is_empty() && chunks::dictionary_id(&bytes) != 0 => {
if dictionary_pays(&chunk_plaintexts, &bytes, self.zstd_level)? {
Some(ZstdDictionary {
zdic_version: 1,
dict_id: chunks::dictionary_id(&bytes),
dict_bytes: bytes,
})
} else {
None
}
}
_ => None,
}
} else {
None
};
let dict_bytes = dictionary.as_ref().map(|d| d.dict_bytes.as_slice());
let workers = self.workers();
let level = self.zstd_level;
let alternatives = &chunk_alternatives;
let choices: Vec<(Vec<u8>, bool)> =
map_ordered(&chunk_plaintexts, workers, |index, plaintext| {
let frame = |text: &[u8]| -> Result<Vec<u8>> {
let frame = chunks::compress(text, level, dict_bytes)?;
chunks::frame_content_size(&frame)?;
Ok(frame)
};
let primary = frame(plaintext)?;
match &alternatives[index] {
Some(candidate) => match frame(&candidate.plaintext) {
Ok(other) if other.len() < primary.len() => Ok((other, true)),
Ok(_) | Err(_) => Ok((primary, false)),
},
None => Ok((primary, false)),
}
})?;
let mut placement: HashMap<(u32, u32), (u32, u64, u32)> = HashMap::new();
let mut frames: Vec<Vec<u8>> = Vec::with_capacity(choices.len());
for (index, (frame, chose_alternative)) in choices.into_iter().enumerate() {
let sector = chunk_sectors[index];
let entries = if chose_alternative {
let candidate = chunk_alternatives[index]
.take()
.expect("an alternative that was compressed");
chunk_plaintexts[index] = candidate.plaintext;
candidate.placement
} else {
std::mem::take(&mut chunk_placements[index])
};
for (layer, offset, size) in entries {
placement.insert((layer, sector), (index as u32, offset, size));
}
frames.push(frame);
}
let hashes = if self.layer_hashes {
let leaves = map_ordered(§or_sets, workers, |index, sectors| {
let mut data = Vec::new();
for sector in sectors {
let key = (index as u32, *sector);
if let Some(&(ordinal, offset, size)) = placement.get(&key) {
let start = offset as usize;
data.extend_from_slice(
&chunk_plaintexts[ordinal as usize][start..start + size as usize],
);
}
}
Ok(lhas::leaf_hash(&data))
})?;
Some(LayerHashes {
hash_algorithm: lhas::HASH_ALGORITHM_SHA256,
hash_size: lhas::HASH_SIZE_SHA256,
layer_count: leaves.len() as u32,
merkle_root: lhas::merkle_root(&leaves),
layer_hashes: leaves,
})
} else {
None
};
let mut flags = 0u32;
if self
.layers
.iter()
.any(|(record, _)| matches!(record, EncodedLayer::Sectors(list) if list.len() > 1))
{
flags |= FLAG_MULTI_SECTOR;
}
let sealed = self.session.is_some() && self.encryption.is_some();
if self.encryption.is_some() {
flags |= FLAG_ENCRYPTED;
}
let mut pending = Vec::new();
pending.push(Pending::raw(ChunkType::HEAD, self.head.to_bytes()));
pending.push(self.seal_if_needed(
ChunkType::META,
json_chunks::meta_to_bytes(&self.meta)?,
true,
)?);
if let Some(profile) = self.profile.as_ref() {
pending.push(self.seal_if_needed(
ChunkType::PROF,
json_chunks::profile_to_bytes(profile)?,
true,
)?);
}
if let Some(auth) = self.build_auth()? {
pending.push(Pending::raw(ChunkType::AUTH, crypto::encode_auth(&auth)));
}
let mut lrov_index: HashMap<(u32, u32), u32> = HashMap::new();
let mut shared: HashMap<Vec<u8>, u32> = HashMap::new();
for over in &self.overrides {
let payload = json_chunks::lrov_to_bytes(&over.timing)?;
let index = match shared.get(&payload) {
Some(index) => *index,
None => {
let index = pending.len() as u32;
pending.push(self.seal_if_needed(ChunkType::LROV, payload.clone(), true)?);
shared.insert(payload, index);
index
}
};
lrov_index.insert((over.layer, over.sector_id), index);
}
if let Some(dictionary) = dictionary.as_ref() {
pending.push(self.seal_if_needed(ChunkType::ZDIC, dictionary.to_bytes(), false)?);
}
for (role, png) in &self.previews {
pending.push(Pending::raw_with_flags(
ChunkType::PREV,
png.clone(),
role.to_flags(),
));
}
if let Some(voxl) = self.voxl.as_ref() {
pending.push(self.seal_if_needed(ChunkType::VOXL, voxl.clone(), true)?);
}
let layr_base = pending.len() as u32 + 1 + u32::from(hashes.is_some());
let entries =
self.table_entries(pushed, §or_sets, &placement, layr_base, &lrov_index)?;
let table = LayerTable::new(entries, pushed)?;
pending.push(Pending::raw(ChunkType::LTBL, table.to_bytes()));
if let Some(hashes) = hashes.as_ref() {
pending.push(Pending::raw(ChunkType::LHAS, hashes.to_bytes()));
}
debug_assert_eq!(
pending.len() as u32,
layr_base,
"the LAYR chunks start where their indices say they do"
);
for (ordinal, frame) in frames.iter().enumerate() {
let directory_index = layr_base + ordinal as u32;
let container = if sealed {
let sealed_frame = crypto::seal(
self.cipher()?,
self.session.as_ref().expect("set with encryption"),
ChunkType::LAYR,
directory_index,
frame,
)?;
layr::to_bytes(&sealed_frame)
} else {
layr::to_bytes(frame)
};
pending.push(Pending {
chunk_type: ChunkType::LAYR,
stored: container,
size_uncompressed: layr::container_len(frame.len()),
size_compressed: if sealed {
layr::container_len(frame.len()) + crypto::UNIT_OVERHEAD as u64
} else {
0
},
flags: if sealed { CHUNK_FLAG_ENCRYPTED } else { 0 },
});
}
for extension in &self.extensions {
pending.push(Pending::raw_with_flags(
ChunkType::EXTD,
extension.to_bytes(),
extension.to_flags(),
));
}
let mut out = Writer::with_capacity(
pending.iter().map(|p| p.stored.len() + 8).sum::<usize>() + 64 + pending.len() * 32,
);
out.bytes(&[0u8; 32]);
let mut directory = Directory::new();
let mut total_uncompressed = 0u64;
for chunk in &pending {
out.pad_to(8);
let offset = out.len() as u64;
out.bytes(&chunk.stored);
total_uncompressed += chunk.size_uncompressed;
directory.push(ChunkDescriptor {
chunk_type: chunk.chunk_type,
offset,
size_uncompressed: chunk.size_uncompressed,
size_compressed: chunk.size_compressed,
flags: chunk.flags,
});
}
out.pad_to(8);
let dir_offset = out.len() as u64;
out.bytes(&directory.to_bytes());
let header = FileHeader {
version: 1,
dir_offset,
chunk_count: directory.descriptors.len() as u32,
flags,
total_uncompressed_size: total_uncompressed,
};
out.patch_bytes(0, &header.to_bytes());
let crc = container::crc32c(out.as_slice());
out.bytes(&container::TRAILER_MAGIC);
out.u32(crc);
Ok(out.into_vec())
}
fn sector_sets(&self) -> Result<Vec<Vec<u32>>> {
let mut sets: Vec<Vec<u32>> = self
.layers
.iter()
.map(|(record, _)| match record {
EncodedLayer::Empty => Vec::new(),
EncodedLayer::Sectors(list) => list.iter().map(|(id, _)| *id).collect(),
})
.collect();
for over in &self.overrides {
let index = over.layer as usize;
let Some(ids) = sets.get_mut(index) else {
return Err(Error::new(
Check::LtblLayerIndexRange,
format!(
"an override names layer {} of {}",
over.layer,
self.layers.len()
),
));
};
if !ids.contains(&over.sector_id) {
ids.push(over.sector_id);
}
}
for ids in &mut sets {
if !ids.contains(&0) {
ids.push(0);
}
ids.sort_unstable();
ids.dedup();
}
Ok(sets)
}
fn table_entries(
&self,
pushed: u32,
sector_sets: &[Vec<u32>],
placement: &HashMap<(u32, u32), (u32, u64, u32)>,
layr_base: u32,
lrov_index: &HashMap<(u32, u32), u32>,
) -> Result<Vec<LayerEntry>> {
let mut entries = Vec::new();
for index in 0..pushed {
let ids = §or_sets[index as usize];
let additional = (ids.len() - 1) as u32;
let home = ids
.iter()
.find_map(|sector| placement.get(&(index, *sector)))
.map_or(layr_base, |(ordinal, _, _)| layr_base + *ordinal);
for (position, sector) in ids.iter().enumerate() {
let run = placement.get(&(index, *sector));
entries.push(LayerEntry {
data_size: run.map_or(0, |(_, _, size)| *size),
first_lrov: lrov_index.get(&(index, *sector)).copied().unwrap_or(0),
first_layr: run.map_or(home, |(ordinal, _, _)| layr_base + *ordinal),
additional_sector_count: if position == 0 { additional } else { 0 },
data_offset: run.map_or(0, |(_, offset, _)| *offset),
sector_id: *sector,
});
}
}
Ok(entries)
}
fn cipher(&self) -> Result<Cipher> {
Ok(self
.encryption
.as_ref()
.expect("cipher asked for without encryption")
.cipher)
}
fn seal_if_needed(
&self,
chunk_type: ChunkType,
plaintext: Vec<u8>,
compress: bool,
) -> Result<Pending> {
let payload = if compress {
chunks::compress(&plaintext, JSON_ZSTD_LEVEL, None)?
} else {
plaintext.clone()
};
let compressed_len = if compress { payload.len() as u64 } else { 0 };
match self.session.as_ref() {
None => Ok(Pending {
chunk_type,
stored: payload,
size_uncompressed: plaintext.len() as u64,
size_compressed: compressed_len,
flags: 0,
}),
Some(key) => {
let sealed = crypto::seal(self.cipher()?, key, chunk_type, 0, &payload)?;
Ok(Pending {
chunk_type,
size_uncompressed: plaintext.len() as u64,
stored: sealed,
size_compressed: payload.len() as u64 + crypto::UNIT_OVERHEAD as u64,
flags: CHUNK_FLAG_ENCRYPTED,
})
}
}
}
fn build_auth(&self) -> Result<Option<Auth>> {
let Some(options) = self.encryption.as_ref() else {
return Ok(None);
};
let session = self.session.as_ref().expect("set with encryption");
let mut mode = 0u32;
let mut password = None;
let mut recipients: Vec<RecipientEntry> = Vec::new();
if let Some(secret) = options.password.as_deref() {
mode |= crypto::MODE_PASSWORD;
let params = options.argon2;
password = Some(crypto::wrap_password(
session,
secret,
params.iterations,
params.memory_kib,
params.parallelism,
)?);
}
for public in &options.recipients {
mode |= crypto::MODE_MACHINE;
recipients.push(crypto::wrap_machine(session, public)?);
}
if mode == 0 {
return Err(Error::new(
Check::CryptModeEmpty,
"encryption was requested with neither a password nor a recipient",
));
}
Ok(Some(Auth {
cipher: options.cipher,
auth_version: 1,
mode,
password: password.map(Option::Some).unwrap_or(None),
recipients,
}))
}
}
fn available_parallelism() -> usize {
#[cfg(target_arch = "wasm32")]
{
1
}
#[cfg(not(target_arch = "wasm32"))]
{
std::thread::available_parallelism().map_or(1, |threads| threads.get())
}
}
fn map_ordered<T, R>(
items: &[T],
workers: usize,
f: impl Fn(usize, &T) -> Result<R> + Sync,
) -> Result<Vec<R>>
where
T: Sync,
R: Send,
{
if workers <= 1 || items.len() <= 1 {
return items
.iter()
.enumerate()
.map(|(index, item)| f(index, item))
.collect();
}
let next = AtomicUsize::new(0);
let mut ordered: Vec<Option<R>> = (0..items.len()).map(|_| None).collect();
std::thread::scope(|scope| -> Result<()> {
let handles: Vec<_> = (0..workers.min(items.len()))
.map(|_| {
scope.spawn(|| -> Result<Vec<(usize, R)>> {
let mut finished = Vec::new();
loop {
let index = next.fetch_add(1, Ordering::Relaxed);
let Some(item) = items.get(index) else {
break;
};
finished.push((index, f(index, item)?));
}
Ok(finished)
})
})
.collect();
for handle in handles {
for (index, value) in handle.join().expect("a worker does not panic")? {
ordered[index] = Some(value);
}
}
Ok(())
})?;
Ok(ordered
.into_iter()
.map(|value| value.expect("every index is taken by one worker"))
.collect())
}
fn check_mask(pixels: &[u8], total_pixels: u32, what: &str) -> Result<()> {
if pixels.len() != total_pixels as usize {
return Err(Error::new(
Check::ReeDataSize,
format!(
"{what} has {} pixels but the display holds {total_pixels}",
pixels.len()
),
));
}
Ok(())
}
impl Pending {
fn raw(chunk_type: ChunkType, payload: Vec<u8>) -> Pending {
let len = payload.len() as u64;
Pending {
chunk_type,
stored: payload,
size_uncompressed: len,
size_compressed: 0,
flags: 0,
}
}
fn raw_with_flags(chunk_type: ChunkType, payload: Vec<u8>, flags: u32) -> Pending {
let mut pending = Pending::raw(chunk_type, payload);
pending.flags = flags;
pending
}
}
fn split_samples(plaintext: &[u8]) -> Vec<&[u8]> {
const WINDOW: usize = 1024;
if plaintext.is_empty() {
return Vec::new();
}
plaintext.chunks(WINDOW).collect()
}
fn dictionary_samples(plaintexts: &[Vec<u8>]) -> Vec<&[u8]> {
let windows: Vec<Vec<&[u8]>> = plaintexts.iter().map(|p| split_samples(p)).collect();
let mut samples = Vec::with_capacity(DICTIONARY_SAMPLE_WINDOWS);
if windows.is_empty() {
return samples;
}
let chunk_stride = windows.len().div_ceil(DICTIONARY_SAMPLE_WINDOWS).max(1);
let per_chunk = (DICTIONARY_SAMPLE_WINDOWS / (windows.len() / chunk_stride).max(1)).max(1);
for (index, chunk) in windows.iter().enumerate() {
if index % chunk_stride != 0 || chunk.is_empty() {
continue;
}
let step = (chunk.len() / per_chunk).max(1);
for window in (0..chunk.len()).step_by(step).take(per_chunk) {
if samples.len() == DICTIONARY_SAMPLE_WINDOWS {
return samples;
}
samples.push(chunk[window]);
}
}
samples
}
fn dictionary_pays(plaintexts: &[Vec<u8>], dictionary: &[u8], level: i32) -> Result<bool> {
let total: usize = plaintexts.iter().map(Vec::len).sum();
if total < DICTIONARY_PROBE_MIN_PLAINTEXT {
return Ok(true);
}
let share = DICTIONARY_PROBE_BYTES / 3;
let last = plaintexts.len() - 1;
let mut probe = Vec::with_capacity(DICTIONARY_PROBE_BYTES);
for index in [0, last / 2, last] {
let plaintext = &plaintexts[index];
probe.extend_from_slice(&plaintext[..plaintext.len().min(share)]);
}
if probe.is_empty() {
return Ok(true);
}
let with = chunks::compress(&probe, level, Some(dictionary))?.len();
let without = chunks::compress(&probe, level, None)?.len();
let charge = dictionary.len() as f64 * probe.len() as f64 / total as f64;
Ok((with as f64 + charge) < without as f64 * (1.0 - DICTIONARY_PROBE_MARGIN))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::json::{Meta, Timing};
const WIDTH: u32 = 8;
const HEIGHT: u32 = 4;
const PIXELS: usize = (WIDTH * HEIGHT) as usize;
fn head(layers: u32) -> Head {
Head {
head_version: 1,
encoder_name: "lumen writer tests".to_string(),
created_unix_sec: 1_750_000_000,
display_width_px: WIDTH,
display_height_px: HEIGHT,
build_width_um: 40_000,
build_depth_um: 20_000,
build_height_um: 20_000,
layer_height_um: 50,
total_layers: layers,
}
}
fn meta() -> Meta {
let mut timing = Timing {
layer_height_um: Some(50),
..Timing::default()
};
timing.extra.clear();
Meta {
meta_version: Some(1),
timing,
..Meta::default()
}
}
fn runs_of(pixels: &[u8]) -> Vec<Run> {
let mut runs: Vec<Run> = Vec::new();
for &pixel in pixels {
match runs.last_mut() {
Some(last) if last.value == pixel => last.length += 1,
_ => runs.push(Run::new(1, pixel)),
}
}
runs
}
fn masks() -> Vec<Vec<u8>> {
let mut stripes = vec![0u8; PIXELS];
for (index, pixel) in stripes.iter_mut().enumerate() {
if index % 3 == 0 {
*pixel = 255;
}
}
let mut edge = vec![0u8; PIXELS];
edge[PIXELS / 2 - 1] = 90;
edge[PIXELS / 2] = 200;
let gradient: Vec<u8> = (0..PIXELS).map(|index| (index * 8) as u8).collect();
vec![vec![0u8; PIXELS], stripes, edge, gradient]
}
#[test]
fn the_run_path_writes_the_same_file_as_the_pixel_path() {
let masks = masks();
let layers = masks.len() as u32;
let mut from_masks = Encoder::new(head(layers), meta());
let mut from_runs = Encoder::new(head(layers), meta());
let mut encoded = Encoder::new(head(layers), meta());
for pixels in &masks {
from_masks.push_layer(pixels).expect("a pushable mask");
from_runs
.push_layer_runs(&runs_of(pixels), EncodeMode::Auto)
.expect("a pushable run list");
let record = match ree::encode(pixels, PIXELS as u32, EncodeMode::Auto).unwrap() {
None => EncodedLayer::Empty,
Some((_tag, stream)) => EncodedLayer::single(stream).expect("a tagged stream"),
};
encoded
.push_encoded_layer(record)
.expect("a storable layer");
}
let expected = from_masks.finish().expect("a writable file");
assert_eq!(from_runs.finish().expect("a writable file"), expected);
assert_eq!(encoded.finish().expect("a writable file"), expected);
}
#[test]
fn the_worker_thread_count_does_not_change_the_file() {
let masks = masks();
let layers = masks.len() as u32;
let write = |threads: usize| {
let mut encoder = Encoder::new(head(layers), meta());
encoder.set_layers_per_chunk(1);
encoder.set_worker_threads(threads);
for pixels in &masks {
encoder.push_layer(pixels).expect("a pushable mask");
}
encoder.finish().expect("a writable file")
};
let serial = write(1);
for threads in [2, 4, 8] {
assert_eq!(write(threads), serial, "with {threads} workers");
}
}
#[test]
fn an_encoded_layer_is_sorted_and_may_be_empty() {
let masks = masks();
let mut sectors = Encoder::new(head(2), meta());
sectors
.push_encoded_layer(EncodedLayer::Sectors(vec![
(
1,
ree::encode(&masks[1], PIXELS as u32, EncodeMode::Auto)
.unwrap()
.unwrap()
.1,
),
(
0,
ree::encode(&masks[2], PIXELS as u32, EncodeMode::Auto)
.unwrap()
.unwrap()
.1,
),
]))
.expect("a storable layer");
sectors
.push_encoded_layer(EncodedLayer::Sectors(Vec::new()))
.expect("a storable empty layer");
let mut from_masks = Encoder::new(head(2), meta());
from_masks
.push_layer_sectors(&[(0, masks[2].clone()), (1, masks[1].clone())])
.expect("a pushable layer");
from_masks
.push_layer_sectors(&[])
.expect("a pushable empty layer");
assert_eq!(
sectors.finish().expect("a writable file"),
from_masks.finish().expect("a writable file")
);
}
fn conforming_meta() -> Meta {
Meta {
meta_version: Some(1),
timing: Timing {
normal_exposure_ms: Some(2500),
bottom_exposure_ms: Some(30000),
bottom_layer_count: Some(2),
transition_layer_count: Some(2),
layer_height_um: Some(50),
lift_slow_distance_um: Some(5000),
lift_slow_speed_um_min: Some(65000),
retract_fast_distance_um: Some(5000),
retract_fast_speed_um_min: Some(150000),
..Timing::default()
},
..Meta::default()
}
}
#[test]
fn equal_override_deltas_share_one_chunk() {
use crate::reader::LumenFile;
use crate::validate::Level;
let exposure = |ms: u32| Timing {
normal_exposure_ms: Some(ms),
..Timing::default()
};
let mut encoder = Encoder::new(head(4), conforming_meta());
for _ in 0..4 {
encoder.push_layer(&[0u8; PIXELS]).expect("a layer");
}
encoder
.set_overrides(vec![
Override {
layer: 0,
sector_id: 0,
timing: exposure(2800),
},
Override {
layer: 1,
sector_id: 0,
timing: exposure(2800),
},
Override {
layer: 2,
sector_id: 0,
timing: exposure(2800),
},
Override {
layer: 3,
sector_id: 0,
timing: exposure(3000),
},
])
.expect("one override set per (layer, sector)");
let bytes = encoder.finish().expect("a writable file");
let file =
LumenFile::open(&bytes, Level::Strict).expect("a shared chunk is a conforming file");
let first_lrov = |layer: u32| {
file.layer_table()
.entry(layer, 0)
.expect("an entry for every layer")
.first_lrov
};
assert_eq!(first_lrov(0), first_lrov(1), "equal deltas share a chunk");
assert_eq!(first_lrov(1), first_lrov(2), "equal deltas share a chunk");
assert_ne!(
first_lrov(2),
first_lrov(3),
"a delta that differs is its own chunk"
);
let chunks = file
.directory()
.entries()
.filter(|d| d.chunk_type == crate::container::ChunkType::LROV)
.count();
assert_eq!(chunks, 2, "four overrides with two distinct deltas");
for layer in 0..3 {
assert_eq!(
file.timing_for(layer, 0).expect("resolvable").exposure_ms,
2800
);
}
assert_eq!(file.timing_for(3, 0).expect("resolvable").exposure_ms, 3000);
}
#[test]
fn dictionary_samples_span_the_print() {
let plaintexts: Vec<Vec<u8>> = (1u8..=16).map(|value| vec![value; 4096]).collect();
let samples = dictionary_samples(&plaintexts);
let chunks: std::collections::BTreeSet<u8> = samples.iter().map(|s| s[0]).collect();
assert_eq!(chunks.len(), 16, "every chunk of the print contributes");
assert!(
samples.len() <= DICTIONARY_SAMPLE_WINDOWS,
"the budget holds"
);
assert!(
samples.iter().all(|sample| sample.len() == 1024),
"a sample is one window"
);
assert_eq!(
samples,
dictionary_samples(&plaintexts),
"sampling is deterministic"
);
}
#[test]
fn a_dictionary_that_does_not_pay_is_not_written() {
let layers = 8u32;
let mut head = head(layers);
head.display_width_px = 1024;
head.display_height_px = 1024;
let pixels = (head.display_width_px * head.display_height_px) as usize;
let mut state = 0x2545_F491_4F6C_DD1Du64;
let mut random = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
state >> 24
};
let mut encoder = Encoder::new(head, conforming_meta());
encoder.set_dictionary(true);
for _ in 0..layers {
let mask: Vec<u8> = (0..pixels).map(|_| random() as u8).collect();
encoder.push_layer(&mask).expect("a layer");
}
let bytes = encoder.finish().expect("a writable file");
let file = crate::reader::LumenFile::open(&bytes, crate::validate::Level::Strict)
.expect("a conforming file");
assert!(
file.dictionary().is_none(),
"a dictionary that makes the frames larger is not written"
);
}
fn moving_band_layers(count: u32) -> Vec<Vec<u8>> {
const SIDE: usize = 64;
let mut layers = Vec::new();
for index in 0..count as usize {
let mut pixels = vec![0u8; SIDE * SIDE];
for (row, line) in pixels.chunks_mut(SIDE).enumerate() {
let half = if row < 32 { row / 2 } else { (63 - row) / 2 };
let centre = 24 + index;
let start = centre.saturating_sub(half);
let end = (centre + half + 1).min(SIDE);
if start < end {
line[start..end].fill(255);
}
}
layers.push(pixels);
}
layers
}
const PIXELS_BAND: u32 = 64 * 64;
fn band_head(layers: u32) -> Head {
Head {
display_width_px: 64,
display_height_px: 64,
..head(layers)
}
}
fn written_with_probe(probe: bool) -> (Vec<u8>, Vec<Vec<u8>>) {
let layers = moving_band_layers(8);
let mut encoder = Encoder::new(band_head(layers.len() as u32), conforming_meta());
encoder.set_tag_probe(probe);
for layer in &layers {
encoder.push_layer(layer).expect("a pushable layer");
}
(encoder.finish().expect("a writable file"), layers)
}
#[test]
fn the_tag_probe_keeps_the_smaller_encoding() {
let (probed, layers) = written_with_probe(true);
let (plain, _) = written_with_probe(false);
for bytes in [&probed, &plain] {
let file = crate::LumenFile::open(bytes, crate::validate::Level::Strict)
.expect("a conforming file");
for (index, mask) in layers.iter().enumerate() {
assert_eq!(&file.layer(index as u32).unwrap().pixels, mask);
}
}
assert!(
probed.len() <= plain.len(),
"the probe must not write a larger file: {} against {}",
probed.len(),
plain.len()
);
assert!(
probed.len() < plain.len(),
"a print whose layers barely move is where the attached form wins: \
{} against {}",
probed.len(),
plain.len()
);
}
#[test]
fn a_pre_encoded_layer_can_carry_its_alternative() {
let layers = moving_band_layers(8);
let mut paired = Encoder::new(band_head(layers.len() as u32), conforming_meta());
let mut plain = Encoder::new(band_head(layers.len() as u32), conforming_meta());
for layer in &layers {
let (tag, primary) = ree::encode(layer, PIXELS_BAND, EncodeMode::Auto)
.unwrap()
.expect("a non-empty layer");
assert_eq!(tag, ree::TAG_BINARY);
let (other_tag, alternative) = ree::encode(layer, PIXELS_BAND, EncodeMode::Attached)
.unwrap()
.expect("a non-empty layer");
assert_eq!(other_tag, ree::TAG_ATTACHED);
paired
.push_encoded_layer_pair(
EncodedLayer::single(primary).unwrap(),
Some(EncodedLayer::single(alternative).unwrap()),
)
.expect("a pair of encodings of one layer");
plain
.push_encoded_layer_pair(
EncodedLayer::single(
ree::encode(layer, PIXELS_BAND, EncodeMode::Auto)
.unwrap()
.expect("a non-empty layer")
.1,
)
.unwrap(),
None,
)
.expect("one encoding");
}
let paired = paired.finish().expect("a writable file");
let plain = plain.finish().expect("a writable file");
for bytes in [&paired, &plain] {
let file = crate::LumenFile::open(bytes, crate::validate::Level::Strict)
.expect("a conforming file");
for (index, mask) in layers.iter().enumerate() {
assert_eq!(&file.layer(index as u32).unwrap().pixels, mask);
}
}
assert!(
paired.len() < plain.len(),
"the pair decides what the probe would have decided: {} against {}",
paired.len(),
plain.len()
);
}
#[test]
fn the_tag_probe_can_be_turned_off() {
let (probed, _) = written_with_probe(true);
let (plain, _) = written_with_probe(false);
assert_eq!(
plain.len(),
written_with_probe(false).0.len(),
"the same input writes the same file"
);
assert_ne!(probed.len(), plain.len());
}
#[test]
fn an_encoded_layer_is_checked_before_it_is_stored() {
let mut encoder = Encoder::new(head(1), meta());
for empty_or_unknown in [Vec::new(), vec![0x04, 0x00], vec![0xFF]] {
assert_eq!(
EncodedLayer::single(empty_or_unknown.clone())
.unwrap_err()
.check(),
Check::ReeTag
);
assert_eq!(
encoder
.push_encoded_layer(EncodedLayer::Sectors(vec![(0, empty_or_unknown)]))
.unwrap_err()
.check(),
Check::ReeTag
);
}
let stream = vec![ree::TAG_GRAYSCALE, 0x01, 0x00, PIXELS as u8];
assert_eq!(
encoder
.push_encoded_layer(EncodedLayer::Sectors(vec![
(2, stream.clone()),
(2, stream)
]))
.unwrap_err()
.check(),
Check::LtblSectorIdUnique
);
assert!(encoder.layers.is_empty());
}
}