3 files changed, 72 insertions, 5 deletions

diff --git a/llvm/unittests/CAS/CMakeLists.txt b/llvm/unittests/CAS/CMakeLists.txt
index 0f8fcb9..ee40e6c 100644
--- a/llvm/unittests/CAS/CMakeLists.txt
+++ b/llvm/unittests/CAS/CMakeLists.txt
@@ -8,6 +8,7 @@ add_llvm_unittest(CASTests
   ActionCacheTest.cpp
   CASTestConfig.cpp
   ObjectStoreTest.cpp
+  OnDiskDataAllocatorTest.cpp
   OnDiskTrieRawHashMapTest.cpp
   ProgramTest.cpp
   )
diff --git a/llvm/unittests/CAS/OnDiskDataAllocatorTest.cpp b/llvm/unittests/CAS/OnDiskDataAllocatorTest.cpp
new file mode 100644
index 0000000..966fa03
--- /dev/null
+++ b/llvm/unittests/CAS/OnDiskDataAllocatorTest.cpp
@@ -0,0 +1,66 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CAS/OnDiskDataAllocator.h"
+#include "llvm/CAS/MappedFileRegionArena.h"
+#include "llvm/Config/llvm-config.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/Testing/Support/Error.h"
+#include "llvm/Testing/Support/SupportHelpers.h"
+
+#if LLVM_ENABLE_ONDISK_CAS
+
+using namespace llvm;
+using namespace llvm::cas;
+
+TEST(OnDiskDataAllocatorTest, Allocate) {
+  unittest::TempDir Temp("data-allocator", /*Unique=*/true);
+  constexpr size_t MB = 1024u * 1024u;
+
+  std::optional<OnDiskDataAllocator> Allocator;
+  ASSERT_THAT_ERROR(OnDiskDataAllocator::create(
+                        Temp.path("allocator"), "data", /*MaxFileSize=*/MB,
+                        /*NewFileInitialSize=*/std::nullopt)
+                        .moveInto(Allocator),
+                    Succeeded());
+
+  // Allocate.
+  {
+    for (size_t Size = 1; Size < 16; ++Size) {
+      OnDiskDataAllocator::OnDiskPtr P;
+      ASSERT_THAT_ERROR(Allocator->allocate(Size).moveInto(P), Succeeded());
+      EXPECT_TRUE(
+          isAligned(MappedFileRegionArena::getAlign(), P.getOffset().get()));
+    }
+  }
+
+  // Out of space.
+  {
+    OnDiskDataAllocator::OnDiskPtr P;
+    ASSERT_THAT_ERROR(Allocator->allocate(MB).moveInto(P), Failed());
+  }
+
+  // Check size and capacity.
+  {
+    ASSERT_EQ(Allocator->capacity(), MB);
+    ASSERT_LE(Allocator->size(), MB);
+  }
+
+  // Get.
+  {
+    OnDiskDataAllocator::OnDiskPtr P;
+    ASSERT_THAT_ERROR(Allocator->allocate(32).moveInto(P), Succeeded());
+    ArrayRef<char> Data;
+    ASSERT_THAT_ERROR(Allocator->get(P.getOffset(), 16).moveInto(Data),
+                      Succeeded());
+    ASSERT_THAT_ERROR(Allocator->get(P.getOffset(), 1025).moveInto(Data),
+                      Failed());
+  }
+}
+
+#endif // LLVM_ENABLE_ONDISK_CAS
diff --git a/llvm/unittests/CAS/OnDiskTrieRawHashMapTest.cpp b/llvm/unittests/CAS/OnDiskTrieRawHashMapTest.cpp
index 7bedfe4..6034c70 100644
--- a/llvm/unittests/CAS/OnDiskTrieRawHashMapTest.cpp
+++ b/llvm/unittests/CAS/OnDiskTrieRawHashMapTest.cpp
@@ -71,7 +71,7 @@ TEST_P(OnDiskTrieRawHashMapTestFixture, General) {
   std::optional<FileOffset> Offset;
   std::optional<MutableArrayRef<char>> Data;
   {
-    std::optional<OnDiskTrieRawHashMap::pointer> Insertion;
+    std::optional<OnDiskTrieRawHashMap::OnDiskPtr> Insertion;
     ASSERT_THAT_ERROR(Trie1->insert({Hash0, Data0v1}).moveInto(Insertion),
                       Succeeded());
     EXPECT_EQ(Hash0, (*Insertion)->Hash);
@@ -128,7 +128,7 @@ TEST_P(OnDiskTrieRawHashMapTestFixture, General) {
 
   // Recover from an offset.
   {
-    OnDiskTrieRawHashMap::const_pointer Recovered;
+    OnDiskTrieRawHashMap::ConstOnDiskPtr Recovered;
     ASSERT_THAT_ERROR(Trie1->recoverFromFileOffset(*Offset).moveInto(Recovered),
                       Succeeded());
     ASSERT_TRUE(Recovered);
@@ -140,14 +140,14 @@ TEST_P(OnDiskTrieRawHashMapTestFixture, General) {
   // Recover from a bad offset.
   {
     FileOffset BadOffset(1);
-    OnDiskTrieRawHashMap::const_pointer Recovered;
+    OnDiskTrieRawHashMap::ConstOnDiskPtr Recovered;
     ASSERT_THAT_ERROR(
         Trie1->recoverFromFileOffset(BadOffset).moveInto(Recovered), Failed());
   }
 
   // Insert another thing.
   {
-    std::optional<OnDiskTrieRawHashMap::pointer> Insertion;
+    std::optional<OnDiskTrieRawHashMap::OnDiskPtr> Insertion;
     ASSERT_THAT_ERROR(Trie1->insert({Hash1, Data1}).moveInto(Insertion),
                       Succeeded());
     EXPECT_EQ(Hash1, (*Insertion)->Hash);
@@ -210,7 +210,7 @@ TEST(OnDiskTrieRawHashMapTest, OutOfSpace) {
   auto Hash0 = ArrayRef(Hash0Bytes);
   constexpr StringLiteral Data0v1Bytes = "data0.v1";
   ArrayRef<char> Data0v1 = ArrayRef(Data0v1Bytes.data(), Data0v1Bytes.size());
-  std::optional<OnDiskTrieRawHashMap::pointer> Insertion;
+  std::optional<OnDiskTrieRawHashMap::OnDiskPtr> Insertion;
   ASSERT_THAT_ERROR(Trie->insert({Hash0, Data0v1}).moveInto(Insertion),
                     Failed());
 }