zrip 0.1.2 → 0.2.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/CHANGELOG.md +12 -0
- data/Cargo.lock +5 -41
- data/README.md +13 -40
- data/ext/zrip/Cargo.toml +2 -2
- data/ext/zrip/build.rs +23 -0
- data/ext/zrip/src/lib.rs +697 -300
- data/ext/zrip/src/rb.rs +615 -0
- data/lib/zrip/block_codec.rb +49 -0
- data/lib/zrip/dict_trainer.rb +32 -0
- data/lib/zrip/dictionary.rb +10 -0
- data/lib/zrip/frame_codec.rb +58 -0
- data/lib/zrip/version.rb +1 -1
- data/lib/zrip.rb +14 -0
- metadata +6 -3
data/ext/zrip/src/lib.rs
CHANGED
|
@@ -1,44 +1,80 @@
|
|
|
1
|
-
|
|
2
|
-
exception::ExceptionClass, function, method, prelude::*, r_string::RString, value::Opaque,
|
|
3
|
-
Error, Ruby,
|
|
4
|
-
};
|
|
5
|
-
use std::cell::RefCell;
|
|
6
|
-
use std::sync::{Mutex, OnceLock};
|
|
1
|
+
mod rb;
|
|
7
2
|
|
|
3
|
+
use std::ffi::c_void;
|
|
4
|
+
use std::io::Read;
|
|
5
|
+
use std::panic::{catch_unwind, AssertUnwindSafe};
|
|
6
|
+
use std::sync::{Mutex, OnceLock, TryLockError};
|
|
7
|
+
|
|
8
|
+
use rb_sys::{rb_data_type_struct__bindgen_ty_1, rb_data_type_t, size_t, VALUE};
|
|
8
9
|
use zstd::dict::fastcover::FastCoverParams;
|
|
9
10
|
use zstd::{CompressContext, DecompressContext, Dictionary};
|
|
10
11
|
|
|
11
|
-
|
|
12
|
-
static COMPRESS_ERROR: OnceLock<Opaque<ExceptionClass>> = OnceLock::new();
|
|
13
|
-
static MISSING_CONTENT_SIZE_ERROR: OnceLock<Opaque<ExceptionClass>> = OnceLock::new();
|
|
14
|
-
static OUTPUT_SIZE_LIMIT_ERROR: OnceLock<Opaque<ExceptionClass>> = OnceLock::new();
|
|
12
|
+
use crate::rb::{RbResult, RubyErr};
|
|
15
13
|
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
14
|
+
static DECOMPRESS_ERROR: OnceLock<GlobalValue> = OnceLock::new();
|
|
15
|
+
static COMPRESS_ERROR: OnceLock<GlobalValue> = OnceLock::new();
|
|
16
|
+
static MISSING_CONTENT_SIZE_ERROR: OnceLock<GlobalValue> = OnceLock::new();
|
|
17
|
+
static OUTPUT_SIZE_LIMIT_ERROR: OnceLock<GlobalValue> = OnceLock::new();
|
|
18
|
+
|
|
19
|
+
const GVL_COMPRESS_THRESHOLD: usize = 64 * 1024;
|
|
20
|
+
const GVL_FRAME_DECOMPRESS_THRESHOLD: usize = 64 * 1024;
|
|
21
|
+
const DECOMPRESS_READ_CHUNK: usize = 16 * 1024;
|
|
22
|
+
|
|
23
|
+
#[derive(Copy, Clone)]
|
|
24
|
+
struct GlobalValue(VALUE);
|
|
25
|
+
|
|
26
|
+
unsafe impl Send for GlobalValue {}
|
|
27
|
+
unsafe impl Sync for GlobalValue {}
|
|
28
|
+
|
|
29
|
+
fn stored_value(lock: &OnceLock<GlobalValue>, name: &str) -> VALUE {
|
|
30
|
+
lock.get()
|
|
31
|
+
.unwrap_or_else(|| panic!("{name} not initialized"))
|
|
32
|
+
.0
|
|
22
33
|
}
|
|
23
34
|
|
|
24
|
-
fn
|
|
25
|
-
|
|
35
|
+
fn decompress_error() -> VALUE {
|
|
36
|
+
stored_value(&DECOMPRESS_ERROR, "DecompressError")
|
|
26
37
|
}
|
|
27
38
|
|
|
28
|
-
fn
|
|
29
|
-
|
|
30
|
-
*MISSING_CONTENT_SIZE_ERROR
|
|
31
|
-
.get()
|
|
32
|
-
.expect("MissingContentSizeError not initialized"),
|
|
33
|
-
)
|
|
39
|
+
fn compress_error() -> VALUE {
|
|
40
|
+
stored_value(&COMPRESS_ERROR, "CompressError")
|
|
34
41
|
}
|
|
35
42
|
|
|
36
|
-
fn output_size_limit_error(
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
43
|
+
fn output_size_limit_error() -> VALUE {
|
|
44
|
+
stored_value(&OUTPUT_SIZE_LIMIT_ERROR, "OutputSizeLimitError")
|
|
45
|
+
}
|
|
46
|
+
|
|
47
|
+
fn should_release_compress_gvl(input_len: usize) -> bool {
|
|
48
|
+
input_len >= GVL_COMPRESS_THRESHOLD
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
fn should_release_frame_decompress_gvl(output_len: usize) -> bool {
|
|
52
|
+
output_len >= GVL_FRAME_DECOMPRESS_THRESHOLD
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
fn with_mutex<T, R, F>(mutex: &Mutex<T>, release_gvl: bool, name: &str, func: F) -> RbResult<R>
|
|
56
|
+
where
|
|
57
|
+
F: FnOnce(&mut T) -> RbResult<R>,
|
|
58
|
+
{
|
|
59
|
+
if release_gvl {
|
|
60
|
+
return rb::maybe_without_gvl(true, || {
|
|
61
|
+
let mut guard = mutex
|
|
62
|
+
.lock()
|
|
63
|
+
.map_err(|_| RubyErr::runtime(format!("{name} mutex poisoned")))?;
|
|
64
|
+
func(&mut guard)
|
|
65
|
+
});
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
match mutex.try_lock() {
|
|
69
|
+
Ok(mut guard) => func(&mut guard),
|
|
70
|
+
Err(TryLockError::WouldBlock) => rb::maybe_without_gvl(true, || {
|
|
71
|
+
let mut guard = mutex
|
|
72
|
+
.lock()
|
|
73
|
+
.map_err(|_| RubyErr::runtime(format!("{name} mutex poisoned")))?;
|
|
74
|
+
func(&mut guard)
|
|
75
|
+
}),
|
|
76
|
+
Err(TryLockError::Poisoned(_)) => Err(RubyErr::runtime(format!("{name} mutex poisoned"))),
|
|
77
|
+
}
|
|
42
78
|
}
|
|
43
79
|
|
|
44
80
|
// ---------- frame header parsing ----------
|
|
@@ -48,8 +84,7 @@ const ZSTD_FRAME_MAGIC: [u8; 4] = [0x28, 0xB5, 0x2F, 0xFD];
|
|
|
48
84
|
#[derive(Debug)]
|
|
49
85
|
enum BoundedError {
|
|
50
86
|
BadMagic,
|
|
51
|
-
|
|
52
|
-
OutputSizeLimit { declared: u64, limit: u64 },
|
|
87
|
+
OutputSizeLimit { limit: u64 },
|
|
53
88
|
DecoderFailed(String),
|
|
54
89
|
}
|
|
55
90
|
|
|
@@ -98,83 +133,246 @@ fn parse_frame_content_size(input: &[u8]) -> Result<Option<u64>, BoundedError> {
|
|
|
98
133
|
Ok(Some(value))
|
|
99
134
|
}
|
|
100
135
|
|
|
136
|
+
fn io_error_is_output_too_small(err: &std::io::Error) -> bool {
|
|
137
|
+
matches!(
|
|
138
|
+
err.get_ref()
|
|
139
|
+
.and_then(|inner| inner.downcast_ref::<zstd::error::DecompressError>()),
|
|
140
|
+
Some(zstd::error::DecompressError::OutputTooSmall)
|
|
141
|
+
)
|
|
142
|
+
}
|
|
143
|
+
|
|
144
|
+
fn read_decoder_bounded<R: Read>(
|
|
145
|
+
decoder: &mut R,
|
|
146
|
+
max_output: usize,
|
|
147
|
+
) -> Result<Vec<u8>, BoundedError> {
|
|
148
|
+
let mut output = Vec::new();
|
|
149
|
+
let mut buf = [0u8; DECOMPRESS_READ_CHUNK];
|
|
150
|
+
|
|
151
|
+
loop {
|
|
152
|
+
let n = decoder.read(&mut buf).map_err(|e| {
|
|
153
|
+
if io_error_is_output_too_small(&e) {
|
|
154
|
+
BoundedError::OutputSizeLimit {
|
|
155
|
+
limit: max_output as u64,
|
|
156
|
+
}
|
|
157
|
+
} else {
|
|
158
|
+
BoundedError::DecoderFailed(format!("{e}"))
|
|
159
|
+
}
|
|
160
|
+
})?;
|
|
161
|
+
if n == 0 {
|
|
162
|
+
return Ok(output);
|
|
163
|
+
}
|
|
164
|
+
|
|
165
|
+
let Some(next_len) = output.len().checked_add(n) else {
|
|
166
|
+
return Err(BoundedError::OutputSizeLimit {
|
|
167
|
+
limit: max_output as u64,
|
|
168
|
+
});
|
|
169
|
+
};
|
|
170
|
+
if next_len > max_output {
|
|
171
|
+
return Err(BoundedError::OutputSizeLimit {
|
|
172
|
+
limit: max_output as u64,
|
|
173
|
+
});
|
|
174
|
+
}
|
|
175
|
+
|
|
176
|
+
output.extend_from_slice(&buf[..n]);
|
|
177
|
+
}
|
|
178
|
+
}
|
|
179
|
+
|
|
101
180
|
fn decompress_bounded(
|
|
102
181
|
compressed: &[u8],
|
|
103
182
|
max_output: usize,
|
|
104
183
|
dctx: &mut DecompressContext,
|
|
184
|
+
dict: Option<&Dictionary>,
|
|
105
185
|
) -> Result<Vec<u8>, BoundedError> {
|
|
106
186
|
if compressed.len() < ZSTD_FRAME_MAGIC.len() || compressed[..4] != ZSTD_FRAME_MAGIC {
|
|
107
187
|
return Err(BoundedError::BadMagic);
|
|
108
188
|
}
|
|
109
189
|
|
|
110
|
-
let
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
if n > u64::from(u32::MAX) {
|
|
119
|
-
return Err(BoundedError::OutputSizeLimit {
|
|
120
|
-
declared: n,
|
|
121
|
-
limit: u64::from(u32::MAX),
|
|
122
|
-
});
|
|
123
|
-
}
|
|
124
|
-
n as usize
|
|
125
|
-
}
|
|
126
|
-
None => {
|
|
127
|
-
if max_output != 0 {
|
|
128
|
-
return Err(BoundedError::MissingContentSize);
|
|
190
|
+
let result = if max_output == 0 {
|
|
191
|
+
dctx.decompress_with_limit(compressed, usize::MAX)
|
|
192
|
+
.map(|out| out.into_owned())
|
|
193
|
+
.map_err(|e| BoundedError::DecoderFailed(format!("{e}")))?
|
|
194
|
+
} else {
|
|
195
|
+
let mut decoder = match dict {
|
|
196
|
+
Some(dict) => {
|
|
197
|
+
zstd::FrameDecoder::with_dict_and_limit(compressed, dict.clone(), usize::MAX)
|
|
129
198
|
}
|
|
130
|
-
|
|
131
|
-
}
|
|
199
|
+
None => zstd::FrameDecoder::with_limit(compressed, usize::MAX),
|
|
200
|
+
};
|
|
201
|
+
read_decoder_bounded(&mut decoder, max_output)?
|
|
132
202
|
};
|
|
133
203
|
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
.map_err(|e| BoundedError::DecoderFailed(format!("{e}")))?;
|
|
204
|
+
Ok(result)
|
|
205
|
+
}
|
|
137
206
|
|
|
138
|
-
|
|
207
|
+
fn decompress_work_size(compressed: &[u8], max_output: usize) -> usize {
|
|
208
|
+
match parse_frame_content_size(compressed) {
|
|
209
|
+
Ok(Some(n)) => usize::try_from(n).unwrap_or(usize::MAX),
|
|
210
|
+
Ok(None) if max_output != 0 => max_output,
|
|
211
|
+
_ => compressed.len(),
|
|
212
|
+
}
|
|
139
213
|
}
|
|
140
214
|
|
|
141
|
-
fn
|
|
215
|
+
fn bounded_err(err: BoundedError, prefix: &str) -> RubyErr {
|
|
142
216
|
match err {
|
|
143
|
-
BoundedError::BadMagic =>
|
|
144
|
-
decompress_error(
|
|
217
|
+
BoundedError::BadMagic => RubyErr::new(
|
|
218
|
+
decompress_error(),
|
|
145
219
|
format!("{prefix}: bad magic (input is not a Zstd frame)"),
|
|
146
220
|
),
|
|
147
|
-
BoundedError::
|
|
148
|
-
|
|
149
|
-
format!("{prefix}:
|
|
150
|
-
),
|
|
151
|
-
BoundedError::OutputSizeLimit { declared, limit } => Error::new(
|
|
152
|
-
output_size_limit_error(ruby),
|
|
153
|
-
format!("{prefix}: declared content size {declared} exceeds limit {limit}"),
|
|
221
|
+
BoundedError::OutputSizeLimit { limit } => RubyErr::new(
|
|
222
|
+
output_size_limit_error(),
|
|
223
|
+
format!("{prefix}: decompressed output exceeds limit {limit}"),
|
|
154
224
|
),
|
|
155
225
|
BoundedError::DecoderFailed(msg) => {
|
|
156
|
-
|
|
226
|
+
RubyErr::new(decompress_error(), format!("{prefix}: {msg}"))
|
|
157
227
|
}
|
|
158
228
|
}
|
|
159
229
|
}
|
|
160
230
|
|
|
161
231
|
// ---------- dict helper ----------
|
|
162
232
|
|
|
163
|
-
fn load_dict(
|
|
233
|
+
fn load_dict(bytes: &[u8]) -> RbResult<Dictionary> {
|
|
164
234
|
Dictionary::from_bytes(bytes).map_err(|_| {
|
|
165
|
-
|
|
166
|
-
ruby.exception_runtime_error(),
|
|
167
|
-
"dictionary must be in ZDICT format (use DictTrainer to train one)",
|
|
168
|
-
)
|
|
235
|
+
RubyErr::runtime("dictionary must be in ZDICT format (use DictTrainer to train one)")
|
|
169
236
|
})
|
|
170
237
|
}
|
|
171
238
|
|
|
239
|
+
// ---------- typed data ----------
|
|
240
|
+
|
|
241
|
+
struct NativeDataType(rb_data_type_t);
|
|
242
|
+
|
|
243
|
+
unsafe impl Send for NativeDataType {}
|
|
244
|
+
unsafe impl Sync for NativeDataType {}
|
|
245
|
+
|
|
246
|
+
static FRAME_CODEC_DATA_TYPE: OnceLock<NativeDataType> = OnceLock::new();
|
|
247
|
+
static BLOCK_CODEC_DATA_TYPE: OnceLock<NativeDataType> = OnceLock::new();
|
|
248
|
+
static DICT_TRAINER_DATA_TYPE: OnceLock<NativeDataType> = OnceLock::new();
|
|
249
|
+
|
|
250
|
+
fn frame_codec_data_type() -> *const rb_data_type_t {
|
|
251
|
+
&FRAME_CODEC_DATA_TYPE
|
|
252
|
+
.get_or_init(|| NativeDataType(make_frame_codec_data_type()))
|
|
253
|
+
.0
|
|
254
|
+
}
|
|
255
|
+
|
|
256
|
+
fn block_codec_data_type() -> *const rb_data_type_t {
|
|
257
|
+
&BLOCK_CODEC_DATA_TYPE
|
|
258
|
+
.get_or_init(|| NativeDataType(make_block_codec_data_type()))
|
|
259
|
+
.0
|
|
260
|
+
}
|
|
261
|
+
|
|
262
|
+
fn dict_trainer_data_type() -> *const rb_data_type_t {
|
|
263
|
+
&DICT_TRAINER_DATA_TYPE
|
|
264
|
+
.get_or_init(|| NativeDataType(make_dict_trainer_data_type()))
|
|
265
|
+
.0
|
|
266
|
+
}
|
|
267
|
+
|
|
268
|
+
fn make_frame_codec_data_type() -> rb_data_type_t {
|
|
269
|
+
rb_data_type_t {
|
|
270
|
+
wrap_struct_name: c"zrip_frame_codec".as_ptr(),
|
|
271
|
+
function: rb_data_type_struct__bindgen_ty_1 {
|
|
272
|
+
dmark: None,
|
|
273
|
+
dfree: Some(frame_codec_free),
|
|
274
|
+
dsize: Some(frame_codec_native_size),
|
|
275
|
+
dcompact: None,
|
|
276
|
+
reserved: [std::ptr::null_mut(); 1],
|
|
277
|
+
},
|
|
278
|
+
parent: std::ptr::null(),
|
|
279
|
+
data: std::ptr::null_mut(),
|
|
280
|
+
flags: 1,
|
|
281
|
+
}
|
|
282
|
+
}
|
|
283
|
+
|
|
284
|
+
fn make_block_codec_data_type() -> rb_data_type_t {
|
|
285
|
+
rb_data_type_t {
|
|
286
|
+
wrap_struct_name: c"zrip_block_codec".as_ptr(),
|
|
287
|
+
function: rb_data_type_struct__bindgen_ty_1 {
|
|
288
|
+
dmark: None,
|
|
289
|
+
dfree: Some(block_codec_free),
|
|
290
|
+
dsize: Some(block_codec_native_size),
|
|
291
|
+
dcompact: None,
|
|
292
|
+
reserved: [std::ptr::null_mut(); 1],
|
|
293
|
+
},
|
|
294
|
+
parent: std::ptr::null(),
|
|
295
|
+
data: std::ptr::null_mut(),
|
|
296
|
+
flags: 1,
|
|
297
|
+
}
|
|
298
|
+
}
|
|
299
|
+
|
|
300
|
+
fn make_dict_trainer_data_type() -> rb_data_type_t {
|
|
301
|
+
rb_data_type_t {
|
|
302
|
+
wrap_struct_name: c"zrip_dict_trainer".as_ptr(),
|
|
303
|
+
function: rb_data_type_struct__bindgen_ty_1 {
|
|
304
|
+
dmark: None,
|
|
305
|
+
dfree: Some(dict_trainer_free),
|
|
306
|
+
dsize: Some(dict_trainer_native_size),
|
|
307
|
+
dcompact: None,
|
|
308
|
+
reserved: [std::ptr::null_mut(); 1],
|
|
309
|
+
},
|
|
310
|
+
parent: std::ptr::null(),
|
|
311
|
+
data: std::ptr::null_mut(),
|
|
312
|
+
flags: 1,
|
|
313
|
+
}
|
|
314
|
+
}
|
|
315
|
+
|
|
316
|
+
unsafe extern "C" fn frame_codec_free(ptr: *mut c_void) {
|
|
317
|
+
if ptr.is_null() {
|
|
318
|
+
return;
|
|
319
|
+
}
|
|
320
|
+
|
|
321
|
+
let _ = catch_unwind(AssertUnwindSafe(|| unsafe {
|
|
322
|
+
drop(Box::from_raw(ptr as *mut FrameCodec));
|
|
323
|
+
}));
|
|
324
|
+
}
|
|
325
|
+
|
|
326
|
+
unsafe extern "C" fn block_codec_free(ptr: *mut c_void) {
|
|
327
|
+
if ptr.is_null() {
|
|
328
|
+
return;
|
|
329
|
+
}
|
|
330
|
+
|
|
331
|
+
let _ = catch_unwind(AssertUnwindSafe(|| unsafe {
|
|
332
|
+
drop(Box::from_raw(ptr as *mut BlockCodec));
|
|
333
|
+
}));
|
|
334
|
+
}
|
|
335
|
+
|
|
336
|
+
unsafe extern "C" fn dict_trainer_free(ptr: *mut c_void) {
|
|
337
|
+
if ptr.is_null() {
|
|
338
|
+
return;
|
|
339
|
+
}
|
|
340
|
+
|
|
341
|
+
let _ = catch_unwind(AssertUnwindSafe(|| unsafe {
|
|
342
|
+
drop(Box::from_raw(ptr as *mut RbDictTrainer));
|
|
343
|
+
}));
|
|
344
|
+
}
|
|
345
|
+
|
|
346
|
+
unsafe extern "C" fn frame_codec_native_size(_ptr: *const c_void) -> size_t {
|
|
347
|
+
std::mem::size_of::<FrameCodec>() as size_t
|
|
348
|
+
}
|
|
349
|
+
|
|
350
|
+
unsafe extern "C" fn block_codec_native_size(_ptr: *const c_void) -> size_t {
|
|
351
|
+
std::mem::size_of::<BlockCodec>() as size_t
|
|
352
|
+
}
|
|
353
|
+
|
|
354
|
+
unsafe extern "C" fn dict_trainer_native_size(_ptr: *const c_void) -> size_t {
|
|
355
|
+
std::mem::size_of::<RbDictTrainer>() as size_t
|
|
356
|
+
}
|
|
357
|
+
|
|
358
|
+
unsafe fn frame_codec_ref(value: VALUE) -> RbResult<&'static FrameCodec> {
|
|
359
|
+
unsafe { rb::typed_data_ref(value, frame_codec_data_type(), "Zrip::FrameCodec") }
|
|
360
|
+
}
|
|
361
|
+
|
|
362
|
+
unsafe fn block_codec_ref(value: VALUE) -> RbResult<&'static BlockCodec> {
|
|
363
|
+
unsafe { rb::typed_data_ref(value, block_codec_data_type(), "Zrip::BlockCodec") }
|
|
364
|
+
}
|
|
365
|
+
|
|
366
|
+
unsafe fn dict_trainer_ref(value: VALUE) -> RbResult<&'static RbDictTrainer> {
|
|
367
|
+
unsafe { rb::typed_data_ref(value, dict_trainer_data_type(), "Zrip::DictTrainer") }
|
|
368
|
+
}
|
|
369
|
+
|
|
172
370
|
// ---------- FrameCodec ----------
|
|
173
371
|
|
|
174
|
-
#[magnus::wrap(class = "Zrip::FrameCodec", free_immediately, size)]
|
|
175
372
|
struct FrameCodec {
|
|
176
373
|
dict_len: usize,
|
|
177
374
|
dict_id: Option<u32>,
|
|
375
|
+
dict: Option<Dictionary>,
|
|
178
376
|
level: i32,
|
|
179
377
|
cctx: Mutex<CompressContext>,
|
|
180
378
|
dctx: Mutex<DecompressContext>,
|
|
@@ -183,198 +381,303 @@ struct FrameCodec {
|
|
|
183
381
|
unsafe impl Send for FrameCodec {}
|
|
184
382
|
unsafe impl Sync for FrameCodec {}
|
|
185
383
|
|
|
186
|
-
fn
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
)
|
|
192
|
-
|
|
384
|
+
fn frame_codec_new_impl(class: VALUE, rb_dict: VALUE, id: VALUE, level: VALUE) -> RbResult<VALUE> {
|
|
385
|
+
let rb_dict = if rb_dict == rb::qnil() {
|
|
386
|
+
None
|
|
387
|
+
} else {
|
|
388
|
+
let rb_dict = rb::string_value(rb_dict)?;
|
|
389
|
+
rb::freeze_value(rb_dict)?;
|
|
390
|
+
Some(rb::value_to_bytes(rb_dict)?)
|
|
391
|
+
};
|
|
392
|
+
let id = rb::value_to_u32(id)?;
|
|
393
|
+
let level = rb::value_to_i32(level)?;
|
|
394
|
+
let (dict_len, dict_id, dict, cctx, dctx) = match rb_dict {
|
|
193
395
|
None => {
|
|
194
396
|
let cctx = CompressContext::new(level).map_err(|e| {
|
|
195
|
-
|
|
196
|
-
compress_error(
|
|
397
|
+
RubyErr::new(
|
|
398
|
+
compress_error(),
|
|
197
399
|
format!("CompressContext::new failed: {e}"),
|
|
198
400
|
)
|
|
199
401
|
})?;
|
|
200
|
-
(0, None, cctx, DecompressContext::new())
|
|
402
|
+
(0, None, None, cctx, DecompressContext::new())
|
|
201
403
|
}
|
|
202
|
-
Some(
|
|
203
|
-
let
|
|
204
|
-
s.freeze();
|
|
205
|
-
let dict = load_dict(ruby, &bytes)?;
|
|
404
|
+
Some(bytes) => {
|
|
405
|
+
let dict = load_dict(&bytes)?;
|
|
206
406
|
let dl = bytes.len();
|
|
207
407
|
let cctx = CompressContext::with_dict(level, dict.clone()).map_err(|e| {
|
|
208
|
-
|
|
209
|
-
compress_error(
|
|
408
|
+
RubyErr::new(
|
|
409
|
+
compress_error(),
|
|
210
410
|
format!("CompressContext::with_dict failed: {e}"),
|
|
211
411
|
)
|
|
212
412
|
})?;
|
|
213
|
-
let dctx = DecompressContext::with_dict(dict);
|
|
214
|
-
(dl, Some(id), cctx, dctx)
|
|
413
|
+
let dctx = DecompressContext::with_dict(dict.clone());
|
|
414
|
+
(dl, Some(id), Some(dict), cctx, dctx)
|
|
215
415
|
}
|
|
216
416
|
};
|
|
217
417
|
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
let input: &[u8] = unsafe { rb_input.as_slice() };
|
|
233
|
-
let mut cctx = rb_self.cctx.lock().expect("FrameCodec CCtx mutex poisoned");
|
|
234
|
-
let out = cctx
|
|
235
|
-
.compress(input)
|
|
236
|
-
.map_err(|e| Error::new(compress_error(ruby), format!("zstd compress failed: {e}")))?;
|
|
237
|
-
Ok(ruby.str_from_slice(&out))
|
|
418
|
+
unsafe {
|
|
419
|
+
rb::wrap_typed_data(
|
|
420
|
+
class,
|
|
421
|
+
Box::new(FrameCodec {
|
|
422
|
+
dict_len,
|
|
423
|
+
dict_id,
|
|
424
|
+
dict,
|
|
425
|
+
level,
|
|
426
|
+
cctx: Mutex::new(cctx),
|
|
427
|
+
dctx: Mutex::new(dctx),
|
|
428
|
+
}),
|
|
429
|
+
frame_codec_data_type(),
|
|
430
|
+
)
|
|
431
|
+
}
|
|
238
432
|
}
|
|
239
433
|
|
|
240
|
-
fn
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
434
|
+
fn frame_codec_compress_impl(rb_self: VALUE, rb_input: VALUE) -> RbResult<VALUE> {
|
|
435
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
436
|
+
let mut input = rb::input_bytes(rb_input)?;
|
|
437
|
+
let release_gvl = should_release_compress_gvl(input.len());
|
|
438
|
+
input.lock_for_without_gvl(release_gvl)?;
|
|
439
|
+
let out = with_mutex(&rb_self.cctx, release_gvl, "FrameCodec CCtx", |cctx| {
|
|
440
|
+
cctx.compress(input.as_slice())
|
|
441
|
+
.map(|out| out.into_owned())
|
|
442
|
+
.map_err(|e| RubyErr::new(compress_error(), format!("zstd compress failed: {e}")))
|
|
443
|
+
})?;
|
|
444
|
+
rb::new_binary_string(&out)
|
|
445
|
+
}
|
|
446
|
+
|
|
447
|
+
fn frame_codec_decompress_impl(
|
|
448
|
+
rb_self: VALUE,
|
|
449
|
+
rb_input: VALUE,
|
|
450
|
+
max_output: VALUE,
|
|
451
|
+
) -> RbResult<VALUE> {
|
|
452
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
453
|
+
let mut compressed = rb::input_bytes(rb_input)?;
|
|
454
|
+
let max_output = rb::value_to_usize(max_output)?;
|
|
455
|
+
let release_gvl = should_release_frame_decompress_gvl(decompress_work_size(
|
|
456
|
+
compressed.as_slice(),
|
|
457
|
+
max_output,
|
|
458
|
+
));
|
|
459
|
+
compressed.lock_for_without_gvl(release_gvl)?;
|
|
460
|
+
let dict = rb_self.dict.clone();
|
|
461
|
+
let out = with_mutex(&rb_self.dctx, release_gvl, "FrameCodec DCtx", |dctx| {
|
|
462
|
+
decompress_bounded(compressed.as_slice(), max_output, dctx, dict.as_ref())
|
|
463
|
+
.map_err(|e| bounded_err(e, "zstd frame decode failed"))
|
|
464
|
+
})?;
|
|
465
|
+
rb::new_binary_string(&out)
|
|
251
466
|
}
|
|
252
467
|
|
|
253
|
-
fn
|
|
254
|
-
rb_self
|
|
468
|
+
fn frame_codec_size_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
469
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
470
|
+
rb::usize_value(rb_self.dict_len)
|
|
255
471
|
}
|
|
256
472
|
|
|
257
|
-
fn
|
|
258
|
-
rb_self
|
|
473
|
+
fn frame_codec_has_dict_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
474
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
475
|
+
Ok(rb::bool_value(rb_self.dict_id.is_some()))
|
|
259
476
|
}
|
|
260
477
|
|
|
261
|
-
fn
|
|
262
|
-
rb_self
|
|
478
|
+
fn frame_codec_id_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
479
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
480
|
+
rb::u32_option_value(rb_self.dict_id)
|
|
263
481
|
}
|
|
264
482
|
|
|
265
|
-
fn
|
|
266
|
-
rb_self
|
|
483
|
+
fn frame_codec_level_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
484
|
+
let rb_self = unsafe { frame_codec_ref(rb_self)? };
|
|
485
|
+
Ok(rb::i32_value(rb_self.level))
|
|
267
486
|
}
|
|
268
487
|
|
|
269
|
-
fn
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
)
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
Ok(v) => Ok(v),
|
|
276
|
-
Err(BoundedError::BadMagic) => Err(Error::new(
|
|
277
|
-
decompress_error(ruby),
|
|
488
|
+
fn frame_codec_get_frame_content_size_impl(rb_input: VALUE) -> RbResult<VALUE> {
|
|
489
|
+
let bytes = rb::input_bytes(rb_input)?;
|
|
490
|
+
match parse_frame_content_size(bytes.as_slice()) {
|
|
491
|
+
Ok(v) => rb::u64_option_value(v),
|
|
492
|
+
Err(BoundedError::BadMagic) => Err(RubyErr::new(
|
|
493
|
+
decompress_error(),
|
|
278
494
|
"zstd frame header parse failed: bad magic (input is not a Zstd frame)",
|
|
279
495
|
)),
|
|
280
|
-
Err(e) => Err(
|
|
496
|
+
Err(e) => Err(bounded_err(e, "zstd frame header parse failed")),
|
|
281
497
|
}
|
|
282
498
|
}
|
|
283
499
|
|
|
500
|
+
unsafe extern "C" fn frame_codec_new(
|
|
501
|
+
class: VALUE,
|
|
502
|
+
rb_dict: VALUE,
|
|
503
|
+
id: VALUE,
|
|
504
|
+
level: VALUE,
|
|
505
|
+
) -> VALUE {
|
|
506
|
+
rb::wrap(|| frame_codec_new_impl(class, rb_dict, id, level))
|
|
507
|
+
}
|
|
508
|
+
|
|
509
|
+
unsafe extern "C" fn frame_codec_compress(rb_self: VALUE, rb_input: VALUE) -> VALUE {
|
|
510
|
+
rb::wrap(|| frame_codec_compress_impl(rb_self, rb_input))
|
|
511
|
+
}
|
|
512
|
+
|
|
513
|
+
unsafe extern "C" fn frame_codec_decompress(
|
|
514
|
+
rb_self: VALUE,
|
|
515
|
+
rb_input: VALUE,
|
|
516
|
+
max_output: VALUE,
|
|
517
|
+
) -> VALUE {
|
|
518
|
+
rb::wrap(|| frame_codec_decompress_impl(rb_self, rb_input, max_output))
|
|
519
|
+
}
|
|
520
|
+
|
|
521
|
+
unsafe extern "C" fn frame_codec_size(rb_self: VALUE) -> VALUE {
|
|
522
|
+
rb::wrap(|| frame_codec_size_impl(rb_self))
|
|
523
|
+
}
|
|
524
|
+
|
|
525
|
+
unsafe extern "C" fn frame_codec_has_dict(rb_self: VALUE) -> VALUE {
|
|
526
|
+
rb::wrap(|| frame_codec_has_dict_impl(rb_self))
|
|
527
|
+
}
|
|
528
|
+
|
|
529
|
+
unsafe extern "C" fn frame_codec_id(rb_self: VALUE) -> VALUE {
|
|
530
|
+
rb::wrap(|| frame_codec_id_impl(rb_self))
|
|
531
|
+
}
|
|
532
|
+
|
|
533
|
+
unsafe extern "C" fn frame_codec_level(rb_self: VALUE) -> VALUE {
|
|
534
|
+
rb::wrap(|| frame_codec_level_impl(rb_self))
|
|
535
|
+
}
|
|
536
|
+
|
|
537
|
+
unsafe extern "C" fn frame_codec_get_frame_content_size(_class: VALUE, rb_input: VALUE) -> VALUE {
|
|
538
|
+
rb::wrap(|| frame_codec_get_frame_content_size_impl(rb_input))
|
|
539
|
+
}
|
|
540
|
+
|
|
284
541
|
// ---------- BlockCodec ----------
|
|
285
542
|
|
|
286
|
-
#[magnus::wrap(class = "Zrip::BlockCodec", free_immediately, size)]
|
|
287
543
|
struct BlockCodec {
|
|
288
544
|
dict_len: usize,
|
|
289
545
|
dict_id: Option<u32>,
|
|
546
|
+
dict: Option<Dictionary>,
|
|
290
547
|
level: i32,
|
|
291
|
-
cctx:
|
|
292
|
-
dctx:
|
|
548
|
+
cctx: Mutex<CompressContext>,
|
|
549
|
+
dctx: Mutex<DecompressContext>,
|
|
293
550
|
}
|
|
294
551
|
|
|
295
|
-
fn
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
) -> Result<BlockCodec, Error> {
|
|
301
|
-
let (dict_len, dict_id, cctx, dctx) = match rb_dict {
|
|
552
|
+
fn block_codec_new_impl(class: VALUE, rb_dict: VALUE, id: VALUE, level: VALUE) -> RbResult<VALUE> {
|
|
553
|
+
let rb_dict = rb::value_to_option_bytes(rb_dict)?;
|
|
554
|
+
let id = rb::value_to_u32(id)?;
|
|
555
|
+
let level = rb::value_to_i32(level)?;
|
|
556
|
+
let (dict_len, dict_id, dict, cctx, dctx) = match rb_dict {
|
|
302
557
|
None => {
|
|
303
558
|
let cctx = CompressContext::new(level).map_err(|e| {
|
|
304
|
-
|
|
305
|
-
compress_error(
|
|
559
|
+
RubyErr::new(
|
|
560
|
+
compress_error(),
|
|
306
561
|
format!("CompressContext::new failed: {e}"),
|
|
307
562
|
)
|
|
308
563
|
})?;
|
|
309
|
-
(0, None, cctx, DecompressContext::new())
|
|
564
|
+
(0, None, None, cctx, DecompressContext::new())
|
|
310
565
|
}
|
|
311
|
-
Some(
|
|
312
|
-
let
|
|
313
|
-
let dict = load_dict(ruby, &bytes)?;
|
|
566
|
+
Some(bytes) => {
|
|
567
|
+
let dict = load_dict(&bytes)?;
|
|
314
568
|
let dl = bytes.len();
|
|
315
569
|
let cctx = CompressContext::with_dict(level, dict.clone()).map_err(|e| {
|
|
316
|
-
|
|
317
|
-
compress_error(
|
|
570
|
+
RubyErr::new(
|
|
571
|
+
compress_error(),
|
|
318
572
|
format!("CompressContext::with_dict failed: {e}"),
|
|
319
573
|
)
|
|
320
574
|
})?;
|
|
321
|
-
let dctx = DecompressContext::with_dict(dict);
|
|
322
|
-
(dl, Some(id), cctx, dctx)
|
|
575
|
+
let dctx = DecompressContext::with_dict(dict.clone());
|
|
576
|
+
(dl, Some(id), Some(dict), cctx, dctx)
|
|
323
577
|
}
|
|
324
578
|
};
|
|
325
579
|
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
580
|
+
unsafe {
|
|
581
|
+
rb::wrap_typed_data(
|
|
582
|
+
class,
|
|
583
|
+
Box::new(BlockCodec {
|
|
584
|
+
dict_len,
|
|
585
|
+
dict_id,
|
|
586
|
+
dict,
|
|
587
|
+
level,
|
|
588
|
+
cctx: Mutex::new(cctx),
|
|
589
|
+
dctx: Mutex::new(dctx),
|
|
590
|
+
}),
|
|
591
|
+
block_codec_data_type(),
|
|
592
|
+
)
|
|
593
|
+
}
|
|
594
|
+
}
|
|
595
|
+
|
|
596
|
+
fn block_codec_compress_impl(rb_self: VALUE, rb_input: VALUE) -> RbResult<VALUE> {
|
|
597
|
+
let rb_self = unsafe { block_codec_ref(rb_self)? };
|
|
598
|
+
let mut input = rb::input_bytes(rb_input)?;
|
|
599
|
+
let release_gvl = should_release_compress_gvl(input.len());
|
|
600
|
+
input.lock_for_without_gvl(release_gvl)?;
|
|
601
|
+
let out = with_mutex(&rb_self.cctx, release_gvl, "BlockCodec CCtx", |cctx| {
|
|
602
|
+
cctx.compress(input.as_slice())
|
|
603
|
+
.map(|out| out.into_owned())
|
|
604
|
+
.map_err(|e| RubyErr::new(compress_error(), format!("zstd compress failed: {e}")))
|
|
605
|
+
})?;
|
|
606
|
+
rb::new_binary_string(&out)
|
|
607
|
+
}
|
|
608
|
+
|
|
609
|
+
fn block_codec_decompress_impl(
|
|
610
|
+
rb_self: VALUE,
|
|
611
|
+
rb_input: VALUE,
|
|
612
|
+
max_output: VALUE,
|
|
613
|
+
) -> RbResult<VALUE> {
|
|
614
|
+
let rb_self = unsafe { block_codec_ref(rb_self)? };
|
|
615
|
+
let compressed = rb::input_bytes(rb_input)?;
|
|
616
|
+
let max_output = rb::value_to_usize(max_output)?;
|
|
617
|
+
let out = with_mutex(&rb_self.dctx, false, "BlockCodec DCtx", |dctx| {
|
|
618
|
+
decompress_bounded(
|
|
619
|
+
compressed.as_slice(),
|
|
620
|
+
max_output,
|
|
621
|
+
dctx,
|
|
622
|
+
rb_self.dict.as_ref(),
|
|
623
|
+
)
|
|
624
|
+
.map_err(|e| bounded_err(e, "zstd block decode failed"))
|
|
625
|
+
})?;
|
|
626
|
+
rb::new_binary_string(&out)
|
|
333
627
|
}
|
|
334
628
|
|
|
335
|
-
fn
|
|
336
|
-
|
|
337
|
-
rb_self
|
|
338
|
-
rb_input: RString,
|
|
339
|
-
) -> Result<RString, Error> {
|
|
340
|
-
let input: &[u8] = unsafe { rb_input.as_slice() };
|
|
341
|
-
let mut cctx = rb_self.cctx.borrow_mut();
|
|
342
|
-
let out = cctx
|
|
343
|
-
.compress(input)
|
|
344
|
-
.map_err(|e| Error::new(compress_error(ruby), format!("zstd compress failed: {e}")))?;
|
|
345
|
-
Ok(ruby.str_from_slice(&out))
|
|
629
|
+
fn block_codec_size_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
630
|
+
let rb_self = unsafe { block_codec_ref(rb_self)? };
|
|
631
|
+
rb::usize_value(rb_self.dict_len)
|
|
346
632
|
}
|
|
347
633
|
|
|
348
|
-
fn
|
|
349
|
-
|
|
350
|
-
rb_self
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
) ->
|
|
354
|
-
let
|
|
355
|
-
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
634
|
+
fn block_codec_has_dict_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
635
|
+
let rb_self = unsafe { block_codec_ref(rb_self)? };
|
|
636
|
+
Ok(rb::bool_value(rb_self.dict_id.is_some()))
|
|
637
|
+
}
|
|
638
|
+
|
|
639
|
+
fn block_codec_level_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
640
|
+
let rb_self = unsafe { block_codec_ref(rb_self)? };
|
|
641
|
+
Ok(rb::i32_value(rb_self.level))
|
|
642
|
+
}
|
|
643
|
+
|
|
644
|
+
unsafe extern "C" fn block_codec_new(
|
|
645
|
+
class: VALUE,
|
|
646
|
+
rb_dict: VALUE,
|
|
647
|
+
id: VALUE,
|
|
648
|
+
level: VALUE,
|
|
649
|
+
) -> VALUE {
|
|
650
|
+
rb::wrap(|| block_codec_new_impl(class, rb_dict, id, level))
|
|
651
|
+
}
|
|
652
|
+
|
|
653
|
+
unsafe extern "C" fn block_codec_compress(rb_self: VALUE, rb_input: VALUE) -> VALUE {
|
|
654
|
+
rb::wrap(|| block_codec_compress_impl(rb_self, rb_input))
|
|
655
|
+
}
|
|
656
|
+
|
|
657
|
+
unsafe extern "C" fn block_codec_decompress(
|
|
658
|
+
rb_self: VALUE,
|
|
659
|
+
rb_input: VALUE,
|
|
660
|
+
max_output: VALUE,
|
|
661
|
+
) -> VALUE {
|
|
662
|
+
rb::wrap(|| block_codec_decompress_impl(rb_self, rb_input, max_output))
|
|
359
663
|
}
|
|
360
664
|
|
|
361
|
-
fn block_codec_size(rb_self:
|
|
362
|
-
rb_self
|
|
665
|
+
unsafe extern "C" fn block_codec_size(rb_self: VALUE) -> VALUE {
|
|
666
|
+
rb::wrap(|| block_codec_size_impl(rb_self))
|
|
363
667
|
}
|
|
364
668
|
|
|
365
|
-
fn block_codec_has_dict(rb_self:
|
|
366
|
-
rb_self
|
|
669
|
+
unsafe extern "C" fn block_codec_has_dict(rb_self: VALUE) -> VALUE {
|
|
670
|
+
rb::wrap(|| block_codec_has_dict_impl(rb_self))
|
|
367
671
|
}
|
|
368
672
|
|
|
369
|
-
fn block_codec_level(rb_self:
|
|
370
|
-
rb_self
|
|
673
|
+
unsafe extern "C" fn block_codec_level(rb_self: VALUE) -> VALUE {
|
|
674
|
+
rb::wrap(|| block_codec_level_impl(rb_self))
|
|
371
675
|
}
|
|
372
676
|
|
|
373
677
|
// ---------- DictTrainer ----------
|
|
374
678
|
|
|
375
|
-
#[magnus::wrap(class = "Zrip::DictTrainer", free_immediately, size)]
|
|
376
679
|
struct RbDictTrainer {
|
|
377
|
-
inner:
|
|
680
|
+
inner: Mutex<Option<TrainerState>>,
|
|
378
681
|
max_dict_size: usize,
|
|
379
682
|
}
|
|
380
683
|
|
|
@@ -383,67 +686,78 @@ struct TrainerState {
|
|
|
383
686
|
total_bytes: usize,
|
|
384
687
|
}
|
|
385
688
|
|
|
386
|
-
fn
|
|
387
|
-
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
689
|
+
fn dict_trainer_new_impl(class: VALUE, max_dict_size: VALUE) -> RbResult<VALUE> {
|
|
690
|
+
let max_dict_size = rb::value_to_usize(max_dict_size)?;
|
|
691
|
+
unsafe {
|
|
692
|
+
rb::wrap_typed_data(
|
|
693
|
+
class,
|
|
694
|
+
Box::new(RbDictTrainer {
|
|
695
|
+
max_dict_size,
|
|
696
|
+
inner: Mutex::new(Some(TrainerState {
|
|
697
|
+
samples: Vec::new(),
|
|
698
|
+
total_bytes: 0,
|
|
699
|
+
})),
|
|
700
|
+
}),
|
|
701
|
+
dict_trainer_data_type(),
|
|
702
|
+
)
|
|
393
703
|
}
|
|
394
704
|
}
|
|
395
705
|
|
|
396
|
-
fn
|
|
397
|
-
|
|
398
|
-
rb_self
|
|
399
|
-
|
|
400
|
-
|
|
401
|
-
|
|
402
|
-
let state = borrow
|
|
403
|
-
|
|
404
|
-
|
|
405
|
-
|
|
406
|
-
)
|
|
407
|
-
})?;
|
|
408
|
-
let data: Vec<u8> = unsafe { rb_data.as_slice().to_vec() };
|
|
706
|
+
fn dict_trainer_add_sample_impl(rb_self: VALUE, rb_data: VALUE) -> RbResult<VALUE> {
|
|
707
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
708
|
+
let mut borrow = rb_self
|
|
709
|
+
.inner
|
|
710
|
+
.lock()
|
|
711
|
+
.map_err(|_| RubyErr::runtime("DictTrainer mutex poisoned"))?;
|
|
712
|
+
let state = borrow
|
|
713
|
+
.as_mut()
|
|
714
|
+
.ok_or_else(|| RubyErr::runtime("DictTrainer already consumed by #train"))?;
|
|
715
|
+
let data = rb::value_to_bytes(rb_data)?;
|
|
409
716
|
if data.len() < 4 {
|
|
410
|
-
return Ok(());
|
|
717
|
+
return Ok(rb::qnil());
|
|
411
718
|
}
|
|
412
719
|
state.total_bytes += data.len();
|
|
413
720
|
state.samples.push(data);
|
|
414
|
-
Ok(())
|
|
415
|
-
}
|
|
416
|
-
|
|
417
|
-
fn
|
|
418
|
-
let
|
|
419
|
-
borrow
|
|
420
|
-
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
|
|
434
|
-
|
|
435
|
-
|
|
436
|
-
|
|
437
|
-
|
|
438
|
-
|
|
439
|
-
|
|
440
|
-
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
|
|
721
|
+
Ok(rb::qnil())
|
|
722
|
+
}
|
|
723
|
+
|
|
724
|
+
fn dict_trainer_sample_count_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
725
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
726
|
+
let borrow = rb_self
|
|
727
|
+
.inner
|
|
728
|
+
.lock()
|
|
729
|
+
.map_err(|_| RubyErr::runtime("DictTrainer mutex poisoned"))?;
|
|
730
|
+
let value = borrow
|
|
731
|
+
.as_ref()
|
|
732
|
+
.map(|s| s.samples.len())
|
|
733
|
+
.ok_or_else(|| RubyErr::runtime("DictTrainer already consumed by #train"))?;
|
|
734
|
+
rb::usize_value(value)
|
|
735
|
+
}
|
|
736
|
+
|
|
737
|
+
fn dict_trainer_total_bytes_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
738
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
739
|
+
let borrow = rb_self
|
|
740
|
+
.inner
|
|
741
|
+
.lock()
|
|
742
|
+
.map_err(|_| RubyErr::runtime("DictTrainer mutex poisoned"))?;
|
|
743
|
+
let value = borrow
|
|
744
|
+
.as_ref()
|
|
745
|
+
.map(|s| s.total_bytes)
|
|
746
|
+
.ok_or_else(|| RubyErr::runtime("DictTrainer already consumed by #train"))?;
|
|
747
|
+
rb::usize_value(value)
|
|
748
|
+
}
|
|
749
|
+
|
|
750
|
+
fn dict_trainer_train_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
751
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
752
|
+
let state = rb_self
|
|
753
|
+
.inner
|
|
754
|
+
.lock()
|
|
755
|
+
.map_err(|_| RubyErr::runtime("DictTrainer mutex poisoned"))?
|
|
756
|
+
.take()
|
|
757
|
+
.ok_or_else(|| RubyErr::runtime("DictTrainer already consumed by #train"))?;
|
|
444
758
|
|
|
445
759
|
if state.samples.len() < 2 {
|
|
446
|
-
return
|
|
760
|
+
return rb::new_binary_string(b"");
|
|
447
761
|
}
|
|
448
762
|
|
|
449
763
|
let refs: Vec<&[u8]> = state.samples.iter().map(|s| s.as_slice()).collect();
|
|
@@ -456,81 +770,146 @@ fn dict_trainer_train(ruby: &Ruby, rb_self: &RbDictTrainer) -> Result<RString, E
|
|
|
456
770
|
let dict_bytes =
|
|
457
771
|
zstd::dict::finalize::finalize_dictionary(&content, &refs, rb_self.max_dict_size);
|
|
458
772
|
|
|
459
|
-
|
|
773
|
+
rb::new_binary_string(&dict_bytes)
|
|
774
|
+
}
|
|
775
|
+
|
|
776
|
+
fn dict_trainer_max_dict_size_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
777
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
778
|
+
rb::usize_value(rb_self.max_dict_size)
|
|
779
|
+
}
|
|
780
|
+
|
|
781
|
+
fn dict_trainer_trained_impl(rb_self: VALUE) -> RbResult<VALUE> {
|
|
782
|
+
let rb_self = unsafe { dict_trainer_ref(rb_self)? };
|
|
783
|
+
let borrow = rb_self
|
|
784
|
+
.inner
|
|
785
|
+
.lock()
|
|
786
|
+
.map_err(|_| RubyErr::runtime("DictTrainer mutex poisoned"))?;
|
|
787
|
+
Ok(rb::bool_value(borrow.is_none()))
|
|
788
|
+
}
|
|
789
|
+
|
|
790
|
+
unsafe extern "C" fn dict_trainer_new(class: VALUE, max_dict_size: VALUE) -> VALUE {
|
|
791
|
+
rb::wrap(|| dict_trainer_new_impl(class, max_dict_size))
|
|
460
792
|
}
|
|
461
793
|
|
|
462
|
-
fn
|
|
463
|
-
rb_self
|
|
794
|
+
unsafe extern "C" fn dict_trainer_add_sample(rb_self: VALUE, rb_data: VALUE) -> VALUE {
|
|
795
|
+
rb::wrap(|| dict_trainer_add_sample_impl(rb_self, rb_data))
|
|
464
796
|
}
|
|
465
797
|
|
|
466
|
-
fn
|
|
467
|
-
|
|
798
|
+
unsafe extern "C" fn dict_trainer_sample_count(rb_self: VALUE) -> VALUE {
|
|
799
|
+
rb::wrap(|| dict_trainer_sample_count_impl(rb_self))
|
|
800
|
+
}
|
|
801
|
+
|
|
802
|
+
unsafe extern "C" fn dict_trainer_total_bytes(rb_self: VALUE) -> VALUE {
|
|
803
|
+
rb::wrap(|| dict_trainer_total_bytes_impl(rb_self))
|
|
804
|
+
}
|
|
805
|
+
|
|
806
|
+
unsafe extern "C" fn dict_trainer_train(rb_self: VALUE) -> VALUE {
|
|
807
|
+
rb::wrap(|| dict_trainer_train_impl(rb_self))
|
|
808
|
+
}
|
|
809
|
+
|
|
810
|
+
unsafe extern "C" fn dict_trainer_max_dict_size(rb_self: VALUE) -> VALUE {
|
|
811
|
+
rb::wrap(|| dict_trainer_max_dict_size_impl(rb_self))
|
|
812
|
+
}
|
|
813
|
+
|
|
814
|
+
unsafe extern "C" fn dict_trainer_trained(rb_self: VALUE) -> VALUE {
|
|
815
|
+
rb::wrap(|| dict_trainer_trained_impl(rb_self))
|
|
468
816
|
}
|
|
469
817
|
|
|
470
818
|
// ---------- module init ----------
|
|
471
819
|
|
|
472
|
-
#
|
|
473
|
-
|
|
474
|
-
|
|
820
|
+
/// # Safety
|
|
821
|
+
///
|
|
822
|
+
/// Ruby calls this function while loading the native extension. The Ruby VM
|
|
823
|
+
/// must be initialized, and the symbol must only be entered by Ruby's extension
|
|
824
|
+
/// loader.
|
|
825
|
+
#[no_mangle]
|
|
826
|
+
pub unsafe extern "C" fn Init_zrip() {
|
|
827
|
+
rb::wrap_init(init);
|
|
828
|
+
}
|
|
829
|
+
|
|
830
|
+
fn init() -> RbResult<()> {
|
|
831
|
+
#[cfg(ruby_engine = "mri")]
|
|
832
|
+
unsafe {
|
|
833
|
+
rb_sys::rb_ext_ractor_safe(true);
|
|
834
|
+
}
|
|
475
835
|
|
|
476
|
-
let module =
|
|
836
|
+
let module = unsafe { rb::define_module(c"Zrip")? };
|
|
477
837
|
|
|
478
838
|
let decompress_error_class =
|
|
479
|
-
module
|
|
839
|
+
unsafe { rb::define_error_under(module, c"DecompressError", rb_sys::rb_eStandardError)? };
|
|
480
840
|
DECOMPRESS_ERROR
|
|
481
|
-
.set(
|
|
841
|
+
.set(GlobalValue(decompress_error_class))
|
|
482
842
|
.unwrap_or_else(|_| panic!("init called more than once"));
|
|
483
843
|
|
|
484
844
|
let compress_error_class =
|
|
485
|
-
module
|
|
845
|
+
unsafe { rb::define_error_under(module, c"CompressError", rb_sys::rb_eStandardError)? };
|
|
486
846
|
COMPRESS_ERROR
|
|
487
|
-
.set(
|
|
847
|
+
.set(GlobalValue(compress_error_class))
|
|
488
848
|
.unwrap_or_else(|_| panic!("init called more than once"));
|
|
489
849
|
|
|
490
|
-
let missing_content_size_error_class =
|
|
491
|
-
module
|
|
850
|
+
let missing_content_size_error_class = unsafe {
|
|
851
|
+
rb::define_error_under(module, c"MissingContentSizeError", decompress_error_class)?
|
|
852
|
+
};
|
|
492
853
|
MISSING_CONTENT_SIZE_ERROR
|
|
493
|
-
.set(
|
|
854
|
+
.set(GlobalValue(missing_content_size_error_class))
|
|
494
855
|
.unwrap_or_else(|_| panic!("init called more than once"));
|
|
495
856
|
|
|
496
857
|
let output_size_limit_error_class =
|
|
497
|
-
module
|
|
858
|
+
unsafe { rb::define_error_under(module, c"OutputSizeLimitError", decompress_error_class)? };
|
|
498
859
|
OUTPUT_SIZE_LIMIT_ERROR
|
|
499
|
-
.set(
|
|
860
|
+
.set(GlobalValue(output_size_limit_error_class))
|
|
500
861
|
.unwrap_or_else(|_| panic!("init called more than once"));
|
|
501
862
|
|
|
502
|
-
|
|
503
|
-
|
|
504
|
-
|
|
505
|
-
|
|
506
|
-
"
|
|
507
|
-
|
|
508
|
-
|
|
509
|
-
|
|
510
|
-
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
|
|
514
|
-
|
|
515
|
-
|
|
516
|
-
|
|
517
|
-
|
|
518
|
-
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
|
|
529
|
-
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
863
|
+
let frame_codec_class =
|
|
864
|
+
unsafe { rb::define_class_under(module, c"FrameCodec", rb_sys::rb_cObject)? };
|
|
865
|
+
unsafe {
|
|
866
|
+
rb::undef_alloc_func(frame_codec_class)?;
|
|
867
|
+
rb::define_singleton_method_3(frame_codec_class, c"_native_new", frame_codec_new)?;
|
|
868
|
+
rb::define_singleton_method_1(
|
|
869
|
+
frame_codec_class,
|
|
870
|
+
c"get_frame_content_size",
|
|
871
|
+
frame_codec_get_frame_content_size,
|
|
872
|
+
)?;
|
|
873
|
+
rb::define_method_1(frame_codec_class, c"compress", frame_codec_compress)?;
|
|
874
|
+
rb::define_method_2(
|
|
875
|
+
frame_codec_class,
|
|
876
|
+
c"_native_decompress",
|
|
877
|
+
frame_codec_decompress,
|
|
878
|
+
)?;
|
|
879
|
+
rb::define_method_0(frame_codec_class, c"size", frame_codec_size)?;
|
|
880
|
+
rb::define_method_0(frame_codec_class, c"has_dict?", frame_codec_has_dict)?;
|
|
881
|
+
rb::define_method_0(frame_codec_class, c"id", frame_codec_id)?;
|
|
882
|
+
rb::define_method_0(frame_codec_class, c"level", frame_codec_level)?;
|
|
883
|
+
}
|
|
884
|
+
|
|
885
|
+
let block_codec_class =
|
|
886
|
+
unsafe { rb::define_class_under(module, c"BlockCodec", rb_sys::rb_cObject)? };
|
|
887
|
+
unsafe {
|
|
888
|
+
rb::undef_alloc_func(block_codec_class)?;
|
|
889
|
+
rb::define_singleton_method_3(block_codec_class, c"_native_new", block_codec_new)?;
|
|
890
|
+
rb::define_method_1(block_codec_class, c"compress", block_codec_compress)?;
|
|
891
|
+
rb::define_method_2(
|
|
892
|
+
block_codec_class,
|
|
893
|
+
c"_native_decompress",
|
|
894
|
+
block_codec_decompress,
|
|
895
|
+
)?;
|
|
896
|
+
rb::define_method_0(block_codec_class, c"size", block_codec_size)?;
|
|
897
|
+
rb::define_method_0(block_codec_class, c"has_dict?", block_codec_has_dict)?;
|
|
898
|
+
rb::define_method_0(block_codec_class, c"level", block_codec_level)?;
|
|
899
|
+
}
|
|
900
|
+
|
|
901
|
+
let trainer_class =
|
|
902
|
+
unsafe { rb::define_class_under(module, c"DictTrainer", rb_sys::rb_cObject)? };
|
|
903
|
+
unsafe {
|
|
904
|
+
rb::undef_alloc_func(trainer_class)?;
|
|
905
|
+
rb::define_singleton_method_1(trainer_class, c"_native_new", dict_trainer_new)?;
|
|
906
|
+
rb::define_method_1(trainer_class, c"add_sample", dict_trainer_add_sample)?;
|
|
907
|
+
rb::define_method_0(trainer_class, c"sample_count", dict_trainer_sample_count)?;
|
|
908
|
+
rb::define_method_0(trainer_class, c"total_bytes", dict_trainer_total_bytes)?;
|
|
909
|
+
rb::define_method_0(trainer_class, c"train", dict_trainer_train)?;
|
|
910
|
+
rb::define_method_0(trainer_class, c"max_dict_size", dict_trainer_max_dict_size)?;
|
|
911
|
+
rb::define_method_0(trainer_class, c"trained?", dict_trainer_trained)?;
|
|
912
|
+
}
|
|
534
913
|
|
|
535
914
|
Ok(())
|
|
536
915
|
}
|
|
@@ -573,4 +952,22 @@ mod tests {
|
|
|
573
952
|
let fcs = parse_frame_content_size(&compressed).unwrap();
|
|
574
953
|
assert_eq!(fcs, Some(data.len() as u64));
|
|
575
954
|
}
|
|
955
|
+
|
|
956
|
+
#[test]
|
|
957
|
+
fn bounded_decompress_uses_total_concatenated_limit() {
|
|
958
|
+
let mut dctx = DecompressContext::new();
|
|
959
|
+
let mut compressed = zstd::compress(&[b'a'; 100], 1).unwrap();
|
|
960
|
+
compressed.extend_from_slice(&zstd::compress(&[b'b'; 100], 1).unwrap());
|
|
961
|
+
|
|
962
|
+
assert!(matches!(
|
|
963
|
+
decompress_bounded(&compressed, 199, &mut dctx, None),
|
|
964
|
+
Err(BoundedError::OutputSizeLimit { .. })
|
|
965
|
+
));
|
|
966
|
+
assert_eq!(
|
|
967
|
+
decompress_bounded(&compressed, 200, &mut dctx, None)
|
|
968
|
+
.unwrap()
|
|
969
|
+
.len(),
|
|
970
|
+
200
|
|
971
|
+
);
|
|
972
|
+
}
|
|
576
973
|
}
|