diff --git a/include/mgard-x/Lossless/Zstd.hpp b/include/mgard-x/Lossless/Zstd.hpp index a6efa1e80..2f44262d4 100644 --- a/include/mgard-x/Lossless/Zstd.hpp +++ b/include/mgard-x/Lossless/Zstd.hpp @@ -111,7 +111,19 @@ template class Zstd { uint32_t actual_out_count = 0; actual_out_count = *reinterpret_cast(in_data); - Resize(actual_out_count, compressionLevel, queue_idx); + + // FIX: Only resize out_data if needed, do not touch in_data which contains + // the compressed data we just copied. The original code called Resize() + // which would free and reallocate BOTH in_data and out_data, destroying + // the compressed data when actual_out_count > buffer_size. + if (this->buffer_size < actual_out_count) { + size_t const estimated_out_size = ZSTD_compressBound(actual_out_count); + if (out_data != nullptr) { + MemoryManager::FreeHost(out_data, queue_idx); + } + MemoryManager::MallocHost( + out_data, estimated_out_size + sizeof(size_t), queue_idx); + } DeviceRuntime::SyncQueue(queue_idx); ZSTD_decompress(out_data, actual_out_count, in_data + sizeof(size_t), @@ -137,4 +149,4 @@ template class Zstd { } // namespace mgard_x -#endif \ No newline at end of file +#endif