Initial commit for change supporting streaming dictionary compression

Author: AndrewCEmil

ext/zstdruby/streaming_compress.c

@@ -58,10 +58,11 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
   rb_scan_args(argc, argv, "01:", &compression_level_value, &kwargs);
   int compression_level = convert_compression_level(compression_level_value);
 
-  ID kwargs_keys[1];
+  ID kwargs_keys[2];
   kwargs_keys[0] = rb_intern("no_gvl");
-  VALUE kwargs_values[1];
-  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+  kwargs_keys[1] = rb_intern("dict");
+  VALUE kwargs_values[2];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 2, kwargs_values);
 
   struct streaming_compress_t* sc;
   TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
@@ -72,6 +73,14 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
   if (ctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createCCtx error");
   }
+  if (kwargs_values[1] != Qundef && kwargs_values[1] != Qnil) {
+    char* dict_buffer = RSTRING_PTR(kwargs_values[1]);
+    size_t dict_size = RSTRING_LEN(kwargs_values[1]);
+    size_t load_dict_ret = ZSTD_CCtx_loadDictionary(ctx, dict_buffer, dict_size);
+    if (ZSTD_isError(load_dict_ret)) {
+      rb_raise(rb_eRuntimeError, "%s", "ZSTD_CCtx_loadDictionary failed");
+    }
+  }
   ZSTD_CCtx_setParameter(ctx, ZSTD_c_compressionLevel, compression_level);
   sc->ctx = ctx;
   sc->buf = rb_str_new(NULL, buffOutSize);

ext/zstdruby/streaming_decompress.c

@@ -55,10 +55,11 @@ rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj)
   VALUE kwargs;
   rb_scan_args(argc, argv, "00:", &kwargs);
 
-  ID kwargs_keys[1];
+  ID kwargs_keys[2];
   kwargs_keys[0] = rb_intern("no_gvl");
-  VALUE kwargs_values[1];
-  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+  kwargs_keys[1] = rb_intern("dict");
+  VALUE kwargs_values[2];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 2, kwargs_values);
 
   struct streaming_decompress_t* sd;
   TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
@@ -69,6 +70,14 @@ rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj)
   if (ctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx error");
   }
+  if (kwargs_values[1] != Qundef && kwargs_values[1] != Qnil) {
+    char* dict_buffer = RSTRING_PTR(kwargs_values[1]);
+    size_t dict_size = RSTRING_LEN(kwargs_values[1]);
+    size_t load_dict_ret = ZSTD_DCtx_loadDictionary(ctx, dict_buffer, dict_size);
+    if (ZSTD_isError(load_dict_ret)) {
+      rb_raise(rb_eRuntimeError, "%s", "ZSTD_DCtx_loadDictionary failed");
+    }
+  }
   sd->ctx = ctx;
   sd->buf = rb_str_new(NULL, buffOutSize);
   sd->buf_size = buffOutSize;

spec/zstd-ruby-streaming-compress_spec.rb

@@ -53,6 +53,41 @@
     end
   end
 
+  describe 'dictionary' do
+    let(:dictionary) do
+      IO.read("#{__dir__}/dictionary")
+    end
+    let(:user_json) do
+      IO.read("#{__dir__}/user_springmt.json")
+    end
+    it 'shoud work' do
+      dict_stream = Zstd::StreamingCompress.new(5, dict: dictionary, no_gvl: no_gvl)
+      dict_stream << user_json
+      dict_res = dict_stream.finish
+      stream = Zstd::StreamingCompress.new(5, no_gvl: no_gvl)
+      stream << user_json
+      res = stream.finish
+
+      expect(dict_res.length).to be < res.length
+    end
+  end
+
+  describe 'nil dictionary' do
+    let(:user_json) do
+      IO.read("#{__dir__}/user_springmt.json")
+    end
+    it 'shoud work' do
+      dict_stream = Zstd::StreamingCompress.new(5, dict: nil, no_gvl: no_gvl)
+      dict_stream << user_json
+      dict_res = dict_stream.finish
+      stream = Zstd::StreamingCompress.new(5, no_gvl: no_gvl)
+      stream << user_json
+      res = stream.finish
+
+      expect(dict_res.length).to eq(res.length)
+    end
+  end
+
   if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0')
     describe 'Ractor' do
       it 'should be supported' do

spec/zstd-ruby-streaming-decompress_spec.rb

@@ -32,6 +32,44 @@
     end
   end
 
+  describe 'dictionary streaming decompress + GC.compact' do
+    let(:dictionary) do
+      IO.read("#{__dir__}/dictionary")
+    end
+    let(:user_json) do
+      IO.read("#{__dir__}/user_springmt.json")
+    end
+    it 'shoud work' do
+      compressed_json = Zstd.compress_using_dict(user_json, dictionary)
+      stream = Zstd::StreamingDecompress.new(dict: dictionary, no_gvl: no_gvl)
+      result = ''
+      result << stream.decompress(compressed_json[0, 5])
+      result << stream.decompress(compressed_json[5, 5])
+      GC.compact
+      result << stream.decompress(compressed_json[10..-1])
+      expect(result).to eq(user_json)
+    end
+  end
+
+  describe 'nil dictionary streaming decompress + GC.compact' do
+    let(:dictionary) do
+      IO.read("#{__dir__}/dictionary")
+    end
+    let(:user_json) do
+      IO.read("#{__dir__}/user_springmt.json")
+    end
+    it 'shoud work' do
+      compressed_json = Zstd.compress(user_json)
+      stream = Zstd::StreamingDecompress.new(dict: nil, no_gvl: no_gvl)
+      result = ''
+      result << stream.decompress(compressed_json[0, 5])
+      result << stream.decompress(compressed_json[5, 5])
+      GC.compact
+      result << stream.decompress(compressed_json[10..-1])
+      expect(result).to eq(user_json)
+    end
+  end
+
   if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0')
     describe 'Ractor' do
       it 'should be supported' do