[llvm] 6dedbcd - Make BinaryStreamWriter::padToAlignment write blocks vs bytes.

[Stella Laurenzo via llvm-commits]

Author: Stella Laurenzo
Date: 2022-05-07T17:37:18-07:00
New Revision: 6dedbcd5e96ff36ea7d38534068a2d3e4746a929

URL: https://github.com/llvm/llvm-project/commit/6dedbcd5e96ff36ea7d38534068a2d3e4746a929
DIFF: https://github.com/llvm/llvm-project/commit/6dedbcd5e96ff36ea7d38534068a2d3e4746a929.diff

LOG: Make BinaryStreamWriter::padToAlignment write blocks vs bytes.

While I think this is a performance improvement over the original, this actually fixes a correctness issue: For an appendable underlying stream, padToAlignment would fail if the additional padding would have caused the stream to grow since it was doing its own check on bounds. By deferring to the regular writeArray method this takes the same path as everything else, which does the correct bounds check in WritableBinaryStreamRef::checkOffsetForWrite (i.e. skips the extension check if BSF_Append is set). I had started to fix the existing bounds check in BinaryStreamWriter but deferred to this because it layered better and is more efficient/consistent.

It didn't look like this method was tested at all, so I added a unit test.

Differential Revision: https://reviews.llvm.org/D124746

Added: 
    

Modified: 
    llvm/lib/Support/BinaryStreamWriter.cpp
    llvm/unittests/Support/BinaryStreamTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Support/BinaryStreamWriter.cpp b/llvm/lib/Support/BinaryStreamWriter.cpp
index 8c9efa0ed9a92..2a2dd2e9ef084 100644
--- a/llvm/lib/Support/BinaryStreamWriter.cpp
+++ b/llvm/lib/Support/BinaryStreamWriter.cpp
@@ -94,10 +94,11 @@ BinaryStreamWriter::split(uint64_t Off) const {
 
 Error BinaryStreamWriter::padToAlignment(uint32_t Align) {
   uint64_t NewOffset = alignTo(Offset, Align);
-  if (NewOffset > getLength())
-    return make_error<BinaryStreamError>(stream_error_code::stream_too_short);
+  const uint64_t ZerosSize = 64;
+  static constexpr char Zeros[ZerosSize] = {};
   while (Offset < NewOffset)
-    if (auto EC = writeInteger('\0'))
-      return EC;
+    if (auto E = writeArray(
+            ArrayRef<char>(Zeros, std::min(ZerosSize, NewOffset - Offset))))
+      return E;
   return Error::success();
 }

diff  --git a/llvm/unittests/Support/BinaryStreamTest.cpp b/llvm/unittests/Support/BinaryStreamTest.cpp
index f98b7f327682b..5c5f724b2ed24 100644
--- a/llvm/unittests/Support/BinaryStreamTest.cpp
+++ b/llvm/unittests/Support/BinaryStreamTest.cpp
@@ -807,7 +807,7 @@ TEST_F(BinaryStreamTest, StreamWriterIntegerArrays) {
   }
 }
 
-TEST_F(BinaryStreamTest, StringWriterStrings) {
+TEST_F(BinaryStreamTest, StreamWriterStrings) {
   StringRef Strings[] = {"First", "Second", "Third", "Fourth"};
 
   size_t Length = 0;
@@ -831,6 +831,47 @@ TEST_F(BinaryStreamTest, StringWriterStrings) {
   }
 }
 
+TEST_F(BinaryStreamTest, StreamWriterPadToAlignment) {
+  // This test may seem excessive but it is checking for past bugs and corner
+  // cases by making sure that the stream is allowed to grow and that
+  // both multiple pad chunks and single chunk extensions work.
+  AppendingBinaryByteStream Stream(support::little);
+  BinaryStreamWriter Writer(Stream);
+
+  // Offset 0: '0'
+  EXPECT_THAT_ERROR(Writer.writeInteger('0'), Succeeded());
+  // Offset 1..110: 0
+  EXPECT_THAT_ERROR(Writer.padToAlignment(111), Succeeded());
+  // Offset 111: '*'
+  EXPECT_THAT_ERROR(Writer.writeInteger('*'), Succeeded());
+  // Offset 112..120: 0
+  EXPECT_THAT_ERROR(Writer.padToAlignment(11), Succeeded());
+
+  BinaryStreamReader Reader(Stream);
+  char c;
+  // Offset 0
+  EXPECT_THAT_ERROR(Reader.readInteger<char>(c), Succeeded());
+  EXPECT_EQ('0', c);
+  // Offset 1..110
+  for (int i = 0; i < 110; ++i) {
+    char c;
+    EXPECT_THAT_ERROR(Reader.readInteger<char>(c), Succeeded());
+    EXPECT_EQ('\0', c);
+  }
+  // Offset 111
+  EXPECT_THAT_ERROR(Reader.readInteger<char>(c), Succeeded());
+  EXPECT_EQ('*', c);
+  // Offset 112..120
+  for (int i = 0; i < 9; ++i) {
+    char c;
+    EXPECT_THAT_ERROR(Reader.readInteger<char>(c), Succeeded());
+    EXPECT_EQ('\0', c);
+  }
+
+  // EOF.
+  EXPECT_THAT_ERROR(Reader.readInteger<char>(c), Failed());
+}
+
 TEST_F(BinaryStreamTest, StreamWriterAppend) {
   StringRef Strings[] = {"First", "Second", "Third", "Fourth"};
   AppendingBinaryByteStream Stream(support::little);