zstd-ruby 1.4.4.0 → 1.5.5.0

data/ext/zstdruby/libzstd/{common/zstd_errors.h → zstd_errors.h}

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-present, Yann Collet, Facebook, Inc.
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
  *
  * This source code is licensed under both the BSD-style license (found in the
@@ -20,19 +20,31 @@ extern "C" {
 
 
 /* =====   ZSTDERRORLIB_API : control library symbols visibility   ===== */
-#ifndef ZSTDERRORLIB_VISIBILITY
-#  if defined(__GNUC__) && (__GNUC__ >= 4)
-#    define ZSTDERRORLIB_VISIBILITY __attribute__ ((visibility ("default")))
+#ifndef ZSTDERRORLIB_VISIBLE
+   /* Backwards compatibility with old macro name */
+#  ifdef ZSTDERRORLIB_VISIBILITY
+#    define ZSTDERRORLIB_VISIBLE ZSTDERRORLIB_VISIBILITY
+#  elif defined(__GNUC__) && (__GNUC__ >= 4) && !defined(__MINGW32__)
+#    define ZSTDERRORLIB_VISIBLE __attribute__ ((visibility ("default")))
 #  else
-#    define ZSTDERRORLIB_VISIBILITY
+#    define ZSTDERRORLIB_VISIBLE
 #  endif
 #endif
+
+#ifndef ZSTDERRORLIB_HIDDEN
+#  if defined(__GNUC__) && (__GNUC__ >= 4) && !defined(__MINGW32__)
+#    define ZSTDERRORLIB_HIDDEN __attribute__ ((visibility ("hidden")))
+#  else
+#    define ZSTDERRORLIB_HIDDEN
+#  endif
+#endif
+
 #if defined(ZSTD_DLL_EXPORT) && (ZSTD_DLL_EXPORT==1)
-#  define ZSTDERRORLIB_API __declspec(dllexport) ZSTDERRORLIB_VISIBILITY
+#  define ZSTDERRORLIB_API __declspec(dllexport) ZSTDERRORLIB_VISIBLE
 #elif defined(ZSTD_DLL_IMPORT) && (ZSTD_DLL_IMPORT==1)
-#  define ZSTDERRORLIB_API __declspec(dllimport) ZSTDERRORLIB_VISIBILITY /* It isn't required but allows to generate better code, saving a function pointer load from the IAT and an indirect jump.*/
+#  define ZSTDERRORLIB_API __declspec(dllimport) ZSTDERRORLIB_VISIBLE /* It isn't required but allows to generate better code, saving a function pointer load from the IAT and an indirect jump.*/
 #else
-#  define ZSTDERRORLIB_API ZSTDERRORLIB_VISIBILITY
+#  define ZSTDERRORLIB_API ZSTDERRORLIB_VISIBLE
 #endif
 
 /*-*********************************************
@@ -58,14 +70,17 @@ typedef enum {
   ZSTD_error_frameParameter_windowTooLarge = 16,
   ZSTD_error_corruption_detected = 20,
   ZSTD_error_checksum_wrong      = 22,
+  ZSTD_error_literals_headerWrong = 24,
   ZSTD_error_dictionary_corrupted      = 30,
   ZSTD_error_dictionary_wrong          = 32,
   ZSTD_error_dictionaryCreation_failed = 34,
   ZSTD_error_parameter_unsupported   = 40,
+  ZSTD_error_parameter_combination_unsupported = 41,
   ZSTD_error_parameter_outOfBound    = 42,
   ZSTD_error_tableLog_tooLarge       = 44,
   ZSTD_error_maxSymbolValue_tooLarge = 46,
   ZSTD_error_maxSymbolValue_tooSmall = 48,
+  ZSTD_error_stabilityCondition_notRespected = 50,
   ZSTD_error_stage_wrong       = 60,
   ZSTD_error_init_missing      = 62,
   ZSTD_error_memory_allocation = 64,
@@ -73,9 +88,15 @@ typedef enum {
   ZSTD_error_dstSize_tooSmall = 70,
   ZSTD_error_srcSize_wrong    = 72,
   ZSTD_error_dstBuffer_null   = 74,
+  ZSTD_error_noForwardProgress_destFull = 80,
+  ZSTD_error_noForwardProgress_inputEmpty = 82,
   /* following error codes are __NOT STABLE__, they can be removed or changed in future versions */
   ZSTD_error_frameIndex_tooLarge = 100,
   ZSTD_error_seekableIO          = 102,
+  ZSTD_error_dstBuffer_wrong     = 104,
+  ZSTD_error_srcBuffer_wrong     = 105,
+  ZSTD_error_sequenceProducer_failed = 106,
+  ZSTD_error_externalSequences_invalid = 107,
   ZSTD_error_maxCode = 120  /* never EVER use this value directly, it can change in future versions! Use ZSTD_isError() instead */
 } ZSTD_ErrorCode;
 

data/ext/zstdruby/main.c

@@ -0,0 +1,20 @@
+#include <common.h>
+VALUE rb_mZstd;
+void zstd_ruby_init(void);
+void zstd_ruby_skippable_frame_init(void);
+void zstd_ruby_streaming_compress_init(void);
+void zstd_ruby_streaming_decompress_init(void);
+
+void
+Init_zstdruby(void)
+{
+#ifdef HAVE_RB_EXT_RACTOR_SAFE
+  rb_ext_ractor_safe(true);
+#endif
+
+  rb_mZstd = rb_define_module("Zstd");
+  zstd_ruby_init();
+  zstd_ruby_skippable_frame_init();
+  zstd_ruby_streaming_compress_init();
+  zstd_ruby_streaming_decompress_init();
+}

data/ext/zstdruby/skippable_frame.c

@@ -0,0 +1,63 @@
+#include <common.h>
+
+extern VALUE rb_mZstd;
+
+static VALUE rb_write_skippable_frame(int argc, VALUE *argv, VALUE self)
+{
+  VALUE input_value;
+  VALUE skip_value;
+  VALUE kwargs;
+  rb_scan_args(argc, argv, "2:", &input_value, &skip_value, &kwargs);
+
+  ID kwargs_keys[1];
+  kwargs_keys[0] = rb_intern("magic_variant");
+  VALUE kwargs_values[1];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+  unsigned magic_variant = (kwargs_values[0] != Qundef) ? (NUM2INT(kwargs_values[0])) : 0;
+
+  StringValue(input_value);
+  StringValue(skip_value);
+  char* input_data = RSTRING_PTR(input_value);
+  size_t input_size = RSTRING_LEN(input_value);
+  char* skip_data = RSTRING_PTR(skip_value);
+  size_t skip_size = RSTRING_LEN(skip_value);
+
+  size_t dst_size = input_size + ZSTD_SKIPPABLEHEADERSIZE + skip_size;
+  VALUE output = rb_str_new(input_data, dst_size);
+  char* output_data = RSTRING_PTR(output);
+  size_t output_size = ZSTD_writeSkippableFrame((void*)output_data, dst_size, (const void*)skip_data, skip_size, magic_variant);
+  if (ZSTD_isError(output_size)) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "write skippable frame failed", ZSTD_getErrorName(output_size));
+  }
+
+  rb_str_resize(output, output_size);
+  return output;
+}
+
+static VALUE rb_read_skippable_frame(VALUE self, VALUE input_value)
+{
+  char* input_data = RSTRING_PTR(input_value);
+  size_t input_size = RSTRING_LEN(input_value);
+
+  if (ZSTD_isSkippableFrame(input_data, input_size) == 0) {
+    return Qnil;
+  }
+  // ref https://github.com/facebook/zstd/blob/321490cd5b9863433b3d44816d04012874e5ecdb/tests/fuzzer.c#L2096
+  size_t const skipLen = 129 * 1024;
+  VALUE output = rb_str_new(NULL, skipLen);
+  char* output_data = RSTRING_PTR(output);
+  unsigned readMagic;
+  size_t output_size = ZSTD_readSkippableFrame((void*)output_data, skipLen, &readMagic, (const void*)input_data, input_size);
+  if (ZSTD_isError(output_size)) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "read skippable frame failed", ZSTD_getErrorName(output_size));
+  }
+  rb_str_resize(output, output_size);
+  return output;
+}
+
+void
+zstd_ruby_skippable_frame_init(void)
+{
+  rb_define_module_function(rb_mZstd, "write_skippable_frame", rb_write_skippable_frame, -1);
+  rb_define_module_function(rb_mZstd, "read_skippable_frame", rb_read_skippable_frame, 1);
+}

data/ext/zstdruby/streaming_compress.c

@@ -0,0 +1,177 @@
+#include <common.h>
+#include <streaming_compress.h>
+
+struct streaming_compress_t {
+  ZSTD_CCtx* ctx;
+  VALUE buf;
+  size_t buf_size;
+};
+
+static void
+streaming_compress_mark(void *p)
+{
+  struct streaming_compress_t *sc = p;
+  rb_gc_mark(sc->buf);
+}
+
+static void
+streaming_compress_free(void *p)
+{
+  struct streaming_compress_t *sc = p;
+  ZSTD_CCtx* ctx = sc->ctx;
+  if (ctx != NULL) {
+    ZSTD_freeCCtx(ctx);
+  }
+  xfree(sc);
+}
+
+static size_t
+streaming_compress_memsize(const void *p)
+{
+    return sizeof(struct streaming_compress_t);
+}
+
+static const rb_data_type_t streaming_compress_type = {
+    "streaming_compress",
+    { streaming_compress_mark, streaming_compress_free, streaming_compress_memsize, },
+     0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE
+rb_streaming_compress_allocate(VALUE klass)
+{
+  struct streaming_compress_t* sc;
+  VALUE obj = TypedData_Make_Struct(klass, struct streaming_compress_t, &streaming_compress_type, sc);
+  sc->ctx = NULL;
+  sc->buf = Qnil;
+  sc->buf_size = 0;
+  return obj;
+}
+
+static VALUE
+rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
+{
+  VALUE compression_level_value;
+  rb_scan_args(argc, argv, "01", &compression_level_value);
+  int compression_level = convert_compression_level(compression_level_value);
+
+  struct streaming_compress_t* sc;
+  TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  size_t const buffOutSize = ZSTD_CStreamOutSize();
+
+  ZSTD_CCtx* ctx = ZSTD_createCCtx();
+  if (ctx == NULL) {
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error");
+  }
+  ZSTD_CCtx_setParameter(ctx, ZSTD_c_compressionLevel, compression_level);
+  sc->ctx = ctx;
+  sc->buf = rb_str_new(NULL, buffOutSize);
+  sc->buf_size = buffOutSize;
+
+  return obj;
+}
+
+#define FIXNUMARG(val, ifnil) \
+    (NIL_P((val)) ? (ifnil) \
+    : (FIX2INT((val))))
+#define ARG_CONTINUE(val)     FIXNUMARG((val), ZSTD_e_continue)
+
+static VALUE
+no_compress(struct streaming_compress_t* sc, ZSTD_EndDirective endOp)
+{
+  ZSTD_inBuffer input = { NULL, 0, 0 };
+  const char* output_data = RSTRING_PTR(sc->buf);
+  VALUE result = rb_str_new(0, 0);
+  size_t ret;
+  do {
+    ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
+
+    size_t const ret = ZSTD_compressStream2(sc->ctx, &output, &input, endOp);
+    if (ZSTD_isError(ret)) {
+      rb_raise(rb_eRuntimeError, "flush error error code: %s", ZSTD_getErrorName(ret));
+    }
+    rb_str_cat(result, output.dst, output.pos);
+  } while (ret > 0);
+  return result;
+}
+
+static VALUE
+rb_streaming_compress_compress(VALUE obj, VALUE src)
+{
+  StringValue(src);
+  const char* input_data = RSTRING_PTR(src);
+  size_t input_size = RSTRING_LEN(src);
+  ZSTD_inBuffer input = { input_data, input_size, 0 };
+
+  struct streaming_compress_t* sc;
+  TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  const char* output_data = RSTRING_PTR(sc->buf);
+  VALUE result = rb_str_new(0, 0);
+  while (input.pos < input.size) {
+    ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
+    size_t const ret = ZSTD_compressStream2(sc->ctx, &output, &input, ZSTD_e_continue);
+    if (ZSTD_isError(ret)) {
+      rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret));
+    }
+    rb_str_cat(result, output.dst, output.pos);
+  }
+  return result;
+}
+
+static VALUE
+rb_streaming_compress_addstr(VALUE obj, VALUE src)
+{
+  StringValue(src);
+  const char* input_data = RSTRING_PTR(src);
+  size_t input_size = RSTRING_LEN(src);
+  ZSTD_inBuffer input = { input_data, input_size, 0 };
+
+  struct streaming_compress_t* sc;
+  TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  const char* output_data = RSTRING_PTR(sc->buf);
+
+  while (input.pos < input.size) {
+    ZSTD_outBuffer output = { (void*)output_data, sc->buf_size, 0 };
+    size_t const result = ZSTD_compressStream2(sc->ctx, &output, &input, ZSTD_e_continue);
+    if (ZSTD_isError(result)) {
+      rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(result));
+    }
+  }
+  return obj;
+}
+
+static VALUE
+rb_streaming_compress_flush(VALUE obj)
+{
+  struct streaming_compress_t* sc;
+  TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  VALUE result = no_compress(sc, ZSTD_e_flush);
+  return result;
+}
+
+static VALUE
+rb_streaming_compress_finish(VALUE obj)
+{
+  struct streaming_compress_t* sc;
+  TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
+  VALUE result = no_compress(sc, ZSTD_e_end);
+  return result;
+}
+
+extern VALUE rb_mZstd, cStreamingCompress;
+void
+zstd_ruby_streaming_compress_init(void)
+{
+  VALUE cStreamingCompress = rb_define_class_under(rb_mZstd, "StreamingCompress", rb_cObject);
+  rb_define_alloc_func(cStreamingCompress, rb_streaming_compress_allocate);
+  rb_define_method(cStreamingCompress, "initialize", rb_streaming_compress_initialize, -1);
+  rb_define_method(cStreamingCompress, "compress", rb_streaming_compress_compress, 1);
+  rb_define_method(cStreamingCompress, "<<", rb_streaming_compress_addstr, 1);
+  rb_define_method(cStreamingCompress, "flush", rb_streaming_compress_flush, 0);
+  rb_define_method(cStreamingCompress, "finish", rb_streaming_compress_finish, 0);
+
+  rb_define_const(cStreamingCompress, "CONTINUE", INT2FIX(ZSTD_e_continue));
+  rb_define_const(cStreamingCompress, "FLUSH", INT2FIX(ZSTD_e_flush));
+  rb_define_const(cStreamingCompress, "END", INT2FIX(ZSTD_e_end));
+}
+

data/ext/zstdruby/streaming_compress.h

@@ -0,0 +1,5 @@
+#if !defined(STREAMING_COMPRESS_H)
+#define STREAMING_COMPRESS_H
+
+
+#endif // STREAMING_COMPRESS_H

data/ext/zstdruby/streaming_decompress.c

@@ -0,0 +1,123 @@
+#include <common.h>
+
+struct streaming_decompress_t {
+  ZSTD_DCtx* ctx;
+  VALUE buf;
+  size_t buf_size;
+};
+
+static void
+streaming_decompress_mark(void *p)
+{
+  struct streaming_decompress_t *sd = p;
+  rb_gc_mark(sd->buf);
+}
+
+static void
+streaming_decompress_free(void *p)
+{
+  struct streaming_decompress_t *sd = p;
+  ZSTD_DCtx* ctx = sd->ctx;
+  if (ctx != NULL) {
+    ZSTD_freeDCtx(ctx);
+  }
+  xfree(sd);
+}
+
+static size_t
+streaming_decompress_memsize(const void *p)
+{
+    return sizeof(struct streaming_decompress_t);
+}
+
+static const rb_data_type_t streaming_decompress_type = {
+    "streaming_decompress",
+    { streaming_decompress_mark, streaming_decompress_free, streaming_decompress_memsize, },
+     0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE
+rb_streaming_decompress_allocate(VALUE klass)
+{
+  struct streaming_decompress_t* sd;
+  VALUE obj = TypedData_Make_Struct(klass, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  sd->ctx = NULL;
+  sd->buf = Qnil;
+  sd->buf_size = 0;
+  return obj;
+}
+
+static VALUE
+rb_streaming_decompress_initialize(VALUE obj)
+{
+  struct streaming_decompress_t* sd;
+  TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  size_t const buffOutSize = ZSTD_DStreamOutSize();
+
+  ZSTD_DCtx* ctx = ZSTD_createDCtx();
+  if (ctx == NULL) {
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx error");
+  }
+  sd->ctx = ctx;
+  sd->buf = rb_str_new(NULL, buffOutSize);
+  sd->buf_size = buffOutSize;
+
+  return obj;
+}
+
+static VALUE
+rb_streaming_decompress_decompress(VALUE obj, VALUE src)
+{
+  StringValue(src);
+  const char* input_data = RSTRING_PTR(src);
+  size_t input_size = RSTRING_LEN(src);
+  ZSTD_inBuffer input = { input_data, input_size, 0 };
+
+  struct streaming_decompress_t* sd;
+  TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  const char* output_data = RSTRING_PTR(sd->buf);
+  VALUE result = rb_str_new(0, 0);
+  while (input.pos < input.size) {
+    ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
+    size_t const ret = ZSTD_decompressStream(sd->ctx, &output, &input);
+    if (ZSTD_isError(ret)) {
+      rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(ret));
+    }
+    rb_str_cat(result, output.dst, output.pos);
+  }
+  return result;
+}
+
+static VALUE
+rb_streaming_decompress_addstr(VALUE obj, VALUE src)
+{
+  StringValue(src);
+  const char* input_data = RSTRING_PTR(src);
+  size_t input_size = RSTRING_LEN(src);
+  ZSTD_inBuffer input = { input_data, input_size, 0 };
+
+  struct streaming_decompress_t* sd;
+  TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
+  const char* output_data = RSTRING_PTR(sd->buf);
+
+  while (input.pos < input.size) {
+    ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 };
+    size_t const result = ZSTD_decompressStream(sd->ctx, &output, &input);
+    if (ZSTD_isError(result)) {
+      rb_raise(rb_eRuntimeError, "compress error error code: %s", ZSTD_getErrorName(result));
+    }
+  }
+  return obj;
+}
+
+extern VALUE rb_mZstd, cStreamingDecompress;
+void
+zstd_ruby_streaming_decompress_init(void)
+{
+  VALUE cStreamingDecompress = rb_define_class_under(rb_mZstd, "StreamingDecompress", rb_cObject);
+  rb_define_alloc_func(cStreamingDecompress, rb_streaming_decompress_allocate);
+  rb_define_method(cStreamingDecompress, "initialize", rb_streaming_decompress_initialize, 0);
+  rb_define_method(cStreamingDecompress, "decompress", rb_streaming_decompress_decompress, 1);
+  rb_define_method(cStreamingDecompress, "<<", rb_streaming_decompress_addstr, 1);
+}
+

data/ext/zstdruby/zstdruby.c

@@ -1,5 +1,6 @@
-#include "zstdruby.h"
-#include "./libzstd/zstd.h"
+#include <common.h>
+
+extern VALUE rb_mZstd;
 
 static VALUE zstdVersion(VALUE self)
 {
@@ -7,41 +8,74 @@ static VALUE zstdVersion(VALUE self)
   return INT2NUM(version);
 }
 
-static VALUE compress(int argc, VALUE *argv, VALUE self)
+static VALUE rb_compress(int argc, VALUE *argv, VALUE self)
 {
   VALUE input_value;
   VALUE compression_level_value;
   rb_scan_args(argc, argv, "11", &input_value, &compression_level_value);
+  int compression_level = convert_compression_level(compression_level_value);
 
-  Check_Type(input_value, RUBY_T_STRING);
-  const char* input_data = RSTRING_PTR(input_value);
+  StringValue(input_value);
+  char* input_data = RSTRING_PTR(input_value);
   size_t input_size = RSTRING_LEN(input_value);
+  size_t max_compressed_size = ZSTD_compressBound(input_size);
 
-  int compression_level;
-  if (NIL_P(compression_level_value)) {
-    compression_level = 0; // The default. See ZSTD_CLEVEL_DEFAULT in zstd_compress.c
-  } else {
-    compression_level = NUM2INT(compression_level_value);
+  VALUE output = rb_str_new(NULL, max_compressed_size);
+  char* output_data = RSTRING_PTR(output);
+  size_t compressed_size = ZSTD_compress((void*)output_data, max_compressed_size,
+                                         (void*)input_data, input_size, compression_level);
+  if (ZSTD_isError(compressed_size)) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "compress failed", ZSTD_getErrorName(compressed_size));
   }
 
-  // do compress
+  rb_str_resize(output, compressed_size);
+  return output;
+}
+
+static VALUE rb_compress_using_dict(int argc, VALUE *argv, VALUE self)
+{
+  VALUE input_value;
+  VALUE dict;
+  VALUE compression_level_value;
+  rb_scan_args(argc, argv, "21", &input_value, &dict, &compression_level_value);
+  int compression_level = convert_compression_level(compression_level_value);
+
+  StringValue(input_value);
+  char* input_data = RSTRING_PTR(input_value);
+  size_t input_size = RSTRING_LEN(input_value);
   size_t max_compressed_size = ZSTD_compressBound(input_size);
 
+  char* dict_buffer = RSTRING_PTR(dict);
+  size_t dict_size = RSTRING_LEN(dict);
+
+  ZSTD_CDict* const cdict = ZSTD_createCDict(dict_buffer, dict_size, compression_level);
+  if (cdict == NULL) {
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCDict failed");
+  }
+  ZSTD_CCtx* const ctx = ZSTD_createCCtx();
+  if (ctx == NULL) {
+    ZSTD_freeCDict(cdict);
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx failed");
+  }
+
   VALUE output = rb_str_new(NULL, max_compressed_size);
   char* output_data = RSTRING_PTR(output);
-
-  size_t compressed_size = ZSTD_compress((void*)output_data, max_compressed_size,
-                                         (const void*)input_data, input_size, compression_level);
+  size_t const compressed_size = ZSTD_compress_usingCDict(ctx, (void*)output_data, max_compressed_size,
+                                             (void*)input_data, input_size, cdict);
 
   if (ZSTD_isError(compressed_size)) {
+    ZSTD_freeCDict(cdict);
+    ZSTD_freeCCtx(ctx);
     rb_raise(rb_eRuntimeError, "%s: %s", "compress failed", ZSTD_getErrorName(compressed_size));
-  } else {
-    rb_str_resize(output, compressed_size);
   }
 
+  rb_str_resize(output, compressed_size);
+  ZSTD_freeCDict(cdict);
+  ZSTD_freeCCtx(ctx);
   return output;
 }
 
+
 static VALUE decompress_buffered(const char* input_data, size_t input_size)
 {
   const size_t outputBufferSize = 4096;
@@ -57,7 +91,6 @@ static VALUE decompress_buffered(const char* input_data, size_t input_size)
     rb_raise(rb_eRuntimeError, "%s: %s", "ZSTD_initDStream failed", ZSTD_getErrorName(initResult));
   }
 
-
   VALUE output_string = rb_str_new(NULL, 0);
   ZSTD_outBuffer output = { NULL, 0, 0 };
 
@@ -79,23 +112,24 @@ static VALUE decompress_buffered(const char* input_data, size_t input_size)
   return output_string;
 }
 
-static VALUE decompress(VALUE self, VALUE input)
+static VALUE rb_decompress(VALUE self, VALUE input_value)
 {
-  Check_Type(input, T_STRING);
-  const char* input_data = RSTRING_PTR(input);
-  size_t input_size = RSTRING_LEN(input);
-
-  uint64_t uncompressed_size = ZSTD_getDecompressedSize(input_data, input_size);
+  StringValue(input_value);
+  char* input_data = RSTRING_PTR(input_value);
+  size_t input_size = RSTRING_LEN(input_value);
 
-  if (uncompressed_size == 0) {
+  unsigned long long const uncompressed_size = ZSTD_getFrameContentSize(input_data, input_size);
+  if (uncompressed_size == ZSTD_CONTENTSIZE_ERROR) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "not compressed by zstd", ZSTD_getErrorName(uncompressed_size));
+  }
+  if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
     return decompress_buffered(input_data, input_size);
   }
 
   VALUE output = rb_str_new(NULL, uncompressed_size);
   char* output_data = RSTRING_PTR(output);
-
-  size_t decompress_size = ZSTD_decompress((void*)output_data, uncompressed_size,
-                                           (const void*)input_data, input_size);
+  size_t const decompress_size = ZSTD_decompress((void*)output_data, uncompressed_size,
+                                           (void*)input_data, input_size);
 
   if (ZSTD_isError(decompress_size)) {
     rb_raise(rb_eRuntimeError, "%s: %s", "decompress error", ZSTD_getErrorName(decompress_size));
@@ -104,13 +138,61 @@ static VALUE decompress(VALUE self, VALUE input)
   return output;
 }
 
-VALUE rb_mZstd;
+static VALUE rb_decompress_using_dict(int argc, VALUE *argv, VALUE self)
+{
+  VALUE input_value;
+  VALUE dict;
+  rb_scan_args(argc, argv, "20", &input_value, &dict);
+
+  StringValue(input_value);
+  char* input_data = RSTRING_PTR(input_value);
+  size_t input_size = RSTRING_LEN(input_value);
+  unsigned long long const uncompressed_size = ZSTD_getFrameContentSize(input_data, input_size);
+  if (uncompressed_size == ZSTD_CONTENTSIZE_ERROR) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "not compressed by zstd", ZSTD_getErrorName(uncompressed_size));
+  }
+  if (uncompressed_size == ZSTD_CONTENTSIZE_UNKNOWN) {
+    return decompress_buffered(input_data, input_size);
+  }
+  VALUE output = rb_str_new(NULL, uncompressed_size);
+  char* output_data = RSTRING_PTR(output);
+
+  char* dict_buffer = RSTRING_PTR(dict);
+  size_t dict_size = RSTRING_LEN(dict);
+  ZSTD_DDict* const ddict = ZSTD_createDDict(dict_buffer, dict_size);
+  if (ddict == NULL) {
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDDict failed");
+  }
+
+  unsigned const expected_dict_id = ZSTD_getDictID_fromDDict(ddict);
+  unsigned const actual_dict_id = ZSTD_getDictID_fromFrame(input_data, input_size);
+  if (expected_dict_id != actual_dict_id) {
+    ZSTD_freeDDict(ddict);
+    rb_raise(rb_eRuntimeError, "%s: %s", "DictID mismatch", ZSTD_getErrorName(uncompressed_size));
+  }
+
+  ZSTD_DCtx* const ctx = ZSTD_createDCtx();
+  if (ctx == NULL) {
+    ZSTD_freeDDict(ddict);
+    rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx failed");
+  }
+  size_t const decompress_size = ZSTD_decompress_usingDDict(ctx, output_data, uncompressed_size, input_data, input_size, ddict);
+  if (ZSTD_isError(decompress_size)) {
+    ZSTD_freeDDict(ddict);
+    ZSTD_freeDCtx(ctx);
+    rb_raise(rb_eRuntimeError, "%s: %s", "decompress error", ZSTD_getErrorName(decompress_size));
+  }
+  ZSTD_freeDDict(ddict);
+  ZSTD_freeDCtx(ctx);
+  return output;
+}
 
 void
-Init_zstdruby(void)
+zstd_ruby_init(void)
 {
-  rb_mZstd = rb_define_module("Zstd");
   rb_define_module_function(rb_mZstd, "zstd_version", zstdVersion, 0);
-  rb_define_module_function(rb_mZstd, "compress", compress, -1);
-  rb_define_module_function(rb_mZstd, "decompress", decompress, 1);
+  rb_define_module_function(rb_mZstd, "compress", rb_compress, -1);
+  rb_define_module_function(rb_mZstd, "compress_using_dict", rb_compress_using_dict, -1);
+  rb_define_module_function(rb_mZstd, "decompress", rb_decompress, 1);
+  rb_define_module_function(rb_mZstd, "decompress_using_dict", rb_decompress_using_dict, -1);
 }

data/lib/zstd-ruby/version.rb

@@ -1,3 +1,3 @@
 module Zstd
-  VERSION = "1.4.4.0"
+  VERSION = "1.5.5.0"
 end

data/lib/zstd-ruby.rb

@@ -2,5 +2,4 @@ require "zstd-ruby/version"
 require "zstd-ruby/zstdruby"
 
 module Zstd
-  # Your code goes here...
 end