Add upperBound method

Sources/CompressionCore/Compressor.swift

@@ -16,4 +16,7 @@ public protocol Compressor: Sendable {
     /// Compress `input` into `output`, advancing its initialized count.
     /// Throws if `output` doesn't have enough free capacity.
     func compress(_ input: some CompressibleInput, into output: inout OutputSpan<UInt8>) throws(Failure)
+
+    /// Returns the size of the buffer you should allocate if using the ``compress(_:into:)`` API.
+    func upperBound(forInputSize inputSize: Int) -> Int
 }

Sources/CompressionCore/StreamingDecompressor.swift

@@ -50,7 +50,6 @@ extension StreamingDecompressor where Self: ~Copyable {
                     output.removeAll()
                 }
             } while (hasStoppedOnFullOutput || consumed < chunk.count) && !isFinished
-            return consumed
         }
         return consumed
     }

Sources/CompressionCore/Utilities/withTemporaryAllocation.swift

@@ -1,6 +1,5 @@
 #if compiler(<6.4)
 @usableFromInline
-@discardableResult
 package func withTemporaryAllocation<T: ~Copyable, R: ~Copyable, E: Error>(
     of type: T.Type,
     capacity: Int,

Sources/Zlib/DeflateCompressor.swift

@@ -23,13 +23,35 @@ extension Deflate {
             self.configuration = configuration
         }
 
+        public func upperBound(forInputSize inputSize: Int) -> Int {
+            precondition(inputSize >= 0, "inputSize must not be negative, got \(inputSize)")
+            // Port of deflateBound
+
+            let fixedLength = inputSize + (inputSize >> 3) + (inputSize >> 8) + (inputSize >> 9) + 4
+            let storeLength = inputSize + (inputSize >> 5) + (inputSize >> 7) + (inputSize >> 11) + 7
+
+            let level = configuration.level == .default ? .custom(6) : configuration.level
+
+            let wrapLength =
+                switch configuration.format {
+                case .raw: 0
+                case .zlib: 6
+                case .gzip: 18
+                }
+
+            let hashBits = Int(configuration.memory.rawValue) + 7
+            // w_bits is always 15
+            if hashBits != 8 + 7 {
+                return (15 <= hashBits && level != .none ? fixedLength : storeLength) + wrapLength
+            }
+
+            return inputSize + (inputSize >> 12) + (inputSize >> 14) + (inputSize >> 25) + 13 - 6
+                + wrapLength
+        }
+
         public func compress(_ input: some CompressibleInput) throws(Deflate.Error) -> [UInt8] {
             try input.withSpan { span throws(Deflate.Error) in
-                // compressBound assumes zlib framing (6 bytes of overhead);
-                // pad for formats with larger framing.
-                let bound =
-                    Int(czlib_z_compressBound(czlib_z_uLong(span.count)))
-                    + configuration.format.extraBoundOverhead
+                let bound = upperBound(forInputSize: span.count)
                 return try [UInt8](capacity: bound) { outputSpan throws(Deflate.Error) in
                     try compress(span, into: &outputSpan)
                 }

Sources/Zlib/DeflateConfiguration.swift

@@ -69,7 +69,7 @@ extension Deflate {
 }
 
 extension Deflate {
-    public enum Level: Sendable {
+    public enum Level: Sendable, Equatable {
         /// zlib default (~level 6). Good balance of speed and compression.
         case `default`
         /// No compression. Output is slightly larger than input (header overhead).
@@ -108,16 +108,6 @@ extension Deflate {
             case .gzip: 31  // 15 + 16
             }
         }
-
-        /// Extra output-bound padding on top of `compressBound`, which only
-        /// budgets for zlib framing (2-byte header + 4-byte Adler-32 trailer).
-        /// gzip framing is 18 bytes: 10-byte header + 8-byte CRC-32/size trailer.
-        var extraBoundOverhead: Int {
-            switch self {
-            case .zlib, .raw: 0
-            case .gzip: 12
-            }
-        }
     }
 
     /// Controls the compression algorithm's trade-offs.

Sources/Zstd/ZstdCompressor.swift

@@ -9,9 +9,14 @@ extension Zstd: OneShotCompressionAlgorithm {
             self.configuration = configuration
         }
 
+        public func upperBound(forInputSize inputSize: Int) -> Int {
+            precondition(inputSize >= 0, "inputSize must not be negative, got \(inputSize)")
+            return ZSTD_compressBound(inputSize)
+        }
+
         public func compress(_ input: some CompressionCore.CompressibleInput) throws(Zstd.Error) -> [UInt8] {
             try input.withSpan { inputSpan throws(Zstd.Error) in
-                let bound = try Zstd.check(ZSTD_compressBound(inputSpan.count))
+                let bound = upperBound(forInputSize: inputSpan.count)
 
                 return try [UInt8](capacity: bound) { output throws(Zstd.Error) in
                     try compress(inputSpan, into: &output)

Sources/Zstd/ZstdDecompressor.swift

@@ -14,7 +14,7 @@ extension Zstd: OneShotDecompressionAlgorithm {
                 let declared = try getDecompressedSize(from: inputSpan)
 
                 if declared != ZSTD_CONTENTSIZE_UNKNOWN, let capacity = Int(exactly: declared) {
-                    if let max = configuration.maxDecompressedSize, capacity > Int(max) {
+                    if let max = configuration.maxDecompressedSize, capacity > max {
                         throw .maxDecompressedSizeExceeded
                     }
 

Tests/DeflateTests/DeflateTestSupport.swift

@@ -0,0 +1,10 @@
+import CompressionDeflate
+
+/// Deterministic pseudo-random bytes
+func pseudoRandomBytes(count: Int, seed: UInt64 = 0x5DEE_CE66_D000_0001) -> [UInt8] {
+    var state = seed
+    return (0..<count).map { _ in
+        state = state &* 6_364_136_223_846_793_005 &+ 1_442_695_040_888_963_407
+        return UInt8(truncatingIfNeeded: state >> 33)
+    }
+}

Tests/DeflateTests/DeflateUpperBoundTests.swift

@@ -0,0 +1,44 @@
+import CompressionDeflate
+import Testing
+
+private let allBoundConfigurations: [Deflate.CompressionConfiguration] =
+    [Deflate.Format.zlib, .raw, .gzip].flatMap { format in
+        [Deflate.MemoryUsage.custom(1), .low, .default, .high].flatMap { memory in
+            [Deflate.Level.none, .speed, .best].map { level in
+                Deflate.CompressionConfiguration(level: level, format: format, memory: memory)
+            }
+        }
+    }
+
+@Suite("Deflate upperBound")
+struct DeflateUpperBoundTests {
+    @Test("A buffer of exactly upperBound fits incompressible input", arguments: allBoundConfigurations)
+    func boundFitsIncompressibleInput(configuration: Deflate.CompressionConfiguration) throws {
+        let input = pseudoRandomBytes(count: 256 * 1024)
+        let compressor = Deflate.Compressor(configuration: configuration)
+        let bound = compressor.upperBound(forInputSize: input.count)
+
+        let compressed = try [UInt8](capacity: bound) { output throws(Deflate.Error) in
+            try compressor.compress(input.span, into: &output)
+        }
+
+        let decompressor = Deflate.Decompressor(configuration: .init(format: configuration.format))
+        #expect(try decompressor.decompress(compressed) == input)
+        #expect(bound < input.count + input.count / 4 + 1_024)
+    }
+
+    @Test(
+        "A buffer of exactly upperBound fits any input size",
+        arguments: [0, 1, 2, 3, 39, 40, 100, 65_536], [Deflate.Format.raw, .zlib, .gzip])
+    func boundFitsEverySize(size: Int, format: Deflate.Format) throws {
+        let input = pseudoRandomBytes(count: size)
+        let compressor = Deflate.Compressor(configuration: .init(format: format))
+        let bound = compressor.upperBound(forInputSize: size)
+
+        let compressed = try [UInt8](capacity: bound) { output throws(Deflate.Error) in
+            try compressor.compress(input.span, into: &output)
+        }
+
+        #expect(try Deflate.Decompressor(configuration: .init(format: format)).decompress(compressed) == input)
+    }
+}

Tests/ZstdTests/ZstdUpperBoundTests.swift

@@ -0,0 +1,39 @@
+import Testing
+import Zstandard
+
+private let allBoundConfigurations: [ZstdCompressionConfiguration] =
+    ([.auto] + allStrategies).map { .init(strategy: $0, level: .default) }
+    + [.min, .fastest, .default, .max].map { .init(strategy: .auto, level: $0) }
+
+@Suite("Zstd upperBound")
+struct ZstdUpperBoundTests {
+    @Test("A buffer of exactly upperBound fits incompressible input", arguments: allBoundConfigurations)
+    func boundFitsIncompressibleInput(configuration: ZstdCompressionConfiguration) throws {
+        let input = pseudoRandomBytes(count: 128 * 1024)
+        let compressor = Zstd.Compressor(configuration: configuration)
+        let bound = compressor.upperBound(forInputSize: input.count)
+
+        let compressed = try [UInt8](capacity: bound) { output throws(Zstd.Error) in
+            try compressor.compress(input.span, into: &output)
+        }
+
+        #expect(try Zstd.Decompressor().decompress(compressed) == input)
+        #expect(bound < input.count + 4_096)
+    }
+
+    @Test(
+        "A buffer of exactly upperBound fits any input size",
+        arguments: [0, 1, 2, 3, 100, 1_000, 131_071, 131_072]
+    )
+    func boundFitsEverySize(size: Int) throws {
+        let input = pseudoRandomBytes(count: size)
+        let compressor = Zstd.Compressor()
+        let bound = compressor.upperBound(forInputSize: size)
+
+        let compressed = try [UInt8](capacity: bound) { output throws(Zstd.Error) in
+            try compressor.compress(input.span, into: &output)
+        }
+
+        #expect(try Zstd.Decompressor().decompress(compressed) == input)
+    }
+}