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.
data/ext/zrip/src/lib.rs CHANGED
@@ -1,44 +1,80 @@
1
- use magnus::{
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
- static DECOMPRESS_ERROR: OnceLock<Opaque<ExceptionClass>> = OnceLock::new();
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
- fn decompress_error(ruby: &Ruby) -> ExceptionClass {
17
- ruby.get_inner(
18
- *DECOMPRESS_ERROR
19
- .get()
20
- .expect("DecompressError not initialized"),
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 compress_error(ruby: &Ruby) -> ExceptionClass {
25
- ruby.get_inner(*COMPRESS_ERROR.get().expect("CompressError not initialized"))
35
+ fn decompress_error() -> VALUE {
36
+ stored_value(&DECOMPRESS_ERROR, "DecompressError")
26
37
  }
27
38
 
28
- fn missing_content_size_error(ruby: &Ruby) -> ExceptionClass {
29
- ruby.get_inner(
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(ruby: &Ruby) -> ExceptionClass {
37
- ruby.get_inner(
38
- *OUTPUT_SIZE_LIMIT_ERROR
39
- .get()
40
- .expect("OutputSizeLimitError not initialized"),
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
- MissingContentSize,
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 upper = match parse_frame_content_size(compressed)? {
111
- Some(n) => {
112
- if max_output != 0 && n > max_output as u64 {
113
- return Err(BoundedError::OutputSizeLimit {
114
- declared: n,
115
- limit: max_output as u64,
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
- 1024 * 1024
131
- }
199
+ None => zstd::FrameDecoder::with_limit(compressed, usize::MAX),
200
+ };
201
+ read_decoder_bounded(&mut decoder, max_output)?
132
202
  };
133
203
 
134
- let result = dctx
135
- .decompress_with_limit(compressed, upper)
136
- .map_err(|e| BoundedError::DecoderFailed(format!("{e}")))?;
204
+ Ok(result)
205
+ }
137
206
 
138
- Ok(result.into_owned())
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 raise_bounded(ruby: &Ruby, err: BoundedError, prefix: &str) -> Error {
215
+ fn bounded_err(err: BoundedError, prefix: &str) -> RubyErr {
142
216
  match err {
143
- BoundedError::BadMagic => Error::new(
144
- decompress_error(ruby),
217
+ BoundedError::BadMagic => RubyErr::new(
218
+ decompress_error(),
145
219
  format!("{prefix}: bad magic (input is not a Zstd frame)"),
146
220
  ),
147
- BoundedError::MissingContentSize => Error::new(
148
- missing_content_size_error(ruby),
149
- format!("{prefix}: Frame_Content_Size absent from frame header"),
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
- Error::new(decompress_error(ruby), format!("{prefix}: {msg}"))
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(ruby: &Ruby, bytes: &[u8]) -> Result<Dictionary, Error> {
233
+ fn load_dict(bytes: &[u8]) -> RbResult<Dictionary> {
164
234
  Dictionary::from_bytes(bytes).map_err(|_| {
165
- Error::new(
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 frame_codec_new(
187
- ruby: &Ruby,
188
- rb_dict: Option<RString>,
189
- id: u32,
190
- level: i32,
191
- ) -> Result<FrameCodec, Error> {
192
- let (dict_len, dict_id, cctx, dctx) = match rb_dict {
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
- Error::new(
196
- compress_error(ruby),
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(s) => {
203
- let bytes: Vec<u8> = unsafe { s.as_slice().to_vec() };
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
- Error::new(
209
- compress_error(ruby),
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
- Ok(FrameCodec {
219
- dict_len,
220
- dict_id,
221
- level,
222
- cctx: Mutex::new(cctx),
223
- dctx: Mutex::new(dctx),
224
- })
225
- }
226
-
227
- fn frame_codec_compress(
228
- ruby: &Ruby,
229
- rb_self: &FrameCodec,
230
- rb_input: RString,
231
- ) -> Result<RString, Error> {
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 frame_codec_decompress(
241
- ruby: &Ruby,
242
- rb_self: &FrameCodec,
243
- rb_input: RString,
244
- max_output: usize,
245
- ) -> Result<RString, Error> {
246
- let compressed: &[u8] = unsafe { rb_input.as_slice() };
247
- let mut dctx = rb_self.dctx.lock().expect("FrameCodec DCtx mutex poisoned");
248
- let out = decompress_bounded(compressed, max_output, &mut dctx)
249
- .map_err(|e| raise_bounded(ruby, e, "zstd frame decode failed"))?;
250
- Ok(ruby.str_from_slice(&out))
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 frame_codec_size(rb_self: &FrameCodec) -> usize {
254
- rb_self.dict_len
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 frame_codec_has_dict(rb_self: &FrameCodec) -> bool {
258
- rb_self.dict_id.is_some()
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 frame_codec_id(rb_self: &FrameCodec) -> Option<u32> {
262
- rb_self.dict_id
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 frame_codec_level(rb_self: &FrameCodec) -> i32 {
266
- rb_self.level
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 frame_codec_get_frame_content_size(
270
- ruby: &Ruby,
271
- rb_input: RString,
272
- ) -> Result<Option<u64>, Error> {
273
- let bytes: &[u8] = unsafe { rb_input.as_slice() };
274
- match parse_frame_content_size(bytes) {
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(raise_bounded(ruby, e, "zstd frame header parse failed")),
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: RefCell<CompressContext>,
292
- dctx: RefCell<DecompressContext>,
548
+ cctx: Mutex<CompressContext>,
549
+ dctx: Mutex<DecompressContext>,
293
550
  }
294
551
 
295
- fn block_codec_new(
296
- ruby: &Ruby,
297
- rb_dict: Option<RString>,
298
- id: u32,
299
- level: i32,
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
- Error::new(
305
- compress_error(ruby),
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(s) => {
312
- let bytes: Vec<u8> = unsafe { s.as_slice().to_vec() };
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
- Error::new(
317
- compress_error(ruby),
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
- Ok(BlockCodec {
327
- dict_len,
328
- dict_id,
329
- level,
330
- cctx: RefCell::new(cctx),
331
- dctx: RefCell::new(dctx),
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 block_codec_compress(
336
- ruby: &Ruby,
337
- rb_self: &BlockCodec,
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 block_codec_decompress(
349
- ruby: &Ruby,
350
- rb_self: &BlockCodec,
351
- rb_input: RString,
352
- max_output: usize,
353
- ) -> Result<RString, Error> {
354
- let compressed: &[u8] = unsafe { rb_input.as_slice() };
355
- let mut dctx = rb_self.dctx.borrow_mut();
356
- let out = decompress_bounded(compressed, max_output, &mut dctx)
357
- .map_err(|e| raise_bounded(ruby, e, "zstd block decode failed"))?;
358
- Ok(ruby.str_from_slice(&out))
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: &BlockCodec) -> usize {
362
- rb_self.dict_len
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: &BlockCodec) -> bool {
366
- rb_self.dict_id.is_some()
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: &BlockCodec) -> i32 {
370
- rb_self.level
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: RefCell<Option<TrainerState>>,
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 dict_trainer_new(_ruby: &Ruby, max_dict_size: usize) -> RbDictTrainer {
387
- RbDictTrainer {
388
- max_dict_size,
389
- inner: RefCell::new(Some(TrainerState {
390
- samples: Vec::new(),
391
- total_bytes: 0,
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 dict_trainer_add_sample(
397
- ruby: &Ruby,
398
- rb_self: &RbDictTrainer,
399
- rb_data: RString,
400
- ) -> Result<(), Error> {
401
- let mut borrow = rb_self.inner.borrow_mut();
402
- let state = borrow.as_mut().ok_or_else(|| {
403
- Error::new(
404
- ruby.exception_runtime_error(),
405
- "DictTrainer already consumed by #train",
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 dict_trainer_sample_count(ruby: &Ruby, rb_self: &RbDictTrainer) -> Result<usize, Error> {
418
- let borrow = rb_self.inner.borrow();
419
- borrow.as_ref().map(|s| s.samples.len()).ok_or_else(|| {
420
- Error::new(
421
- ruby.exception_runtime_error(),
422
- "DictTrainer already consumed by #train",
423
- )
424
- })
425
- }
426
-
427
- fn dict_trainer_total_bytes(ruby: &Ruby, rb_self: &RbDictTrainer) -> Result<usize, Error> {
428
- let borrow = rb_self.inner.borrow();
429
- borrow.as_ref().map(|s| s.total_bytes).ok_or_else(|| {
430
- Error::new(
431
- ruby.exception_runtime_error(),
432
- "DictTrainer already consumed by #train",
433
- )
434
- })
435
- }
436
-
437
- fn dict_trainer_train(ruby: &Ruby, rb_self: &RbDictTrainer) -> Result<RString, Error> {
438
- let state = rb_self.inner.borrow_mut().take().ok_or_else(|| {
439
- Error::new(
440
- ruby.exception_runtime_error(),
441
- "DictTrainer already consumed by #train",
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 Ok(ruby.str_from_slice(b""));
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
- Ok(ruby.str_from_slice(&dict_bytes))
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 dict_trainer_max_dict_size(rb_self: &RbDictTrainer) -> usize {
463
- rb_self.max_dict_size
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 dict_trainer_trained(rb_self: &RbDictTrainer) -> bool {
467
- rb_self.inner.borrow().is_none()
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
- #[magnus::init]
473
- fn init(ruby: &Ruby) -> Result<(), Error> {
474
- unsafe { rb_sys::rb_ext_ractor_safe(true) };
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 = ruby.define_module("Zrip")?;
836
+ let module = unsafe { rb::define_module(c"Zrip")? };
477
837
 
478
838
  let decompress_error_class =
479
- module.define_error("DecompressError", ruby.exception_standard_error())?;
839
+ unsafe { rb::define_error_under(module, c"DecompressError", rb_sys::rb_eStandardError)? };
480
840
  DECOMPRESS_ERROR
481
- .set(Opaque::from(decompress_error_class))
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.define_error("CompressError", ruby.exception_standard_error())?;
845
+ unsafe { rb::define_error_under(module, c"CompressError", rb_sys::rb_eStandardError)? };
486
846
  COMPRESS_ERROR
487
- .set(Opaque::from(compress_error_class))
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.define_error("MissingContentSizeError", decompress_error_class)?;
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(Opaque::from(missing_content_size_error_class))
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.define_error("OutputSizeLimitError", decompress_error_class)?;
858
+ unsafe { rb::define_error_under(module, c"OutputSizeLimitError", decompress_error_class)? };
498
859
  OUTPUT_SIZE_LIMIT_ERROR
499
- .set(Opaque::from(output_size_limit_error_class))
860
+ .set(GlobalValue(output_size_limit_error_class))
500
861
  .unwrap_or_else(|_| panic!("init called more than once"));
501
862
 
502
- // FrameCodec
503
- let frame_codec_class = module.define_class("FrameCodec", ruby.class_object())?;
504
- frame_codec_class.define_singleton_method("_native_new", function!(frame_codec_new, 3))?;
505
- frame_codec_class.define_singleton_method(
506
- "get_frame_content_size",
507
- function!(frame_codec_get_frame_content_size, 1),
508
- )?;
509
- frame_codec_class.define_method("compress", method!(frame_codec_compress, 1))?;
510
- frame_codec_class.define_method("_native_decompress", method!(frame_codec_decompress, 2))?;
511
- frame_codec_class.define_method("size", method!(frame_codec_size, 0))?;
512
- frame_codec_class.define_method("has_dict?", method!(frame_codec_has_dict, 0))?;
513
- frame_codec_class.define_method("id", method!(frame_codec_id, 0))?;
514
- frame_codec_class.define_method("level", method!(frame_codec_level, 0))?;
515
-
516
- // BlockCodec
517
- let block_codec_class = module.define_class("BlockCodec", ruby.class_object())?;
518
- block_codec_class.define_singleton_method("_native_new", function!(block_codec_new, 3))?;
519
- block_codec_class.define_method("compress", method!(block_codec_compress, 1))?;
520
- block_codec_class.define_method("_native_decompress", method!(block_codec_decompress, 2))?;
521
- block_codec_class.define_method("size", method!(block_codec_size, 0))?;
522
- block_codec_class.define_method("has_dict?", method!(block_codec_has_dict, 0))?;
523
- block_codec_class.define_method("level", method!(block_codec_level, 0))?;
524
-
525
- // DictTrainer
526
- let trainer_class = module.define_class("DictTrainer", ruby.class_object())?;
527
- trainer_class.define_singleton_method("_native_new", function!(dict_trainer_new, 1))?;
528
- trainer_class.define_method("add_sample", method!(dict_trainer_add_sample, 1))?;
529
- trainer_class.define_method("sample_count", method!(dict_trainer_sample_count, 0))?;
530
- trainer_class.define_method("total_bytes", method!(dict_trainer_total_bytes, 0))?;
531
- trainer_class.define_method("train", method!(dict_trainer_train, 0))?;
532
- trainer_class.define_method("max_dict_size", method!(dict_trainer_max_dict_size, 0))?;
533
- trainer_class.define_method("trained?", method!(dict_trainer_trained, 0))?;
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
  }