feat: add simple dictionay compression and decompression

Author: SpringMT

ext/zstdruby/common.h

@@ -4,5 +4,12 @@
 #include "ruby.h"
 #include "./libzstd/zstd.h"
 
+static int convert_compression_level(VALUE compression_level_value)
+{
+  if (NIL_P(compression_level_value)) {
+    return ZSTD_CLEVEL_DEFAULT;
+  }
+  return NUM2INT(compression_level_value);
+}
 
 #endif /* ZSTD_RUBY_H */

ext/zstdruby/streaming_compress.c

@@ -53,13 +53,7 @@ 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;
-  if (NIL_P(compression_level_value)) {
-    compression_level = ZSTD_CLEVEL_DEFAULT;
-  } else {
-    compression_level = NUM2INT(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);

ext/zstdruby/zstdruby.c

@@ -8,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);
 
   StringValue(input_value);
-  const char* input_data = RSTRING_PTR(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;
@@ -58,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 };
 
@@ -80,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)
 {
-  StringValue(input);
-  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));
@@ -105,10 +138,61 @@ static VALUE decompress(VALUE self, VALUE input)
   return output;
 }
 
+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
 zstd_ruby_init(void)
 {
   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);
 }

lib/zstd-ruby.rb

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

spec/dictionary

@@ -0,0 +1,19 @@
+{
+  "login": "SpringMT",
+  "id": 579005,
+  "avatar_url": "https://avatars.githubusercontent.com/u/579005?v=4",
+  "gravatar_id": "",
+  "url": "https://api.github.com/users/SpringMT",
+  "html_url": "https://github.com/SpringMT",
+  "followers_url": "https://api.github.com/users/SpringMT/followers",
+  "following_url": "https://api.github.com/users/SpringMT/following{/other_user}",
+  "gists_url": "https://api.github.com/users/SpringMT/gists{/gist_id}",
+  "starred_url": "https://api.github.com/users/SpringMT/starred{/owner}{/repo}",
+  "subscriptions_url": "https://api.github.com/users/SpringMT/subscriptions",
+  "organizations_url": "https://api.github.com/users/SpringMT/orgs",
+  "repos_url": "https://api.github.com/users/SpringMT/repos",
+  "events_url": "https://api.github.com/users/SpringMT/events{/privacy}",
+  "received_events_url": "https://api.github.com/users/SpringMT/received_events",
+  "type": "User",
+  "site_admin": false
+}

spec/user_springmt.json

@@ -0,0 +1,43 @@
+require "spec_helper"
+require 'zstd-ruby'
+require 'securerandom'
+
+# Generate dictionay methods
+# https://github.com/facebook/zstd#the-case-for-small-data-compression
+# https://github.com/facebook/zstd/releases/tag/v1.1.3
+
+RSpec.describe Zstd do
+  describe 'compress_using_dict' do
+    let(:user_json) do
+      IO.read("#{__dir__}/user_springmt.json")
+    end
+    let(:dictionary) do
+      IO.read("#{__dir__}/dictionary")
+    end
+
+    it 'should work' do
+      compressed_using_dict = Zstd.compress_using_dict(user_json, dictionary)
+      compressed = Zstd.compress(user_json)
+      expect(compressed_using_dict.length).to be < compressed.length
+      expect(user_json).to eq(Zstd.decompress_using_dict(compressed_using_dict, dictionary))
+    end
+
+    it 'should work with simple string' do
+      compressed_using_dict = Zstd.compress_using_dict("abc", dictionary)
+      expect("abc").to eq(Zstd.decompress_using_dict(compressed_using_dict, dictionary))
+    end
+
+    it 'should work with blank' do
+      compressed_using_dict = Zstd.compress_using_dict("", dictionary)
+      expect("").to eq(Zstd.decompress_using_dict(compressed_using_dict, dictionary))
+    end
+
+    it 'should support compression levels' do
+      compressed_using_dict    = Zstd.compress_using_dict(user_json, dictionary)
+      compressed_using_dict_10 = Zstd.compress_using_dict(user_json, dictionary, 10)
+      expect(compressed_using_dict_10.length).to be < compressed_using_dict.length
+    end
+  end
+
+end
+