Skip to main content

assetpack_core/pipeline/
mod.rs

1use crate::{
2  FastCdcSplitter,
3  codec::{Codec, compress, compress_auto},
4  error::Result,
5  file_transform::{FileHint, TransformSelection, TransformSelector},
6  hash::Hash32,
7  transform::TRANSFORM_ID_NONE,
8};
9
10#[derive(Debug, Clone, Copy)]
11pub enum ChunkCompressionPolicy {
12  Auto,
13  Raw,
14  ForceCodec(Codec),
15}
16
17#[derive(Debug, Clone)]
18pub struct PipelineConfig {
19  pub splitter: FastCdcSplitter,
20  pub chunk_compression: ChunkCompressionPolicy,
21  pub discard_payload: bool,
22}
23
24impl Default for PipelineConfig {
25  fn default() -> Self {
26    Self {
27      splitter: FastCdcSplitter::v1_defaults(),
28      chunk_compression: ChunkCompressionPolicy::Auto,
29      discard_payload: false,
30    }
31  }
32}
33
34#[derive(Debug, Clone)]
35pub struct ChunkPlan {
36  pub hash: Hash32,
37  pub raw_len: u32,
38  pub codec: Codec,
39  pub stored_len: u32,
40  pub payload: Option<Vec<u8>>,
41}
42
43#[derive(Debug, Clone)]
44pub struct PipelinePlan {
45  pub transform_id: u16,
46  pub transform_version: u16,
47  pub original_size: u64,
48  pub original_hash: Hash32,
49  pub stored_stream_size: u64,
50  pub chunks: Vec<ChunkPlan>,
51  pub estimated_bytes: u64,
52}
53
54pub trait ChunkCompressor: Send + Sync {
55  fn compress(&self, hash: Hash32, data: &[u8], extension: Option<&str>, policy: ChunkCompressionPolicy) -> Result<(Codec, Vec<u8>)>;
56}
57
58#[derive(Debug, Default)]
59pub struct DefaultChunkCompressor;
60
61impl ChunkCompressor for DefaultChunkCompressor {
62  fn compress(&self, _hash: Hash32, data: &[u8], extension: Option<&str>, policy: ChunkCompressionPolicy) -> Result<(Codec, Vec<u8>)> {
63    match policy {
64      ChunkCompressionPolicy::Auto => compress_auto(data, extension),
65      ChunkCompressionPolicy::Raw => Ok((Codec::Raw, data.to_vec())),
66      ChunkCompressionPolicy::ForceCodec(codec) => Ok((codec, compress(codec, data)?)),
67    }
68  }
69}
70
71#[derive(Debug, Clone)]
72pub struct Pipeline {
73  config: PipelineConfig,
74}
75
76impl Pipeline {
77  pub fn new(config: PipelineConfig) -> Self {
78    Self { config }
79  }
80
81  pub fn run(&self, data: Vec<u8>, hint: &FileHint, original_hash: Hash32, selector: Option<&TransformSelector>) -> Result<PipelinePlan> {
82    self.run_with_compressor(data, hint, original_hash, selector, &DefaultChunkCompressor)
83  }
84
85  pub fn run_with_compressor(
86    &self,
87    data: Vec<u8>,
88    hint: &FileHint,
89    original_hash: Hash32,
90    selector: Option<&TransformSelector>,
91    compressor: &dyn ChunkCompressor,
92  ) -> Result<PipelinePlan> {
93    let selection = if let Some(selector) = selector {
94      selector.select_bytes(data, hint, original_hash)?
95    } else {
96      TransformSelection::none(data, original_hash)
97    };
98    self.run_with_selection(selection, hint.extension.as_deref(), compressor)
99  }
100
101  pub fn run_with_selection(
102    &self,
103    selection: TransformSelection,
104    extension: Option<&str>,
105    compressor: &dyn ChunkCompressor,
106  ) -> Result<PipelinePlan> {
107    self.run_with_selection_ref(&selection, extension, compressor)
108  }
109
110  pub fn run_with_selection_ref(
111    &self,
112    selection: &TransformSelection,
113    extension: Option<&str>,
114    compressor: &dyn ChunkCompressor,
115  ) -> Result<PipelinePlan> {
116    let stored_stream_size = selection.stored_stream.len() as u64;
117    let compress_ext = if selection.transform_id == TRANSFORM_ID_NONE {
118      extension
119    } else {
120      None
121    };
122
123    let mut chunks = Vec::new();
124    let mut estimated_bytes = 0u64;
125    for chunk in self.config.splitter.split_bytes(&selection.stored_stream) {
126      let start = chunk.offset;
127      let end = start + chunk.length;
128      let slice = &selection.stored_stream[start..end];
129      let hash = Hash32::sha3_256(slice);
130      let (codec, stored) = compressor.compress(hash, slice, compress_ext, self.config.chunk_compression)?;
131      let stored_len = stored.len() as u32;
132      estimated_bytes += stored_len as u64;
133      let payload = if self.config.discard_payload { None } else { Some(stored) };
134      chunks.push(ChunkPlan {
135        hash,
136        raw_len: chunk.length as u32,
137        codec,
138        stored_len,
139        payload,
140      });
141    }
142    estimated_bytes += recipe_bytes_len(chunks.len());
143
144    Ok(PipelinePlan {
145      transform_id: selection.transform_id,
146      transform_version: selection.transform_version,
147      original_size: selection.original_size,
148      original_hash: selection.original_hash,
149      stored_stream_size,
150      chunks,
151      estimated_bytes,
152    })
153  }
154}
155
156fn recipe_bytes_len(chunk_count: usize) -> u64 {
157  let header = 1 + 8 + 32 + 2 + 2 + 4;
158  let per_chunk = 32 + 4;
159  (header + per_chunk * chunk_count) as u64
160}
161
162#[cfg(test)]
163mod tests {
164  use super::*;
165  use crate::codec::Codec;
166
167  #[test]
168  fn pipeline_discard_payload_strips_chunks() {
169    let config = PipelineConfig {
170      splitter: FastCdcSplitter::v1_defaults(),
171      chunk_compression: ChunkCompressionPolicy::Raw,
172      discard_payload: true,
173    };
174    let pipeline = Pipeline::new(config);
175    let data = b"pipeline payload test".to_vec();
176    let hash = Hash32::sha3_256(&data);
177    let selection = TransformSelection::none(data.clone(), hash);
178    let plan = pipeline
179      .run_with_selection(selection, Some("txt"), &DefaultChunkCompressor)
180      .unwrap();
181    assert!(!plan.chunks.is_empty());
182    for chunk in &plan.chunks {
183      assert!(chunk.payload.is_none());
184      assert_eq!(chunk.codec, Codec::Raw);
185      assert_eq!(chunk.stored_len, chunk.raw_len);
186    }
187  }
188}