Merge pull request #62 from AndrewCEmil/main

Add support for dictionary use with streaming operations

Author: SpringMT

README.md

@@ -66,6 +66,15 @@ res << stream.compress("def")
 res << stream.finish
 ```
 
+### Streaming Compression using Dictionary
+```ruby
+stream = Zstd::StreamingCompress.new(dict: IO.read('dictionary_file'))
+stream << "abc" << "def"
+res = stream.flush
+stream << "ghi"
+res << stream.finish
+```
+
 ### Simple Decompression
 
 ```ruby
@@ -87,6 +96,15 @@ result << stream.decompress(cstr[0, 10])
 result << stream.decompress(cstr[10..-1])
 ```
 
+### Streaming Decompression using dictionary
+```ruby
+cstr = "" # Compressed data
+stream = Zstd::StreamingDecompress.new(dict: IO.read('dictionary_file'))
+result = ''
+result << stream.decompress(cstr[0, 10])
+result << stream.decompress(cstr[10..-1])
+```
+
 ### Skippable frame
 
 ```ruby

ext/zstdruby/streaming_compress.c

@@ -75,6 +75,11 @@ rb_streaming_compress_initialize(int argc, VALUE *argv, VALUE obj)
   rb_scan_args(argc, argv, "01", &compression_level_value);
   int compression_level = convert_compression_level(compression_level_value);
 
+  ID kwargs_keys[1];
+  kwargs_keys[0] = rb_intern("dict");
+  VALUE kwargs_values[1];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+
   struct streaming_compress_t* sc;
   TypedData_Get_Struct(obj, struct streaming_compress_t, &streaming_compress_type, sc);
   size_t const buffOutSize = ZSTD_CStreamOutSize();
@@ -83,6 +88,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[0] != Qundef && kwargs_values[0] != Qnil) {
+    char* dict_buffer = RSTRING_PTR(kwargs_values[0]);
+    size_t dict_size = RSTRING_LEN(kwargs_values[0]);
+    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

@@ -70,6 +70,14 @@ rb_streaming_decompress_allocate(VALUE klass)
 static VALUE
 rb_streaming_decompress_initialize(VALUE obj)
 {
+  VALUE kwargs;
+  rb_scan_args(argc, argv, "00:", &kwargs);
+
+  ID kwargs_keys[1];
+  kwargs_keys[0] = rb_intern("dict");
+  VALUE kwargs_values[1];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 1, kwargs_values);
+
   struct streaming_decompress_t* sd;
   TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd);
   size_t const buffOutSize = ZSTD_DStreamOutSize();
@@ -78,6 +86,14 @@ rb_streaming_decompress_initialize(VALUE obj)
   if (ctx == NULL) {
     rb_raise(rb_eRuntimeError, "%s", "ZSTD_createDCtx error");
   }
+  if (kwargs_values[0] != Qundef && kwargs_values[0] != Qnil) {
+    char* dict_buffer = RSTRING_PTR(kwargs_values[0]);
+    size_t dict_size = RSTRING_LEN(kwargs_values[0]);
+    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