// Copyright 2023 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

#include "pw_allocator/allocator.h"

#include <cstddef>

#include "pw_allocator/capability.h"
#include "pw_allocator/testing.h"
#include "pw_unit_test/framework.h"

namespace {

using ::pw::allocator::Capability;
using ::pw::allocator::Layout;
using AllocatorForTest = ::pw::allocator::test::AllocatorForTest<256>;
using BlockType = AllocatorForTest::BlockType;
static_assert(sizeof(uintptr_t) == BlockType::kAlignment);

TEST(AllocatorTest, HasFlags) {
  AllocatorForTest allocator;
  EXPECT_TRUE(
      allocator.HasCapability(Capability::kImplementsGetRequestedLayout));
  EXPECT_TRUE(allocator.HasCapability(Capability::kImplementsGetUsableLayout));
}

TEST(AllocatorTest, ResizeNull) {
  AllocatorForTest allocator;
  EXPECT_FALSE(allocator.Resize(nullptr, sizeof(uintptr_t)));
}

TEST(AllocatorTest, ResizeZero) {
  AllocatorForTest allocator;
  constexpr Layout layout = Layout::Of<uintptr_t>();
  void* ptr = allocator.Allocate(layout);
  ASSERT_NE(ptr, nullptr);
  EXPECT_FALSE(allocator.Resize(ptr, 0));
}

TEST(AllocatorTest, ResizeSame) {
  AllocatorForTest allocator;
  constexpr Layout layout = Layout::Of<uintptr_t>();
  void* ptr = allocator.Allocate(layout);
  ASSERT_NE(ptr, nullptr);
  EXPECT_TRUE(allocator.Resize(ptr, layout.size()));
  EXPECT_EQ(allocator.resize_ptr(), ptr);
  EXPECT_EQ(allocator.resize_old_size(), layout.size());
  EXPECT_EQ(allocator.resize_new_size(), layout.size());
}

TEST(AllocatorTest, ReallocateNull) {
  AllocatorForTest allocator;
  constexpr Layout old_layout = Layout::Of<uintptr_t>();
  constexpr Layout new_layout(old_layout.size(), old_layout.alignment());
  void* new_ptr = allocator.Reallocate(nullptr, new_layout);

  // Resize should fail and Reallocate should call Allocate.
  EXPECT_EQ(allocator.allocate_size(), new_layout.size());

  // Deallocate should not be called.
  EXPECT_EQ(allocator.deallocate_ptr(), nullptr);
  EXPECT_EQ(allocator.deallocate_size(), 0U);

  // Overall, Reallocate should succeed.
  EXPECT_NE(new_ptr, nullptr);
}

TEST(AllocatorTest, ReallocateZeroNewSize) {
  AllocatorForTest allocator;
  constexpr Layout old_layout = Layout::Of<uintptr_t[3]>();
  void* ptr = allocator.Allocate(old_layout);
  ASSERT_EQ(allocator.allocate_size(), old_layout.size());
  ASSERT_NE(ptr, nullptr);
  allocator.ResetParameters();

  constexpr Layout new_layout(0, old_layout.alignment());
  void* new_ptr = allocator.Reallocate(ptr, new_layout);

  // Reallocate does not call Resize, Allocate, or Deallocate.
  EXPECT_EQ(allocator.resize_ptr(), nullptr);
  EXPECT_EQ(allocator.resize_old_size(), 0U);
  EXPECT_EQ(allocator.resize_new_size(), 0U);
  EXPECT_EQ(allocator.allocate_size(), 0U);
  EXPECT_EQ(allocator.deallocate_ptr(), nullptr);
  EXPECT_EQ(allocator.deallocate_size(), 0U);

  // Overall, Reallocate should fail.
  EXPECT_EQ(new_ptr, nullptr);
}

TEST(AllocatorTest, ReallocateSame) {
  AllocatorForTest allocator;
  constexpr Layout layout = Layout::Of<uintptr_t[3]>();
  void* ptr = allocator.Allocate(layout);
  ASSERT_EQ(allocator.allocate_size(), layout.size());
  ASSERT_NE(ptr, nullptr);
  allocator.ResetParameters();

  void* new_ptr = allocator.Reallocate(ptr, layout);

  // Reallocate should call Resize.
  EXPECT_EQ(allocator.resize_ptr(), ptr);
  EXPECT_EQ(allocator.resize_old_size(), layout.size());
  EXPECT_EQ(allocator.resize_new_size(), layout.size());

  // DoAllocate should not be called.
  EXPECT_EQ(allocator.allocate_size(), 0U);

  // DoDeallocate should not be called.
  EXPECT_EQ(allocator.deallocate_ptr(), nullptr);
  EXPECT_EQ(allocator.deallocate_size(), 0U);

  // Overall, Reallocate should succeed.
  EXPECT_EQ(new_ptr, ptr);
}

TEST(AllocatorTest, ReallocateSmaller) {
  AllocatorForTest allocator;
  constexpr Layout old_layout = Layout::Of<uintptr_t[3]>();
  void* ptr = allocator.Allocate(old_layout);
  ASSERT_EQ(allocator.allocate_size(), old_layout.size());
  ASSERT_NE(ptr, nullptr);
  allocator.ResetParameters();

  constexpr Layout new_layout(sizeof(uintptr_t), old_layout.alignment());
  void* new_ptr = allocator.Reallocate(ptr, new_layout);

  // Reallocate should call Resize.
  EXPECT_EQ(allocator.resize_ptr(), ptr);
  EXPECT_EQ(allocator.resize_old_size(), old_layout.size());
  EXPECT_EQ(allocator.resize_new_size(), new_layout.size());

  // Allocate should not be called.
  EXPECT_EQ(allocator.allocate_size(), 0U);

  // Deallocate should not be called.
  EXPECT_EQ(allocator.deallocate_ptr(), nullptr);
  EXPECT_EQ(allocator.deallocate_size(), 0U);

  // Overall, Reallocate should succeed.
  EXPECT_EQ(new_ptr, ptr);
}

TEST(AllocatorTest, ReallocateLarger) {
  AllocatorForTest allocator;
  constexpr Layout old_layout = Layout::Of<uintptr_t>();
  void* ptr = allocator.Allocate(old_layout);
  ASSERT_EQ(allocator.allocate_size(), old_layout.size());
  ASSERT_NE(ptr, nullptr);

  // The abstraction is a bit leaky here: This tests relies on the details of
  // `Resize` in order to get it to fail and fallback to
  // allocate/copy/deallocate. Allocate a second block, which should prevent the
  // first one from being able to grow.
  void* next = allocator.Allocate(old_layout);
  ASSERT_EQ(allocator.allocate_size(), old_layout.size());
  ASSERT_NE(next, nullptr);
  allocator.ResetParameters();

  constexpr Layout new_layout(sizeof(uintptr_t[3]), old_layout.alignment());
  void* new_ptr = allocator.Reallocate(ptr, new_layout);

  // Reallocate should call Resize.
  EXPECT_EQ(allocator.resize_ptr(), ptr);
  EXPECT_EQ(allocator.resize_old_size(), old_layout.size());
  EXPECT_EQ(allocator.resize_new_size(), new_layout.size());

  // Resize should fail and Reallocate should call Allocate.
  EXPECT_EQ(allocator.allocate_size(), new_layout.size());

  // Deallocate should also be called.
  EXPECT_EQ(allocator.deallocate_ptr(), ptr);
  EXPECT_EQ(allocator.deallocate_size(), old_layout.size());

  // Overall, Reallocate should succeed.
  EXPECT_NE(new_ptr, nullptr);
  EXPECT_NE(new_ptr, ptr);
}

// Test fixture for IsEqual tests.
class BaseAllocator : public pw::Allocator {
 public:
  BaseAllocator(void* ptr) : pw::Allocator(Capabilities()), ptr_(ptr) {}

 private:
  void* DoAllocate(Layout) override {
    void* ptr = ptr_;
    ptr_ = nullptr;
    return ptr;
  }

  void DoDeallocate(void*) override {}
  void DoDeallocate(void*, Layout) override {}

  void* ptr_;
};

// Test fixture for IsEqual tests.
class DerivedAllocator : public BaseAllocator {
 public:
  DerivedAllocator(size_t value, void* ptr)
      : BaseAllocator(ptr), value_(value) {}
  size_t value() const { return value_; }

 private:
  size_t value_;
};

TEST(AllocatorTest, IsEqualFailsWithDifferentObjects) {
  std::array<std::byte, 8> buffer;
  DerivedAllocator derived1(1, buffer.data());
  DerivedAllocator derived2(2, buffer.data());
  EXPECT_FALSE(derived1.IsEqual(derived2));
  EXPECT_FALSE(derived2.IsEqual(derived1));
}

TEST(AllocatorTest, IsEqualSucceedsWithSameObject) {
  std::array<std::byte, 8> buffer;
  DerivedAllocator derived(1, buffer.data());
  BaseAllocator* base = &derived;
  EXPECT_TRUE(derived.IsEqual(*base));
  EXPECT_TRUE(base->IsEqual(derived));
}

}  // namespace