// Copyright (C) 2018-2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once
#include <algorithm>
#include "openvino/runtime/allocator.hpp"

// Modified from SharedTensorAllocator in [openvinotoolkit/openvino ›
// samples/cpp/benchmark_app/shared_tensor_allocator.hpp]
class SharedTensorAllocator : public ov::AllocatorImpl {
 public:
  SharedTensorAllocator(size_t sizeBytes) : size(sizeBytes) { data = new char[size]; }

  // Copy Constructor
  SharedTensorAllocator(const SharedTensorAllocator& other) : size(other.size) {
    data = new char[size];
    std::copy(other.data, other.data + size, data);
  }

  // Copy Assignment Operator
  SharedTensorAllocator& operator=(const SharedTensorAllocator& other) {
    if (this != &other) {
      size = other.size;
      delete[] data;
      data = new char[size];
      std::copy(other.data, other.data + size, data);
    }
    return *this;
  }

  ~SharedTensorAllocator() { delete[] data; }

  void* allocate(const size_t bytes, const size_t) override {
    return bytes <= this->size ? (void*)data : nullptr;
  }

  void deallocate(void* handle, const size_t bytes, const size_t) override {
    if (handle == data) {
      delete[] data;
      data = nullptr;
    }
  }

  bool is_equal(const AllocatorImpl& other) const override {
    auto other_tensor_allocator = dynamic_cast<const SharedTensorAllocator*>(&other);
    return other_tensor_allocator != nullptr && other_tensor_allocator == this;
  }

  char* get_buffer() { return data; }

 private:
  char* data;
  size_t size;
};