shape_runtime/stdlib/
compress.rs

1//! Native `compress` module for data compression and decompression.
2//!
3//! Exports: compress.gzip, compress.gunzip, compress.zstd, compress.unzstd,
4//!          compress.deflate, compress.inflate
5
6use crate::module_exports::{ModuleContext, ModuleExports, ModuleFunction, ModuleParam};
7use shape_value::ValueWord;
8use std::sync::Arc;
9use super::byte_utils::{bytes_from_array, bytes_to_array};
10
11/// Create the `compress` module with compression/decompression functions.
12pub fn create_compress_module() -> ModuleExports {
13    let mut module = ModuleExports::new("std::core::compress");
14    module.description = "Data compression and decompression (gzip, zstd, deflate)".to_string();
15
16    // compress.gzip(data: string) -> Array<int>
17    module.add_function_with_schema(
18        "gzip",
19        |args: &[ValueWord], _ctx: &ModuleContext| {
20            use flate2::Compression;
21            use flate2::write::GzEncoder;
22            use std::io::Write;
23
24            let data = args
25                .first()
26                .and_then(|a| a.as_str())
27                .ok_or_else(|| "compress.gzip() requires a string argument".to_string())?;
28
29            let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
30            encoder
31                .write_all(data.as_bytes())
32                .map_err(|e| format!("compress.gzip() failed: {}", e))?;
33            let compressed = encoder
34                .finish()
35                .map_err(|e| format!("compress.gzip() failed: {}", e))?;
36
37            Ok(bytes_to_array(&compressed))
38        },
39        ModuleFunction {
40            description: "Compress a string using gzip, returning a byte array".to_string(),
41            params: vec![ModuleParam {
42                name: "data".to_string(),
43                type_name: "string".to_string(),
44                required: true,
45                description: "String data to compress".to_string(),
46                ..Default::default()
47            }],
48            return_type: Some("Array<int>".to_string()),
49        },
50    );
51
52    // compress.gunzip(data: Array<int>) -> string
53    module.add_function_with_schema(
54        "gunzip",
55        |args: &[ValueWord], _ctx: &ModuleContext| {
56            use flate2::read::GzDecoder;
57            use std::io::Read;
58
59            let input = args
60                .first()
61                .ok_or_else(|| "compress.gunzip() requires an Array<int> argument".to_string())?;
62            let bytes = bytes_from_array(input).map_err(|e| format!("compress.gunzip(): {}", e))?;
63
64            let mut decoder = GzDecoder::new(&bytes[..]);
65            let mut output = String::new();
66            decoder
67                .read_to_string(&mut output)
68                .map_err(|e| format!("compress.gunzip() failed: {}", e))?;
69
70            Ok(ValueWord::from_string(Arc::new(output)))
71        },
72        ModuleFunction {
73            description: "Decompress a gzip byte array back to a string".to_string(),
74            params: vec![ModuleParam {
75                name: "data".to_string(),
76                type_name: "Array<int>".to_string(),
77                required: true,
78                description: "Gzip-compressed byte array".to_string(),
79                ..Default::default()
80            }],
81            return_type: Some("string".to_string()),
82        },
83    );
84
85    // compress.zstd(data: string, level?: int) -> Array<int>
86    module.add_function_with_schema(
87        "zstd",
88        |args: &[ValueWord], _ctx: &ModuleContext| {
89            let data = args
90                .first()
91                .and_then(|a| a.as_str())
92                .ok_or_else(|| "compress.zstd() requires a string argument".to_string())?;
93
94            let level = args
95                .get(1)
96                .and_then(|a| a.as_i64().or_else(|| a.as_f64().map(|n| n as i64)))
97                .unwrap_or(3) as i32;
98
99            let compressed = zstd::encode_all(data.as_bytes(), level)
100                .map_err(|e| format!("compress.zstd() failed: {}", e))?;
101
102            Ok(bytes_to_array(&compressed))
103        },
104        ModuleFunction {
105            description: "Compress a string using Zstandard, returning a byte array".to_string(),
106            params: vec![
107                ModuleParam {
108                    name: "data".to_string(),
109                    type_name: "string".to_string(),
110                    required: true,
111                    description: "String data to compress".to_string(),
112                    ..Default::default()
113                },
114                ModuleParam {
115                    name: "level".to_string(),
116                    type_name: "int".to_string(),
117                    required: false,
118                    description: "Compression level (default: 3)".to_string(),
119                    default_snippet: Some("3".to_string()),
120                    ..Default::default()
121                },
122            ],
123            return_type: Some("Array<int>".to_string()),
124        },
125    );
126
127    // compress.unzstd(data: Array<int>) -> string
128    module.add_function_with_schema(
129        "unzstd",
130        |args: &[ValueWord], _ctx: &ModuleContext| {
131            let input = args
132                .first()
133                .ok_or_else(|| "compress.unzstd() requires an Array<int> argument".to_string())?;
134            let bytes = bytes_from_array(input).map_err(|e| format!("compress.unzstd(): {}", e))?;
135
136            let decompressed = zstd::decode_all(&bytes[..])
137                .map_err(|e| format!("compress.unzstd() failed: {}", e))?;
138
139            let output = String::from_utf8(decompressed)
140                .map_err(|e| format!("compress.unzstd() invalid UTF-8: {}", e))?;
141
142            Ok(ValueWord::from_string(Arc::new(output)))
143        },
144        ModuleFunction {
145            description: "Decompress a Zstandard byte array back to a string".to_string(),
146            params: vec![ModuleParam {
147                name: "data".to_string(),
148                type_name: "Array<int>".to_string(),
149                required: true,
150                description: "Zstd-compressed byte array".to_string(),
151                ..Default::default()
152            }],
153            return_type: Some("string".to_string()),
154        },
155    );
156
157    // compress.deflate(data: string) -> Array<int>
158    module.add_function_with_schema(
159        "deflate",
160        |args: &[ValueWord], _ctx: &ModuleContext| {
161            use flate2::Compression;
162            use flate2::write::DeflateEncoder;
163            use std::io::Write;
164
165            let data = args
166                .first()
167                .and_then(|a| a.as_str())
168                .ok_or_else(|| "compress.deflate() requires a string argument".to_string())?;
169
170            let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default());
171            encoder
172                .write_all(data.as_bytes())
173                .map_err(|e| format!("compress.deflate() failed: {}", e))?;
174            let compressed = encoder
175                .finish()
176                .map_err(|e| format!("compress.deflate() failed: {}", e))?;
177
178            Ok(bytes_to_array(&compressed))
179        },
180        ModuleFunction {
181            description: "Compress a string using raw deflate, returning a byte array".to_string(),
182            params: vec![ModuleParam {
183                name: "data".to_string(),
184                type_name: "string".to_string(),
185                required: true,
186                description: "String data to compress".to_string(),
187                ..Default::default()
188            }],
189            return_type: Some("Array<int>".to_string()),
190        },
191    );
192
193    // compress.inflate(data: Array<int>) -> string
194    module.add_function_with_schema(
195        "inflate",
196        |args: &[ValueWord], _ctx: &ModuleContext| {
197            use flate2::read::DeflateDecoder;
198            use std::io::Read;
199
200            let input = args
201                .first()
202                .ok_or_else(|| "compress.inflate() requires an Array<int> argument".to_string())?;
203            let bytes =
204                bytes_from_array(input).map_err(|e| format!("compress.inflate(): {}", e))?;
205
206            let mut decoder = DeflateDecoder::new(&bytes[..]);
207            let mut output = String::new();
208            decoder
209                .read_to_string(&mut output)
210                .map_err(|e| format!("compress.inflate() failed: {}", e))?;
211
212            Ok(ValueWord::from_string(Arc::new(output)))
213        },
214        ModuleFunction {
215            description: "Decompress a raw deflate byte array back to a string".to_string(),
216            params: vec![ModuleParam {
217                name: "data".to_string(),
218                type_name: "Array<int>".to_string(),
219                required: true,
220                description: "Deflate-compressed byte array".to_string(),
221                ..Default::default()
222            }],
223            return_type: Some("string".to_string()),
224        },
225    );
226
227    module
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    fn test_ctx() -> crate::module_exports::ModuleContext<'static> {
235        let registry = Box::leak(Box::new(crate::type_schema::TypeSchemaRegistry::new()));
236        crate::module_exports::ModuleContext {
237            schemas: registry,
238            invoke_callable: None,
239            raw_invoker: None,
240            function_hashes: None,
241            vm_state: None,
242            granted_permissions: None,
243            scope_constraints: None,
244            set_pending_resume: None,
245            set_pending_frame_resume: None,
246        }
247    }
248
249    #[test]
250    fn test_compress_module_creation() {
251        let module = create_compress_module();
252        assert_eq!(module.name, "std::core::compress");
253        assert!(module.has_export("gzip"));
254        assert!(module.has_export("gunzip"));
255        assert!(module.has_export("zstd"));
256        assert!(module.has_export("unzstd"));
257        assert!(module.has_export("deflate"));
258        assert!(module.has_export("inflate"));
259    }
260
261    #[test]
262    fn test_gzip_roundtrip() {
263        let module = create_compress_module();
264        let ctx = test_ctx();
265        let gzip_fn = module.get_export("gzip").unwrap();
266        let gunzip_fn = module.get_export("gunzip").unwrap();
267
268        let input = ValueWord::from_string(Arc::new("hello world".to_string()));
269        let compressed = gzip_fn(&[input], &ctx).unwrap();
270
271        // Compressed should be an array
272        assert!(compressed.as_any_array().is_some());
273
274        let decompressed = gunzip_fn(&[compressed], &ctx).unwrap();
275        assert_eq!(decompressed.as_str(), Some("hello world"));
276    }
277
278    #[test]
279    fn test_zstd_roundtrip() {
280        let module = create_compress_module();
281        let ctx = test_ctx();
282        let zstd_fn = module.get_export("zstd").unwrap();
283        let unzstd_fn = module.get_export("unzstd").unwrap();
284
285        let input = ValueWord::from_string(Arc::new("hello zstd compression".to_string()));
286        let compressed = zstd_fn(&[input], &ctx).unwrap();
287
288        assert!(compressed.as_any_array().is_some());
289
290        let decompressed = unzstd_fn(&[compressed], &ctx).unwrap();
291        assert_eq!(decompressed.as_str(), Some("hello zstd compression"));
292    }
293
294    #[test]
295    fn test_zstd_with_level() {
296        let module = create_compress_module();
297        let ctx = test_ctx();
298        let zstd_fn = module.get_export("zstd").unwrap();
299        let unzstd_fn = module.get_export("unzstd").unwrap();
300
301        let input = ValueWord::from_string(Arc::new("level test".to_string()));
302        let level = ValueWord::from_i64(1);
303        let compressed = zstd_fn(&[input, level], &ctx).unwrap();
304
305        let decompressed = unzstd_fn(&[compressed], &ctx).unwrap();
306        assert_eq!(decompressed.as_str(), Some("level test"));
307    }
308
309    #[test]
310    fn test_deflate_roundtrip() {
311        let module = create_compress_module();
312        let ctx = test_ctx();
313        let deflate_fn = module.get_export("deflate").unwrap();
314        let inflate_fn = module.get_export("inflate").unwrap();
315
316        let input = ValueWord::from_string(Arc::new("deflate test data".to_string()));
317        let compressed = deflate_fn(&[input], &ctx).unwrap();
318
319        assert!(compressed.as_any_array().is_some());
320
321        let decompressed = inflate_fn(&[compressed], &ctx).unwrap();
322        assert_eq!(decompressed.as_str(), Some("deflate test data"));
323    }
324
325    #[test]
326    fn test_gzip_requires_string() {
327        let module = create_compress_module();
328        let ctx = test_ctx();
329        let gzip_fn = module.get_export("gzip").unwrap();
330
331        let result = gzip_fn(&[ValueWord::from_i64(42)], &ctx);
332        assert!(result.is_err());
333    }
334
335    #[test]
336    fn test_gunzip_invalid_data() {
337        let module = create_compress_module();
338        let ctx = test_ctx();
339        let gunzip_fn = module.get_export("gunzip").unwrap();
340
341        let bad_data = ValueWord::from_array(Arc::new(vec![
342            ValueWord::from_i64(1),
343            ValueWord::from_i64(2),
344            ValueWord::from_i64(3),
345        ]));
346        let result = gunzip_fn(&[bad_data], &ctx);
347        assert!(result.is_err());
348    }
349
350    #[test]
351    fn test_empty_string_roundtrip() {
352        let module = create_compress_module();
353        let ctx = test_ctx();
354        let gzip_fn = module.get_export("gzip").unwrap();
355        let gunzip_fn = module.get_export("gunzip").unwrap();
356
357        let input = ValueWord::from_string(Arc::new(String::new()));
358        let compressed = gzip_fn(&[input], &ctx).unwrap();
359        let decompressed = gunzip_fn(&[compressed], &ctx).unwrap();
360        assert_eq!(decompressed.as_str(), Some(""));
361    }
362
363    #[test]
364    fn test_large_data_roundtrip() {
365        let module = create_compress_module();
366        let ctx = test_ctx();
367        let gzip_fn = module.get_export("gzip").unwrap();
368        let gunzip_fn = module.get_export("gunzip").unwrap();
369
370        let large = "a".repeat(10_000);
371        let input = ValueWord::from_string(Arc::new(large.clone()));
372        let compressed = gzip_fn(&[input], &ctx).unwrap();
373
374        // Compressed should be smaller than original
375        let arr = compressed.as_any_array().unwrap().to_generic();
376        assert!(arr.len() < 10_000);
377
378        let decompressed = gunzip_fn(&[compressed], &ctx).unwrap();
379        assert_eq!(decompressed.as_str(), Some(large.as_str()));
380    }
381
382    #[test]
383    fn test_schemas() {
384        let module = create_compress_module();
385
386        let gzip_schema = module.get_schema("gzip").unwrap();
387        assert_eq!(gzip_schema.params.len(), 1);
388        assert!(gzip_schema.params[0].required);
389        assert_eq!(gzip_schema.return_type.as_deref(), Some("Array<int>"));
390
391        let zstd_schema = module.get_schema("zstd").unwrap();
392        assert_eq!(zstd_schema.params.len(), 2);
393        assert!(zstd_schema.params[0].required);
394        assert!(!zstd_schema.params[1].required);
395
396        let gunzip_schema = module.get_schema("gunzip").unwrap();
397        assert_eq!(gunzip_schema.return_type.as_deref(), Some("string"));
398    }
399}
shape_runtime/stdlib/compress.rs

shape_runtime/stdlib/
compress.rs
