Add test for insufficient alignment.

Author: chsigg

mlir/unittests/Bytecode/BytecodeTest.cpp

@@ -18,14 +18,10 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
-#include <algorithm>
-#include <cstdlib>
-#include <memory>
-
 using namespace llvm;
 using namespace mlir;
 
-using testing::ElementsAre;
+using ::testing::StartsWith;
 
 StringLiteral IRWithResources = R"(
 module @TestDialectResources attributes {
@@ -34,7 +30,7 @@ module @TestDialectResources attributes {
 {-#
   dialect_resources: {
     builtin: {
-      resource: "0x1000000001000000020000000300000004000000"
+      resource: "0x2000000001000000020000000300000004000000"
     }
   }
 #-}
@@ -52,19 +48,19 @@ TEST(Bytecode, MultiModuleWithResource) {
   std::string buffer;
   llvm::raw_string_ostream ostream(buffer);
   ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
-
-  // Make a sufficiently aligned copy of the buffer for reading it back.
   ostream.flush();
-  constexpr std::size_t kAlignment = 16; // AsmResourceBlob alignment.
-  auto deleter = [](char *ptr) { std::free(ptr); };
-  std::unique_ptr<char, decltype(deleter)> aligned_buffer(
-      static_cast<char *>(std::aligned_alloc(kAlignment, buffer.size())),
-      deleter);
-  std::copy(buffer.begin(), buffer.end(), aligned_buffer.get());
+
+  // Create copy of buffer which is aligned to requested resource alignment.
+  constexpr size_t kAlignment = 0x20;
+  size_t buffer_size = buffer.size();
+  buffer.reserve(buffer_size + kAlignment - 1);
+  size_t pad = ~(uintptr_t)buffer.data() + 1 & kAlignment - 1;
+  buffer.insert(0, pad, ' ');
+  StringRef aligned_buffer(buffer.data() + pad, buffer_size);
 
   // Parse it back
-  OwningOpRef<Operation *> roundTripModule = parseSourceString<Operation *>(
-      {aligned_buffer.get(), buffer.size()}, parseConfig);
+  OwningOpRef<Operation *> roundTripModule =
+      parseSourceString<Operation *>(aligned_buffer, parseConfig);
   ASSERT_TRUE(roundTripModule);
 
   // FIXME: Parsing external resources does not work on big-endian
@@ -92,3 +88,39 @@ TEST(Bytecode, MultiModuleWithResource) {
   checkResourceAttribute(*module);
   checkResourceAttribute(*roundTripModule);
 }
+
+TEST(Bytecode, InsufficientAlignmentFailure) {
+  MLIRContext context;
+  Builder builder(&context);
+  ParserConfig parseConfig(&context);
+  OwningOpRef<Operation *> module =
+      parseSourceString<Operation *>(IRWithResources, parseConfig);
+  ASSERT_TRUE(module);
+
+  // Write the module to bytecode
+  std::string buffer;
+  llvm::raw_string_ostream ostream(buffer);
+  ASSERT_TRUE(succeeded(writeBytecodeToFile(module.get(), ostream)));
+  ostream.flush();
+
+  // Create copy of buffer which is insufficiently aligned.
+  constexpr size_t kAlignment = 0x20;
+  size_t buffer_size = buffer.size();
+  buffer.reserve(buffer_size + kAlignment - 1);
+  size_t pad = ~(uintptr_t)buffer.data() + kAlignment / 2 + 1 & kAlignment - 1;
+  buffer.insert(0, pad, ' ');
+  StringRef misaligned_buffer(buffer.data() + pad, buffer_size);
+
+  std::unique_ptr<Diagnostic> diagnostic;
+  context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
+    diagnostic = std::make_unique<Diagnostic>(std::move(diag));
+  });
+
+  // Try to parse it back and check for alignment error.
+  OwningOpRef<Operation *> roundTripModule =
+      parseSourceString<Operation *>(misaligned_buffer, parseConfig);
+  EXPECT_FALSE(roundTripModule);
+  ASSERT_TRUE(diagnostic);
+  EXPECT_THAT(diagnostic->str(),
+              StartsWith("expected bytecode buffer to be aligned to 32"));
+}