From 448a3a7ddc40fc95b2874cadd12ed911516207a4 Mon Sep 17 00:00:00 2001
From: Andrew Aladjev <aladjev.andrew@gmail.com>
Date: Thu, 30 Sep 2021 21:36:45 +0300
Subject: [PATCH] added header size for dictionary

---
 ext/extconf.rb             |  1 +
 ext/zstds_ext/dictionary.c | 11 +++++++++++
 lib/zstds/dictionary.rb    | 12 ++++++++----
 test/dictionary.test.rb    |  2 ++
 4 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/ext/extconf.rb b/ext/extconf.rb
index 3ba33fa..0ff1385 100644
--- a/ext/extconf.rb
+++ b/ext/extconf.rb
@@ -110,6 +110,7 @@ def require_library(name, functions)
   "zstd",
   %w[
     ZDICT_getDictID
+    ZDICT_getDictHeaderSize
     ZDICT_isError
     ZDICT_trainFromBuffer
     ZSTD_CCtx_loadDictionary
diff --git a/ext/zstds_ext/dictionary.c b/ext/zstds_ext/dictionary.c
index 5728cfa..16aa2d0 100644
--- a/ext/zstds_ext/dictionary.c
+++ b/ext/zstds_ext/dictionary.c
@@ -153,6 +153,16 @@ VALUE zstds_ext_get_dictionary_buffer_id(VALUE ZSTDS_EXT_UNUSED(self), VALUE buf
   return UINT2NUM(id);
 }
 
+VALUE zstds_ext_get_dictionary_header_size(VALUE ZSTDS_EXT_UNUSED(self), VALUE buffer)
+{
+  zstds_result_t result = ZDICT_getDictHeaderSize(RSTRING_PTR(buffer), RSTRING_LEN(buffer));
+  if (ZDICT_isError(result)) {
+    zstds_ext_raise_error(zstds_ext_get_error(ZSTD_getErrorCode(result)));
+  }
+
+  return SIZET2NUM(result);
+}
+
 // -- exports --
 
 void zstds_ext_dictionary_exports(VALUE root_module)
@@ -160,5 +170,6 @@ void zstds_ext_dictionary_exports(VALUE root_module)
   VALUE dictionary = rb_define_class_under(root_module, "Dictionary", rb_cObject);
 
   rb_define_singleton_method(dictionary, "get_buffer_id", zstds_ext_get_dictionary_buffer_id, 1);
+  rb_define_singleton_method(dictionary, "get_header_size", zstds_ext_get_dictionary_header_size, 1);
   rb_define_singleton_method(dictionary, "train_buffer", zstds_ext_train_dictionary_buffer, 2);
 }
diff --git a/lib/zstds/dictionary.rb b/lib/zstds/dictionary.rb
index 9d1527a..9eeb82f 100644
--- a/lib/zstds/dictionary.rb
+++ b/lib/zstds/dictionary.rb
@@ -23,10 +23,6 @@ def initialize(buffer)
       @buffer = buffer
     end
 
-    def id
-      self.class.get_buffer_id @buffer
-    end
-
     def self.train(samples, options = {})
       Validation.validate_array samples
 
@@ -45,5 +41,13 @@ def self.train(samples, options = {})
       buffer = train_buffer samples, options
       new buffer
     end
+
+    def id
+      self.class.get_buffer_id @buffer
+    end
+
+    def header_size
+      self.class.get_header_size @buffer
+    end
   end
 end
diff --git a/test/dictionary.test.rb b/test/dictionary.test.rb
index 1dad9d0..044fddb 100644
--- a/test/dictionary.test.rb
+++ b/test/dictionary.test.rb
@@ -62,12 +62,14 @@ def test_basic
           dictionary = Target.train SAMPLES, :capacity => capacity
 
           assert dictionary.id.positive?
+          assert dictionary.header_size.positive?
           refute_nil dictionary.buffer
           refute_empty dictionary.buffer
 
           dictionary_copy = Target.new dictionary.buffer
 
           assert_equal dictionary.id, dictionary_copy.id
+          assert_equal dictionary.header_size, dictionary_copy.header_size
           assert_equal dictionary.buffer, dictionary_copy.buffer
 
           text            = TEXTS.sample