@@ -324,7 +324,8 @@ class ZstdCompressContext final : public ZstdContext {
324324  CompressionError ResetStream ();
325325
326326  //  Zstd specific:
327-   CompressionError Init (uint64_t  pledged_src_size);
327+   CompressionError Init (uint64_t  pledged_src_size,
328+                         std::string_view dictionary = {});
328329  CompressionError SetParameter (int  key, int  value);
329330
330331  //  Wrap ZSTD_freeCCtx to remove the return type.
@@ -349,7 +350,9 @@ class ZstdDecompressContext final : public ZstdContext {
349350  CompressionError ResetStream ();
350351
351352  //  Zstd specific:
352-   CompressionError Init (uint64_t  pledged_src_size);
353+   CompressionError Init (uint64_t  pledged_src_size,
354+                         std::string_view dictionary = {});
355+ 
353356  CompressionError SetParameter (int  key, int  value);
354357
355358  //  Wrap ZSTD_freeDCtx to remove the return type.
@@ -875,8 +878,10 @@ class ZstdStream final : public CompressionStream<CompressionContext> {
875878    Environment* env = Environment::GetCurrent (args);
876879    Local<Context> context = env->context ();
877880
878-     CHECK (args.Length () == 4  &&
879-           " init(params, pledgedSrcSize, writeResult, writeCallback)"  );
881+     CHECK ((args.Length () == 4  || args.Length () == 5 ) &&
882+           " init(params, pledgedSrcSize, writeResult, writeCallback[, " 
883+           " dictionary])"  );
884+ 
880885    ZstdStream* wrap;
881886    ASSIGN_OR_RETURN_UNWRAP (&wrap, args.This ());
882887
@@ -904,7 +909,19 @@ class ZstdStream final : public CompressionStream<CompressionContext> {
904909    }
905910
906911    AllocScope alloc_scope (wrap);
907-     CompressionError err = wrap->context ()->Init (pledged_src_size);
912+     std::string_view dictionary;
913+     ArrayBufferViewContents<char > contents;
914+     if  (args.Length () == 5  && !args[4 ]->IsUndefined ()) {
915+       if  (!args[4 ]->IsArrayBufferView ()) {
916+         THROW_ERR_INVALID_ARG_TYPE (
917+             wrap->env (), " dictionary must be an ArrayBufferView if provided"  );
918+         return ;
919+       }
920+       contents.ReadValue (args[4 ]);
921+       dictionary = std::string_view (contents.data (), contents.length ());
922+     }
923+ 
924+     CompressionError err = wrap->context ()->Init (pledged_src_size, dictionary);
908925    if  (err.IsError ()) {
909926      wrap->EmitError (err);
910927      THROW_ERR_ZLIB_INITIALIZATION_FAILED (wrap->env (), err.message );
@@ -1509,14 +1526,26 @@ CompressionError ZstdCompressContext::SetParameter(int key, int value) {
15091526  return  {};
15101527}
15111528
1512- CompressionError ZstdCompressContext::Init (uint64_t  pledged_src_size) {
1529+ CompressionError ZstdCompressContext::Init (uint64_t  pledged_src_size,
1530+                                            std::string_view dictionary) {
15131531  pledged_src_size_ = pledged_src_size;
15141532  cctx_.reset (ZSTD_createCCtx ());
15151533  if  (!cctx_) {
15161534    return  CompressionError (" Could not initialize zstd instance"  ,
15171535                            " ERR_ZLIB_INITIALIZATION_FAILED"  ,
15181536                            -1 );
15191537  }
1538+ 
1539+   if  (!dictionary.empty ()) {
1540+     size_t  ret = ZSTD_CCtx_loadDictionary (
1541+         cctx_.get (), dictionary.data (), dictionary.size ());
1542+     if  (ZSTD_isError (ret)) {
1543+       return  CompressionError (" Failed to load zstd dictionary"  ,
1544+                               " ERR_ZLIB_DICTIONARY_LOAD_FAILED"  ,
1545+                               -1 );
1546+     }
1547+   }
1548+ 
15201549  size_t  result = ZSTD_CCtx_setPledgedSrcSize (cctx_.get (), pledged_src_size);
15211550  if  (ZSTD_isError (result)) {
15221551    return  CompressionError (
@@ -1549,13 +1578,24 @@ CompressionError ZstdDecompressContext::SetParameter(int key, int value) {
15491578  return  {};
15501579}
15511580
1552- CompressionError ZstdDecompressContext::Init (uint64_t  pledged_src_size) {
1581+ CompressionError ZstdDecompressContext::Init (uint64_t  pledged_src_size,
1582+                                              std::string_view dictionary) {
15531583  dctx_.reset (ZSTD_createDCtx ());
15541584  if  (!dctx_) {
15551585    return  CompressionError (" Could not initialize zstd instance"  ,
15561586                            " ERR_ZLIB_INITIALIZATION_FAILED"  ,
15571587                            -1 );
15581588  }
1589+ 
1590+   if  (!dictionary.empty ()) {
1591+     size_t  ret = ZSTD_DCtx_loadDictionary (
1592+         dctx_.get (), dictionary.data (), dictionary.size ());
1593+     if  (ZSTD_isError (ret)) {
1594+       return  CompressionError (" Failed to load zstd dictionary"  ,
1595+                               " ERR_ZLIB_DICTIONARY_LOAD_FAILED"  ,
1596+                               -1 );
1597+     }
1598+   }
15591599  return  {};
15601600}
15611601
0 commit comments