Merge pull request #43 from SpringMT/feature/add-streaming-decompression

feat: add streaming decompression

Author: SpringMT

.github/workflows/ruby.yml

@@ -15,15 +15,14 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        ruby-version: ['2.6', '2.7', '3.0', '3.1']
+        ruby-version: ['2.7', '3.0', '3.1']
 
     steps:
     - uses: actions/checkout@v2
     - name: Set up Ruby
     # To automatically get bug fixes and new Ruby versions for ruby/setup-ruby,
     # change this to (see https://github.com/ruby/setup-ruby#versioning):
-    # uses: ruby/setup-ruby@v1
-      uses: ruby/setup-ruby@168d6a54b412713c0ed60998a78093a270ca8d84
+      uses: ruby/setup-ruby@v1
       with:
         ruby-version: ${{ matrix.ruby-version }}
         bundler-cache: true # runs 'bundle install' and caches installed gems automatically

.gitignore

@@ -17,3 +17,5 @@ vendor/
 
 # rspec failure tracking
 .rspec_status
+
+.ruby-version

ext/zstdruby/main.c

@@ -2,11 +2,13 @@
 VALUE rb_mZstd;
 void zstd_ruby_init(void);
 void zstd_ruby_streaming_compress_init(void);
+void zstd_ruby_streaming_decompress_init(void);
 
 void
 Init_zstdruby(void)
 {
   rb_mZstd = rb_define_module("Zstd");
   zstd_ruby_init();
   zstd_ruby_streaming_compress_init();
+  zstd_ruby_streaming_decompress_init();
 }

ext/zstdruby/streaming_compress.c

@@ -5,14 +5,13 @@ struct streaming_compress_t {
   ZSTD_CCtx* ctx;
   VALUE buf;
   size_t buf_size;
-  size_t pos;
 };
 
 static void
 streaming_compress_mark(void *p)
 {
   struct streaming_compress_t *sc = p;
-  rb_gc_mark((VALUE)sc->ctx);
+  // rb_gc_mark((VALUE)sc->ctx);
   rb_gc_mark(sc->buf);
   rb_gc_mark(sc->buf_size);
 }

ext/zstdruby/streaming_decompress.c

@@ -0,0 +1,125 @@
+#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((VALUE)sd->ctx);
+  rb_gc_mark(sd->buf);
+  rb_gc_mark(sd->buf_size);
+}
+
+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);
+}
+

spec/zstd-ruby-streaming-compress_spec.rb

@@ -11,6 +11,17 @@
     end
   end
 
+  describe '<< + GC.compat' do
+    it 'shoud work' do
+      stream = Zstd::StreamingCompress.new
+      stream << "abc" << "def"
+      GC.compact
+      stream << "ghi"
+      res = stream.finish
+      expect(Zstd.decompress(res)).to eq('abcdefghi')
+    end
+  end
+
   describe '<< + flush' do
     it 'shoud work' do
       stream = Zstd::StreamingCompress.new

spec/zstd-ruby-streaming-decompress_spec.rb

@@ -0,0 +1,20 @@
+require "spec_helper"
+require 'zstd-ruby'
+require 'securerandom'
+
+RSpec.describe Zstd::StreamingDecompress do
+  describe 'streaming decompress' do
+    it 'shoud work' do
+      # str = SecureRandom.hex(150)
+      str = "foo bar buzz" * 100
+      cstr = Zstd.compress(str)
+      stream = Zstd::StreamingDecompress.new
+      result = ''
+      result << stream.decompress(cstr[0, 5])
+      result << stream.decompress(cstr[5, 5])
+      result << stream.decompress(cstr[10..-1])
+      expect(result).to eq(str)
+    end
+  end
+end
+