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}