Add AlignTo method to make a Blob's contents aligned

After loading a Blob from the cache, Dawn may need it to
match a particular alignment. For example, SPIRV must be 4-byte
aligned, or Dawn may need to cast a Blob to a struct layout.

Bug: dawn:549
Change-Id: Iad4857d1ad2d9b41e61e9f177aa7083b1f078be5
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/94532
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Loko Kung <lokokung@google.com>
Commit-Queue: Austin Eng <enga@chromium.org>
diff --git a/src/dawn/native/Blob.cpp b/src/dawn/native/Blob.cpp
index a3ac2b2..ea4d28d 100644
--- a/src/dawn/native/Blob.cpp
+++ b/src/dawn/native/Blob.cpp
@@ -15,14 +15,24 @@
 #include <utility>
 
 #include "dawn/common/Assert.h"
+#include "dawn/common/Math.h"
 #include "dawn/native/Blob.h"
 
 namespace dawn::native {
 
-Blob CreateBlob(size_t size) {
+Blob CreateBlob(size_t size, size_t alignment) {
+    ASSERT(IsPowerOfTwo(alignment));
+    ASSERT(alignment != 0);
     if (size > 0) {
-        uint8_t* data = new uint8_t[size];
-        return Blob::UnsafeCreateWithDeleter(data, size, [=]() { delete[] data; });
+        // Allocate extra space so that there will be sufficient space for |size| even after
+        // the |data| pointer is aligned.
+        // TODO(crbug.com/dawn/824): Use aligned_alloc when possible. It should be available
+        // with C++17 but on macOS it also requires macOS 10.15 to work.
+        size_t allocatedSize = size + alignment - 1;
+        uint8_t* data = new uint8_t[allocatedSize];
+        uint8_t* ptr = AlignPtr(data, alignment);
+        ASSERT(ptr + size <= data + allocatedSize);
+        return Blob::UnsafeCreateWithDeleter(ptr, size, [=]() { delete[] data; });
     } else {
         return Blob();
     }
@@ -77,4 +87,14 @@
     return mSize;
 }
 
+void Blob::AlignTo(size_t alignment) {
+    if (IsPtrAligned(mData, alignment)) {
+        return;
+    }
+
+    Blob blob = CreateBlob(mSize, alignment);
+    memcpy(blob.Data(), mData, mSize);
+    *this = std::move(blob);
+}
+
 }  // namespace dawn::native
diff --git a/src/dawn/native/Blob.h b/src/dawn/native/Blob.h
index 4bdcef7..c5f1c7d 100644
--- a/src/dawn/native/Blob.h
+++ b/src/dawn/native/Blob.h
@@ -43,6 +43,10 @@
     uint8_t* Data();
     size_t Size() const;
 
+    // If the blob data is not aligned to |alignment|, copy it into a new backing store which
+    // is aligned.
+    void AlignTo(size_t alignment);
+
   private:
     // The constructor should be responsible to take ownership of |data| and releases ownership by
     // calling |deleter|. The deleter function is called at ~Blob() and during std::move.
@@ -53,7 +57,7 @@
     std::function<void()> mDeleter;
 };
 
-Blob CreateBlob(size_t size);
+Blob CreateBlob(size_t size, size_t alignment = 1);
 
 }  // namespace dawn::native
 
diff --git a/src/dawn/tests/unittests/native/BlobTests.cpp b/src/dawn/tests/unittests/native/BlobTests.cpp
index 580f0b7..4bf081d 100644
--- a/src/dawn/tests/unittests/native/BlobTests.cpp
+++ b/src/dawn/tests/unittests/native/BlobTests.cpp
@@ -14,6 +14,7 @@
 
 #include <utility>
 
+#include "dawn/common/Math.h"
 #include "dawn/native/Blob.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -163,6 +164,61 @@
     EXPECT_EQ(memcmp(b2.Data(), data, sizeof(data)), 0);
 }
 
+// Test that an empty blob can be requested to have a particular alignment.
+TEST(BlobTests, EmptyAlignTo) {
+    for (size_t alignment : {1, 2, 4, 8, 16, 32}) {
+        Blob b;
+        EXPECT_TRUE(b.Empty());
+        EXPECT_EQ(b.Size(), 0u);
+        EXPECT_EQ(b.Data(), nullptr);
+        b.AlignTo(alignment);
+        // After aligning, it is still empty.
+        EXPECT_TRUE(b.Empty());
+        EXPECT_EQ(b.Size(), 0u);
+        EXPECT_EQ(b.Data(), nullptr);
+    }
+}
+
+// Test that AlignTo makes a blob have a particular alignment.
+TEST(BlobTests, AlignTo) {
+    uint8_t data[64];
+    for (uint8_t i = 0; i < sizeof(data); ++i) {
+        data[i] = i;
+    }
+    // Test multiple alignments.
+    for (size_t alignment : {1, 2, 4, 8, 16, 32}) {
+        for (size_t offset = 0; offset < alignment; ++offset) {
+            // Make a blob pointing to |data| starting at |offset|.
+            size_t size = sizeof(data) - offset;
+            testing::StrictMock<testing::MockFunction<void()>> mockDeleter;
+            Blob b =
+                Blob::UnsafeCreateWithDeleter(&data[offset], size, [&]() { mockDeleter.Call(); });
+            bool alreadyAligned = IsPtrAligned(&data[offset], alignment);
+
+            // The backing store should be deleted at the end of the scope, or because it was
+            // replaced.
+            EXPECT_CALL(mockDeleter, Call());
+
+            b.AlignTo(alignment);
+            if (!alreadyAligned) {
+                // If the Blob is not aligned, its data will be deleted and replaced by AlignTo.
+                testing::Mock::VerifyAndClearExpectations(&mockDeleter);
+            }
+            // The blob should not have changed in size.
+            EXPECT_EQ(b.Size(), size) << "alignment = " << alignment << " offset = " << offset;
+            // The data should be aligned.
+            EXPECT_TRUE(IsPtrAligned(b.Data(), alignment))
+                << "alignment = " << alignment << " offset = " << offset;
+            // The contents should be the same.
+            EXPECT_EQ(memcmp(b.Data(), &data[offset], size), 0)
+                << "alignment = " << alignment << " offset = " << offset;
+            // If the data was already aligned, the blob should point to the same memory.
+            EXPECT_EQ(alreadyAligned, b.Data() == &data[offset])
+                << "alignment = " << alignment << " offset = " << offset;
+        }
+    }
+}
+
 }  // namespace
 
 }  // namespace dawn::native