Use ZSTD_compressStream2 with endOp argument

src/libzstd.jl

@@ -65,8 +65,37 @@ function reset!(cstream::CStream, srcsize::Integer)
 
 end
 
-function compress!(cstream::CStream)
-    return LibZstd.ZSTD_compressStream(cstream, cstream.obuffer, cstream.ibuffer)
+"""
+    compress!(cstream::CStream; endOp=:continue)
+
+Compress a Zstandard `CStream`. Optionally, provide one one of the following end directives
+* `:continue` - (default) collect more data, encoder decides when to output compressed result, for optimal compression ratio
+* `:flush` - flush any data provided so far
+* `:end` - flush any remaining data _and_ close current frame
+
+If `:end` is provided on the first call to `compress!`, the size of the input data will be recorded in the frame header.
+This is equivalent to calling [`CodecZstd.LibZstd.ZSTD_compress2`](@ref).
+See also [`CodecZstd.LibZstd.ZSTD_CCtx_setPledgedSrcSize`](@ref).
+
+See the Zstd manual for additional details:
+https://facebook.github.io/zstd/zstd_manual.html
+"""
+function compress!(cstream::CStream; endOp::Union{Symbol,LibZstd.ZSTD_EndDirective}=LibZstd.ZSTD_e_continue)
+    return LibZstd.ZSTD_compressStream2(cstream, cstream.obuffer, cstream.ibuffer, endOp)
+end
+
+# Support endOp keyword of compress! above when providing a Symbol
+function Base.convert(::Type{LibZstd.ZSTD_EndDirective}, endOp::Symbol)
+    endOp = if endOp == :continue
+        LibZstd.ZSTD_e_continue
+    elseif endOp == :flush
+        LibZstd.ZSTD_e_flush
+    elseif endOp == :end
+        LibZstd.ZSTD_e_end
+    else
+        throw(ArgumentError("Received value `:$endOp` for `endOp`, but only :continue, :flush, or :end are allowed values."))
+    end
+    return endOp
 end
 
 function finish!(cstream::CStream)

test/compress_endOp.jl

@@ -0,0 +1,61 @@
+using CodecZstd
+using Test
+
+@testset "compress! endOp = :continue" begin
+    data = rand(1:100, 1024*1024)
+    cstream = CodecZstd.CStream()
+    cstream.ibuffer.src = pointer(data)
+    cstream.ibuffer.size = sizeof(data)
+    cstream.ibuffer.pos = 0
+    cstream.obuffer.dst = Base.Libc.malloc(sizeof(data)*2)
+    cstream.obuffer.size = sizeof(data)*2
+    cstream.obuffer.pos = 0
+    try
+        GC.@preserve data begin
+            # default endOp
+            @test CodecZstd.compress!(cstream; endOp=:continue) == 0
+            @test CodecZstd.find_decompressed_size(cstream.obuffer.dst, cstream.obuffer.pos) == CodecZstd.ZSTD_CONTENTSIZE_UNKNOWN
+        end
+    finally
+        Base.Libc.free(cstream.obuffer.dst)
+    end
+end
+
+@testset "compress! endOp = :flush" begin
+    data = rand(1:100, 1024*1024)
+    cstream = CodecZstd.CStream()
+    cstream.ibuffer.src = pointer(data)
+    cstream.ibuffer.size = sizeof(data)
+    cstream.ibuffer.pos = 0
+    cstream.obuffer.dst = Base.Libc.malloc(sizeof(data)*2)
+    cstream.obuffer.size = sizeof(data)*2
+    cstream.obuffer.pos = 0
+    try
+        GC.@preserve data begin
+            @test CodecZstd.compress!(cstream; endOp=:flush) == 0
+            @test CodecZstd.find_decompressed_size(cstream.obuffer.dst, cstream.obuffer.pos) == CodecZstd.ZSTD_CONTENTSIZE_UNKNOWN
+        end
+    finally
+        Base.Libc.free(cstream.obuffer.dst)
+    end
+end
+
+@testset "compress! endOp = :end" begin
+    data = rand(1:100, 1024*1024)
+    cstream = CodecZstd.CStream()
+    cstream.ibuffer.src = pointer(data)
+    cstream.ibuffer.size = sizeof(data)
+    cstream.ibuffer.pos = 0
+    cstream.obuffer.dst = Base.Libc.malloc(sizeof(data)*2)
+    cstream.obuffer.size = sizeof(data)*2
+    cstream.obuffer.pos = 0
+    try
+        GC.@preserve data begin
+            # The frame should contain the decompressed size
+            @test CodecZstd.compress!(cstream; endOp=:end) == 0
+            @test CodecZstd.find_decompressed_size(cstream.obuffer.dst, cstream.obuffer.pos) == sizeof(data)
+        end
+    finally
+        Base.Libc.free(cstream.obuffer.dst)
+    end
+end

test/runtests.jl

@@ -27,4 +27,6 @@ using Test
     TranscodingStreams.test_roundtrip_write(ZstdCompressorStream, ZstdDecompressorStream)
     TranscodingStreams.test_roundtrip_lines(ZstdCompressorStream, ZstdDecompressorStream)
     TranscodingStreams.test_roundtrip_transcode(ZstdCompressor, ZstdDecompressor)
+
+    include("compress_endOp.jl")
 end