zstd_rs/
par.rs

1//! The parallelism boundary. The core stays pure and single-threaded: it plans
2//! independent jobs, compresses or decompresses one job at a time, and assembles
3//! results. *How* jobs run (native threads, Web Workers, or one after another) is
4//! an [`Executor`] the caller supplies; nothing here picks one implicitly.
5//! Output is byte-identical for every executor, because each job's bytes depend
6//! only on (config, input, job, dictionary).
7use crate::decode::{DecoderDictionary, Decompressor};
8pub use crate::encode::Job;
9use crate::encode::{begin_frame, write_checksum, CompressionConfig, Compressor, EncoderDictionary};
10use crate::error::{Error, Result};
11use crate::frame::{parse_block_header, BlockType, FrameHeader, SKIPPABLE_BASE, SKIPPABLE_MASK};
12use alloc::vec::Vec;
13
14/// Runs `f` over `items` and returns the results in input order.
15pub trait Executor {
16    /// Applies `f` to every item. Implementations may run calls concurrently.
17    fn map<T: Sync, R: Send>(&self, items: &[T], f: &(dyn Fn(&T) -> R + Sync)) -> Vec<R>;
18}
19
20/// Runs jobs one after another on the calling thread (an explicit choice, not a default).
21#[derive(Clone, Copy, Debug)]
22pub struct Sequential;
23
24impl Executor for Sequential {
25    fn map<T: Sync, R: Send>(&self, items: &[T], f: &(dyn Fn(&T) -> R + Sync)) -> Vec<R> {
26        items.iter().map(f).collect()
27    }
28}
29
30/// Splits `len` input bytes into jobs of `job_size` bytes, each allowed to
31/// reference up to `overlap` bytes before its start (clamped to the frame
32/// window when the job runs).
33pub fn plan(len: usize, job_size: usize, overlap: usize) -> Result<Vec<Job>> {
34    if job_size < 1024 {
35        return Err(Error::Parameter("job_size must be at least 1024"));
36    }
37    if len == 0 {
38        return Ok(alloc::vec![Job { start: 0, end: 0, prefix_start: 0, first: true, last: true }]);
39    }
40    let mut jobs = Vec::with_capacity(len.div_ceil(job_size));
41    let mut start = 0;
42    while start < len {
43        let end = (start + job_size).min(len);
44        jobs.push(Job { start, end, prefix_start: start.saturating_sub(overlap), first: start == 0, last: end == len });
45        start = end;
46    }
47    Ok(jobs)
48}
49
50fn check_plan(len: usize, jobs: &[Job]) -> Result<()> {
51    let mut at = 0;
52    for (i, j) in jobs.iter().enumerate() {
53        if j.start != at || j.end < j.start || j.first != (i == 0) || j.last != (i + 1 == jobs.len()) || j.prefix_start > j.start {
54            return Err(Error::Parameter("jobs must tile the input in order"));
55        }
56        at = j.end;
57    }
58    if at != len || jobs.is_empty() {
59        return Err(Error::Parameter("jobs must cover the whole input"));
60    }
61    Ok(())
62}
63
64/// Writes the header of a frame holding `len` bytes, for callers that run jobs
65/// themselves (for example on Web Workers): header, then every job's
66/// [`Compressor::compress_job`] output in order, then [`frame_trailer`].
67pub fn frame_header(cfg: &CompressionConfig, len: usize, dict: Option<&EncoderDictionary>, out: &mut Vec<u8>) -> Result<()> {
68    begin_frame(cfg, len, dict, out)
69}
70
71/// Writes the end of a frame over `src` (the checksum, when configured).
72pub fn frame_trailer(cfg: &CompressionConfig, src: &[u8], out: &mut Vec<u8>) {
73    write_checksum(cfg, src, out);
74}
75
76/// Compresses `src` into one frame whose jobs run on `exec`.
77pub fn compress_frame<E: Executor>(
78    exec: &E,
79    cfg: &CompressionConfig,
80    src: &[u8],
81    dict: Option<&EncoderDictionary>,
82    jobs: &[Job],
83    out: &mut Vec<u8>,
84) -> Result<()> {
85    check_plan(src.len(), jobs)?;
86    begin_frame(cfg, src.len(), dict, out)?;
87    let parts = exec.map(jobs, &|job: &Job| -> Result<Vec<u8>> {
88        let mut c = Compressor::new(*cfg)?;
89        let mut o = Vec::new();
90        c.compress_job(src, job, dict, &mut o)?;
91        Ok(o)
92    });
93    for p in parts {
94        out.extend_from_slice(&p?);
95    }
96    write_checksum(cfg, src, out);
97    Ok(())
98}
99
100/// Compresses `src` as independent frames of `frame_size` bytes (each decodable
101/// on its own, and in parallel by [`decompress_frames`]).
102pub fn compress_frames<E: Executor>(
103    exec: &E,
104    cfg: &CompressionConfig,
105    src: &[u8],
106    dict: Option<&EncoderDictionary>,
107    frame_size: usize,
108    out: &mut Vec<u8>,
109) -> Result<()> {
110    if frame_size == 0 {
111        return Err(Error::Parameter("frame_size must be positive"));
112    }
113    let chunks: Vec<&[u8]> = if src.is_empty() { alloc::vec![src] } else { src.chunks(frame_size).collect() };
114    let parts = exec.map(&chunks, &|chunk: &&[u8]| -> Result<Vec<u8>> {
115        let mut o = Vec::new();
116        Compressor::new(*cfg)?.compress(chunk, dict, &mut o)?;
117        Ok(o)
118    });
119    for p in parts {
120        out.extend_from_slice(&p?);
121    }
122    Ok(())
123}
124
125/// Location of one frame inside a buffer of concatenated frames.
126#[derive(Clone, Copy, Debug, PartialEq, Eq)]
127pub struct FrameSpan {
128    /// Offset of the frame's first byte.
129    pub start: usize,
130    /// Offset just past the frame.
131    pub end: usize,
132    /// Declared content size (`None` for skippable frames or frames without one).
133    pub content_size: Option<u64>,
134    /// True for skippable frames (no content).
135    pub skippable: bool,
136}
137
138/// Finds frame boundaries by reading headers and block headers only.
139pub fn split_frames(src: &[u8]) -> Result<Vec<FrameSpan>> {
140    let mut spans = Vec::new();
141    let mut at = 0;
142    while at < src.len() {
143        let rest = &src[at..];
144        let magic = u32::from_le_bytes(rest.get(..4).ok_or(Error::Truncated)?.try_into().map_err(|_| Error::Truncated)?);
145        let (len, content_size, skippable) = if magic & SKIPPABLE_MASK == SKIPPABLE_BASE {
146            let n = u32::from_le_bytes(rest.get(4..8).ok_or(Error::Truncated)?.try_into().map_err(|_| Error::Truncated)?);
147            (8 + n as usize, None, true)
148        } else {
149            let h = FrameHeader::parse(rest)?;
150            let mut p = h.header_len;
151            loop {
152                let (last, ty, size) = parse_block_header(rest.get(p..).ok_or(Error::Truncated)?)?;
153                p += 3 + if ty == BlockType::Rle { 1 } else { size };
154                if last {
155                    break;
156                }
157            }
158            (p + if h.checksum { 4 } else { 0 }, h.content_size, false)
159        };
160        if at + len > src.len() {
161            return Err(Error::Truncated);
162        }
163        spans.push(FrameSpan { start: at, end: at + len, content_size, skippable });
164        at += len;
165    }
166    Ok(spans)
167}
168
169/// Decodes concatenated frames on `exec`, appending to `out`. Every frame must
170/// declare its content size (so each job's output is bounded before it runs);
171/// the total may not exceed `max_output`.
172pub fn decompress_frames<E: Executor>(
173    exec: &E,
174    src: &[u8],
175    dict: Option<&DecoderDictionary>,
176    max_output: usize,
177    out: &mut Vec<u8>,
178) -> Result<usize> {
179    let spans: Vec<FrameSpan> = split_frames(src)?.into_iter().filter(|s| !s.skippable).collect();
180    let mut total: u64 = 0;
181    for s in &spans {
182        total += s.content_size.ok_or(Error::Parameter("parallel decoding needs every frame's content size"))?;
183    }
184    if total > max_output as u64 {
185        return Err(Error::OutputLimit);
186    }
187    let parts = exec.map(&spans, &|s: &FrameSpan| -> Result<Vec<u8>> {
188        let cs = s.content_size.unwrap_or(0) as usize;
189        let mut o = Vec::with_capacity(cs);
190        Decompressor::new().decompress_frame(&src[s.start..s.end], dict, cs, &mut o)?;
191        Ok(o)
192    });
193    let before = out.len();
194    for p in parts {
195        out.extend_from_slice(&p?);
196    }
197    Ok(out.len() - before)
198}
zstd_rs/par.rs

zstd_rs/
par.rs
