Add pledgeinsize

JuliaIO · Sep 7, 2024 · 11cff42 · 11cff42
1 parent 31cf742
commit 11cff42
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 19 deletions.
diff --git a/src/compression.jl b/src/compression.jl
@@ -101,13 +101,31 @@ function TranscodingStreams.finalize(codec::ZstdCompressor)
     return
 end
 
-function TranscodingStreams.startproc(codec::ZstdCompressor, mode::Symbol, error::Error)
-    code = reset!(codec.cstream, 0 #=unknown source size=#)
+function TranscodingStreams.startproc(codec::ZstdCompressor, mode::Symbol, error::Error)::Symbol
+    code = reset!(codec.cstream)
     if iserror(code)
-        error[] = ErrorException("zstd error")
-        return :error
+        error[] = ErrorException("zstd error resetting compression context")
+        :error
+    else
+        :ok
+    end
+end
+
+if isdefined(TranscodingStreams, :pledgeinsize)
+    function TranscodingStreams.pledgeinsize(codec::ZstdCompressor, insize::Int64, error::Error)::Symbol
+        srcsize = if signbit(insize)
+            ZSTD_CONTENTSIZE_UNKNOWN
+        else
+            Culonglong(insize)
+        end
+        code = LibZstd.ZSTD_CCtx_setPledgedSrcSize(codec.cstream, srcsize)
+        if iserror(code)
+            error[] = ErrorException("zstd error setting pledged source size")
+            :error
+        else
+            :ok
+        end
     end
-    return :ok
 end
 
 function TranscodingStreams.process(codec::ZstdCompressor, input::Memory, output::Memory, error::Error)

diff --git a/src/libzstd.jl b/src/libzstd.jl
@@ -60,23 +60,11 @@ function initialize!(cstream::CStream, level::Integer)
     return LibZstd.ZSTD_initCStream(cstream, level)
 end
 
-function reset!(cstream::CStream, srcsize::Integer)
-    # ZSTD_resetCStream is deprecated
-    # https://github.com/facebook/zstd/blob/9d2a45a705e22ad4817b41442949cd0f78597154/lib/zstd.h#L2253-L2272
+function reset!(cstream::CStream)
     res = LibZstd.ZSTD_CCtx_reset(cstream, LibZstd.ZSTD_reset_session_only)
-    if iserror(res)
-        return res
-    end
-    if srcsize == 0
-        # From zstd.h:
-        # Note: ZSTD_resetCStream() interprets pledgedSrcSize == 0 as ZSTD_CONTENTSIZE_UNKNOWN, but
-        # ZSTD_CCtx_setPledgedSrcSize() does not do the same, so ZSTD_CONTENTSIZE_UNKNOWN must be
-        # explicitly specified.
-        srcsize = ZSTD_CONTENTSIZE_UNKNOWN
-    end
     reset!(cstream.ibuffer)
     reset!(cstream.obuffer)
-    return LibZstd.ZSTD_CCtx_setPledgedSrcSize(cstream, srcsize)
+    return res
 end
 
 """