[llvm] fix zlib buffer truncate edge cases and fix nits in tests

[ckissane]

Summary

• add check before truncating (un)compressed data buffer if the buffer is already a perfect length, to avoid triggering truncate assertion in edge case.
• explictly coerce LLVM_ENABLE_ZLIB to a 0 or 1 value in OFF case, to match current ON, FORCE_ON behavior.
• fix code style nits in zlib tests

[ckissane]

@MaskRay in this revision I handle the nits that you had noted about my zstd testing in https://reviews.llvm.org/D128465, that also apply for the existing zlib code.

[phosek]

LGTM

llvm/unittests/Support/CompressionTest.cpp

26

I'd probably omit changes to this file since they're not really necessary (and not even described in your commit message).

[ckissane]

llvm/unittests/Support/CompressionTest.cpp

26

@MaskRay noted them on the matching zstd test code on my other patch so I'm going to add it to my commit message.

llvm/cmake/config-ix.cmake

if(ZLIB_FOUND)
     list(APPEND CMAKE_REQUIRED_LIBRARIES ${ZLIB_LIBRARY})
     check_symbol_exists(compress2 zlib.h HAVE_ZLIB)
     cmake_pop_check_state()
     if(LLVM_ENABLE_ZLIB STREQUAL FORCE_ON AND NOT HAVE_ZLIB)
       message(FATAL_ERROR "Failed to configure zlib")
     endif()
   endif()
   set(LLVM_ENABLE_ZLIB "${HAVE_ZLIB}")
+else()
+  set(LLVM_ENABLE_ZLIB 0)
 endif()
 
 if(LLVM_ENABLE_LIBXML2)
   if(LLVM_ENABLE_LIBXML2 STREQUAL FORCE_ON)
     find_package(LibXml2 REQUIRED)
   elseif(NOT LLVM_USE_SANITIZER MATCHES "Memory.*")
     find_package(LibXml2)
   endif()

llvm/lib/Support/Compression.cpp

int Res =
       ::compress2((Bytef *)CompressedBuffer.data(), &CompressedSize,
                   (const Bytef *)InputBuffer.data(), InputBuffer.size(), Level);
   if (Res == Z_MEM_ERROR)
     report_bad_alloc_error("Allocation failed");
   assert(Res == Z_OK);
   // Tell MemorySanitizer that zlib output buffer is fully initialized.
   // This avoids a false report when running LLVM with uninstrumented ZLib.
   __msan_unpoison(CompressedBuffer.data(), CompressedSize);
+  if (CompressedSize < CompressedBuffer.size())
     CompressedBuffer.truncate(CompressedSize);
 }
 
 Error zlib::uncompress(StringRef InputBuffer, char *UncompressedBuffer,
                        size_t &UncompressedSize) {
   int Res =
       ::uncompress((Bytef *)UncompressedBuffer, (uLongf *)&UncompressedSize,
                    (const Bytef *)InputBuffer.data(), InputBuffer.size());
   // Tell MemorySanitizer that zlib output buffer is fully initialized.
   // This avoids a false report when running LLVM with uninstrumented ZLib.
   __msan_unpoison(UncompressedBuffer, UncompressedSize);
   return Res ? make_error<StringError>(convertZlibCodeToString(Res),
                                        inconvertibleErrorCode())
              : Error::success();
 }
 
 Error zlib::uncompress(StringRef InputBuffer,
                        SmallVectorImpl<char> &UncompressedBuffer,
                        size_t UncompressedSize) {
   UncompressedBuffer.resize_for_overwrite(UncompressedSize);
   Error E = zlib::uncompress(InputBuffer, UncompressedBuffer.data(),
                              UncompressedSize);
+  if (UncompressedSize < UncompressedBuffer.size())
     UncompressedBuffer.truncate(UncompressedSize);
   return E;
 }
 
 #else
 bool zlib::isAvailable() { return false; }
 void zlib::compress(StringRef InputBuffer,
                     SmallVectorImpl<char> &CompressedBuffer, int Level) {
   llvm_unreachable("zlib::compress is unavailable");

llvm/unittests/Support/CompressionTest.cpp

 #include "gtest/gtest.h"
 
 using namespace llvm;
 using namespace llvm::compression;
 
 namespace {
 
 #if LLVM_ENABLE_ZLIB
-
-void TestZlibCompression(StringRef Input, int Level) {
+static void testZlibCompression(StringRef Input, int Level) {
   SmallString<32> Compressed;
   SmallString<32> Uncompressed;
-
   zlib::compress(Input, Compressed, Level);
 
   // Check that uncompressed buffer is the same as original.
   Error E = zlib::uncompress(Compressed, Uncompressed, Input.size());
   consumeError(std::move(E));
 
   EXPECT_EQ(Input, Uncompressed);
   if (Input.size() > 0) {
     // Uncompression fails if expected length is too short.
     E = zlib::uncompress(Compressed, Uncompressed, Input.size() - 1);
     EXPECT_EQ("zlib error: Z_BUF_ERROR", llvm::toString(std::move(E)));
   }
 }
 
 TEST(CompressionTest, Zlib) {
-  TestZlibCompression("", zlib::DefaultCompression);
+  testZlibCompression("", zlib::DefaultCompression);
 
-  TestZlibCompression("hello, world!", zlib::NoCompression);
-  TestZlibCompression("hello, world!", zlib::BestSizeCompression);
-  TestZlibCompression("hello, world!", zlib::BestSpeedCompression);
-  TestZlibCompression("hello, world!", zlib::DefaultCompression);
+  testZlibCompression("hello, world!", zlib::NoCompression);
+  testZlibCompression("hello, world!", zlib::BestSizeCompression);
+  testZlibCompression("hello, world!", zlib::BestSpeedCompression);
+  testZlibCompression("hello, world!", zlib::DefaultCompression);
 
   const size_t kSize = 1024;
   char BinaryData[kSize];
-  for (size_t i = 0; i < kSize; ++i) {
+  for (size_t i = 0; i < kSize; ++i)
     BinaryData[i] = i & 255;
-  }
   StringRef BinaryDataStr(BinaryData, kSize);
 
-  TestZlibCompression(BinaryDataStr, zlib::NoCompression);
-  TestZlibCompression(BinaryDataStr, zlib::BestSizeCompression);
-  TestZlibCompression(BinaryDataStr, zlib::BestSpeedCompression);
-  TestZlibCompression(BinaryDataStr, zlib::DefaultCompression);
+  testZlibCompression(BinaryDataStr, zlib::NoCompression);
+  testZlibCompression(BinaryDataStr, zlib::BestSizeCompression);
+  testZlibCompression(BinaryDataStr, zlib::BestSpeedCompression);
+  testZlibCompression(BinaryDataStr, zlib::DefaultCompression);
 }
-
 #endif
 
 }