//===- llvm/unittest/Support/ThreadSafeAllocatorTest.cpp ------------------===//
//
// 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/Support/ThreadSafeAllocator.h"
#include "llvm/Config/llvm-config.h"
#include "llvm/Support/ThreadPool.h"
#include "gtest/gtest.h"
#include <atomic>
#include <thread>

using namespace llvm;

namespace {

struct AllocCondition {
  std::mutex BusyLock, EndLock;
  std::condition_variable Busy, End;
  bool IsBusy = false, IsEnd = false;
  std::atomic<unsigned> BytesAllocated = 0;

  void startAllocation() {
    {
      std::lock_guard<std::mutex> Lock(BusyLock);
      IsBusy = true;
    }
    Busy.notify_all();
  }
  void waitAllocationStarted() {
    std::unique_lock<std::mutex> LBusy(BusyLock);
    Busy.wait(LBusy, [&]() { return IsBusy; });
    IsBusy = false;
  }
  void finishAllocation() {
    {
      std::lock_guard<std::mutex> Lock(EndLock);
      IsEnd = true;
    }
    End.notify_all();
  }
  void waitAllocationFinished() {
    std::unique_lock<std::mutex> LEnd(EndLock);
    // Wait for end state.
    End.wait(LEnd, [&]() { return IsEnd; });
    IsEnd = false;
  }
};

class MockAllocator : public AllocatorBase<MockAllocator> {
public:
  MockAllocator() = default;

  void *Allocate(size_t Size, size_t Alignment) {
    C.startAllocation();
    C.waitAllocationFinished();
    C.BytesAllocated += Size;
    return Reserved;
  }

  AllocCondition &getAllocCondition() { return C; }

private:
  char Reserved[16];
  AllocCondition C;
};

} // namespace

#if (LLVM_ENABLE_THREADS)
TEST(ThreadSafeAllocatorTest, AllocWait) {
  ThreadSafeAllocator<MockAllocator> Alloc;
  AllocCondition *C;
  // Get the allocation from the allocator first since this requires a lock.
  Alloc.applyLocked(
      [&](MockAllocator &Alloc) { C = &Alloc.getAllocCondition(); });
  ThreadPool Threads;
  // First allocation of 1 byte.
  Threads.async([&Alloc]() {
    char *P = (char *)Alloc.Allocate(1, alignof(char));
    P[0] = 0;
  });
  // No allocation yet.
  EXPECT_EQ(C->BytesAllocated, 0u);
  C->waitAllocationStarted(); // wait till 1st alloocation starts.
  // Second allocation of 2 bytes.
  Threads.async([&Alloc]() {
    char *P = (char *)Alloc.Allocate(2, alignof(char));
    P[1] = 0;
  });
  C->finishAllocation(); // finish 1st allocation.

  C->waitAllocationStarted(); // wait till 2nd allocation starts.
  // still 1 byte allocated since 2nd allocation is not finished yet.
  EXPECT_EQ(C->BytesAllocated, 1u);
  C->finishAllocation(); // finish 2nd allocation.

  Threads.wait(); // all allocations done.
  EXPECT_EQ(C->BytesAllocated, 3u);
}

TEST(ThreadSafeAllocatorTest, AllocWithAlign) {
  ThreadSafeAllocator<BumpPtrAllocator> Alloc;
  ThreadPool Threads;

  for (unsigned Index = 1; Index < 100; ++Index)
    Threads.async(
        [&Alloc](unsigned I) {
          int *P = (int *)Alloc.Allocate(sizeof(int) * I, alignof(int));
          P[I - 1] = I;
        },
        Index);

  Threads.wait();

  Alloc.applyLocked([](BumpPtrAllocator &Alloc) {
    EXPECT_EQ(4950U * sizeof(int), Alloc.getBytesAllocated());
  });
}

TEST(ThreadSafeAllocatorTest, SpecificBumpPtrAllocator) {
  ThreadSafeAllocator<SpecificBumpPtrAllocator<int>> Alloc;
  ThreadPool Threads;

  for (unsigned Index = 1; Index < 100; ++Index)
    Threads.async(
        [&Alloc](unsigned I) {
          int *P = Alloc.Allocate(I);
          P[I - 1] = I;
        },
        Index);

  Threads.wait();
}
#endif