diff --git a/c_src/zstd_nif.c b/c_src/zstd_nif.c index 26ed8c0..f0345f2 100644 --- a/c_src/zstd_nif.c +++ b/c_src/zstd_nif.c @@ -3,11 +3,20 @@ #include #include +ErlNifTSDKey zstdDecompressContextKey; +ErlNifTSDKey zstdCompressContextKey; + static ERL_NIF_TERM zstd_nif_compress(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) { ErlNifBinary bin, ret_bin; size_t buff_size, compressed_size; unsigned int compression_level; + ZSTD_CCtx* ctx = (ZSTD_CCtx*)enif_tsd_get(zstdCompressContextKey); + if (!ctx) { + ctx = ZSTD_createCCtx(); + enif_tsd_set(zstdCompressContextKey, ctx); + } + if(!enif_inspect_binary(env, argv[0], &bin) || !enif_get_uint(env, argv[1], &compression_level) || compression_level > ZSTD_maxCLevel()) @@ -18,7 +27,7 @@ static ERL_NIF_TERM zstd_nif_compress(ErlNifEnv* env, int argc, const ERL_NIF_TE if(!enif_alloc_binary(buff_size, &ret_bin)) return enif_make_atom(env, "error"); - compressed_size = ZSTD_compress(ret_bin.data, buff_size, bin.data, bin.size, compression_level); + compressed_size = ZSTD_compressCCtx(ctx, ret_bin.data, buff_size, bin.data, bin.size, compression_level); if(ZSTD_isError(compressed_size)) return enif_make_atom(env, "error"); @@ -34,6 +43,12 @@ static ERL_NIF_TERM zstd_nif_decompress(ErlNifEnv* env, int argc, const ERL_NIF_ ErlNifBinary bin; unsigned long long uncompressed_size; + ZSTD_DCtx* ctx = (ZSTD_DCtx*)enif_tsd_get(zstdDecompressContextKey); + if (!ctx) { + ctx = ZSTD_createDCtx(); + enif_tsd_set(zstdDecompressContextKey, ctx); + } + if(!enif_inspect_binary(env, argv[0], &bin)) return enif_make_badarg(env); @@ -41,7 +56,7 @@ static ERL_NIF_TERM zstd_nif_decompress(ErlNifEnv* env, int argc, const ERL_NIF_ outp = enif_make_new_binary(env, uncompressed_size, &out); - if(ZSTD_decompress(outp, uncompressed_size, bin.data, bin.size) != uncompressed_size) + if(ZSTD_decompressDCtx(ctx, outp, uncompressed_size, bin.data, bin.size) != uncompressed_size) return enif_make_atom(env, "error"); return out; @@ -52,4 +67,11 @@ static ErlNifFunc nif_funcs[] = { {"decompress", 1, zstd_nif_decompress} }; -ERL_NIF_INIT(zstd, nif_funcs, NULL, NULL, NULL, NULL); +static int load(ErlNifEnv* env, void** priv_data, ERL_NIF_TERM load_info) + { + enif_tsd_key_create("zstd_decompress_context_key", &zstdDecompressContextKey); + enif_tsd_key_create("zstd_compress_context_key", &zstdCompressContextKey); + return 0; + } + +ERL_NIF_INIT(zstd, nif_funcs, load, NULL, NULL, NULL);