// <memory_resource> -*- C++ -*-

// Copyright (C) 2018-2023 Free Software Foundation, Inc.
//
// This file is part of the GNU ISO C++ Library.  This library is free
// software; you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the
// Free Software Foundation; either version 3, or (at your option)
// any later version.

// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// Under Section 7 of GPL version 3, you are granted additional
// permissions described in the GCC Runtime Library Exception, version
// 3.1, as published by the Free Software Foundation.

// You should have received a copy of the GNU General Public License and
// a copy of the GCC Runtime Library Exception along with this program;
// see the files COPYING3 and COPYING.RUNTIME respectively.  If not, see
// <http://www.gnu.org/licenses/>.

/** @file include/memory_resource
 *  This is a Standard C++ Library header.
 */

#ifndef _GLIBCXX_MEMORY_RESOURCE
#define _GLIBCXX_MEMORY_RESOURCE 1

#pragma GCC system_header

#include <bits/requires_hosted.h> // polymorphic allocation

#if __cplusplus >= 201703L

#include <bits/memory_resource.h>
#include <vector>			// vector
#include <shared_mutex>			// shared_mutex
#include <bits/align.h>			// align
#include <debug/assertions.h>

namespace std _GLIBCXX_VISIBILITY(default)
{
_GLIBCXX_BEGIN_NAMESPACE_VERSION
namespace pmr
{
#ifdef _GLIBCXX_HAS_GTHREADS
  // Header and all contents are present.
# define __cpp_lib_memory_resource 201603L
#else
  // The pmr::synchronized_pool_resource type is missing.
# define __cpp_lib_memory_resource 1
#endif

#if __cplusplus >= 202002L
# define __cpp_lib_polymorphic_allocator 201902L
  template<typename _Tp = std::byte>
    class polymorphic_allocator;
#endif

  // Global memory resources

  /// A pmr::memory_resource that uses `new` to allocate memory
  [[nodiscard, __gnu__::__returns_nonnull__, __gnu__::__const__]]
  memory_resource*
  new_delete_resource() noexcept;

  /// A pmr::memory_resource that always throws `bad_alloc`
  [[nodiscard, __gnu__::__returns_nonnull__, __gnu__::__const__]]
  memory_resource*
  null_memory_resource() noexcept;

  /// Replace the default memory resource pointer
  [[__gnu__::__returns_nonnull__]]
  memory_resource*
  set_default_resource(memory_resource* __r) noexcept;

  /// Get the current default memory resource pointer
  [[__gnu__::__returns_nonnull__]]
  memory_resource*
  get_default_resource() noexcept;

  // Pool resource classes
  struct pool_options;
#ifdef _GLIBCXX_HAS_GTHREADS
  class synchronized_pool_resource;
#endif
  class unsynchronized_pool_resource;
  class monotonic_buffer_resource;

  /// Parameters for tuning a pool resource's behaviour.
  struct pool_options
  {
    /** @brief Upper limit on number of blocks in a chunk.
     *
     * A lower value prevents allocating huge chunks that could remain mostly
     * unused, but means pools will need to replenished more frequently.
     */
    size_t max_blocks_per_chunk = 0;

    /* @brief Largest block size (in bytes) that should be served from pools.
     *
     * Larger allocations will be served directly by the upstream resource,
     * not from one of the pools managed by the pool resource.
     */
    size_t largest_required_pool_block = 0;
  };

  // Common implementation details for un-/synchronized pool resources.
  class __pool_resource
  {
    friend class synchronized_pool_resource;
    friend class unsynchronized_pool_resource;

    __pool_resource(const pool_options& __opts, memory_resource* __upstream);

    ~__pool_resource();

    __pool_resource(const __pool_resource&) = delete;
    __pool_resource& operator=(const __pool_resource&) = delete;

    // Allocate a large unpooled block.
    void*
    allocate(size_t __bytes, size_t __alignment);

    // Deallocate a large unpooled block.
    void
    deallocate(void* __p, size_t __bytes, size_t __alignment);


    // Deallocate unpooled memory.
    void release() noexcept;

    memory_resource* resource() const noexcept
    { return _M_unpooled.get_allocator().resource(); }

    struct _Pool;

    _Pool* _M_alloc_pools();

    const pool_options _M_opts;

    struct _BigBlock;
    // Collection of blocks too big for any pool, sorted by address.
    // This also stores the only copy of the upstream memory resource pointer.
    _GLIBCXX_STD_C::pmr::vector<_BigBlock> _M_unpooled;

    const int _M_npools;
  };

#ifdef _GLIBCXX_HAS_GTHREADS
  /// A thread-safe memory resource that manages pools of fixed-size blocks.
  class synchronized_pool_resource : public memory_resource
  {
  public:
    synchronized_pool_resource(const pool_options& __opts,
				 memory_resource* __upstream)
    __attribute__((__nonnull__));

    synchronized_pool_resource()
    : synchronized_pool_resource(pool_options(), get_default_resource())
    { }

    explicit
    synchronized_pool_resource(memory_resource* __upstream)
    __attribute__((__nonnull__))
    : synchronized_pool_resource(pool_options(), __upstream)
    { }

    explicit
    synchronized_pool_resource(const pool_options& __opts)
    : synchronized_pool_resource(__opts, get_default_resource()) { }

    synchronized_pool_resource(const synchronized_pool_resource&) = delete;

    virtual ~synchronized_pool_resource();

    synchronized_pool_resource&
    operator=(const synchronized_pool_resource&) = delete;

    void release();

    memory_resource*
    upstream_resource() const noexcept
    __attribute__((__returns_nonnull__))
    { return _M_impl.resource(); }

    pool_options options() const noexcept { return _M_impl._M_opts; }

  protected:
    void*
    do_allocate(size_t __bytes, size_t __alignment) override;

    void
    do_deallocate(void* __p, size_t __bytes, size_t __alignment) override;

    bool
    do_is_equal(const memory_resource& __other) const noexcept override
    { return this == &__other; }

  public:
    // Thread-specific pools (only public for access by implementation details)
    struct _TPools;

  private:
    _TPools* _M_alloc_tpools(lock_guard<shared_mutex>&);
    _TPools* _M_alloc_shared_tpools(lock_guard<shared_mutex>&);
    auto _M_thread_specific_pools() noexcept;

    __pool_resource _M_impl;
    __gthread_key_t _M_key;
    // Linked list of thread-specific pools. All threads share _M_tpools[0].
    _TPools* _M_tpools = nullptr;
    mutable shared_mutex _M_mx;
  };
#endif

  /// A non-thread-safe memory resource that manages pools of fixed-size blocks.
  class unsynchronized_pool_resource : public memory_resource
  {
  public:
    [[__gnu__::__nonnull__]]
    unsynchronized_pool_resource(const pool_options& __opts,
				 memory_resource* __upstream);

    unsynchronized_pool_resource()
    : unsynchronized_pool_resource(pool_options(), get_default_resource())
    { }

    [[__gnu__::__nonnull__]]
    explicit
    unsynchronized_pool_resource(memory_resource* __upstream)
    : unsynchronized_pool_resource(pool_options(), __upstream)
    { }

    explicit
    unsynchronized_pool_resource(const pool_options& __opts)
    : unsynchronized_pool_resource(__opts, get_default_resource()) { }

    unsynchronized_pool_resource(const unsynchronized_pool_resource&) = delete;

    virtual ~unsynchronized_pool_resource();

    unsynchronized_pool_resource&
    operator=(const unsynchronized_pool_resource&) = delete;

    void release();

    [[__gnu__::__returns_nonnull__]]
    memory_resource*
    upstream_resource() const noexcept
    { return _M_impl.resource(); }

    pool_options options() const noexcept { return _M_impl._M_opts; }

  protected:
    void*
    do_allocate(size_t __bytes, size_t __alignment) override;

    void
    do_deallocate(void* __p, size_t __bytes, size_t __alignment) override;

    bool
    do_is_equal(const memory_resource& __other) const noexcept override
    { return this == &__other; }

  private:
    using _Pool = __pool_resource::_Pool;

    auto _M_find_pool(size_t) noexcept;

    __pool_resource _M_impl;
    _Pool* _M_pools = nullptr;
  };

  class monotonic_buffer_resource : public memory_resource
  {
  public:
    explicit
    monotonic_buffer_resource(memory_resource* __upstream) noexcept
    __attribute__((__nonnull__))
    : _M_upstream(__upstream)
    { _GLIBCXX_DEBUG_ASSERT(__upstream != nullptr); }

    monotonic_buffer_resource(size_t __initial_size,
			      memory_resource* __upstream) noexcept
    __attribute__((__nonnull__))
    : _M_next_bufsiz(__initial_size),
      _M_upstream(__upstream)
    {
      _GLIBCXX_DEBUG_ASSERT(__upstream != nullptr);
      _GLIBCXX_DEBUG_ASSERT(__initial_size > 0);
    }

    monotonic_buffer_resource(void* __buffer, size_t __buffer_size,
			      memory_resource* __upstream) noexcept
    __attribute__((__nonnull__(4)))
    : _M_current_buf(__buffer), _M_avail(__buffer_size),
      _M_next_bufsiz(_S_next_bufsize(__buffer_size)),
      _M_upstream(__upstream),
      _M_orig_buf(__buffer), _M_orig_size(__buffer_size)
    {
      _GLIBCXX_DEBUG_ASSERT(__upstream != nullptr);
      _GLIBCXX_DEBUG_ASSERT(__buffer != nullptr || __buffer_size == 0);
    }

    monotonic_buffer_resource() noexcept
    : monotonic_buffer_resource(get_default_resource())
    { }

    explicit
    monotonic_buffer_resource(size_t __initial_size) noexcept
    : monotonic_buffer_resource(__initial_size, get_default_resource())
    { }

    monotonic_buffer_resource(void* __buffer, size_t __buffer_size) noexcept
    : monotonic_buffer_resource(__buffer, __buffer_size, get_default_resource())
    { }

    monotonic_buffer_resource(const monotonic_buffer_resource&) = delete;

    virtual ~monotonic_buffer_resource(); // key function

    monotonic_buffer_resource&
    operator=(const monotonic_buffer_resource&) = delete;

    void
    release() noexcept
    {
      if (_M_head)
	_M_release_buffers();

      // reset to initial state at contruction:
      if ((_M_current_buf = _M_orig_buf))
	{
	  _M_avail = _M_orig_size;
	  _M_next_bufsiz = _S_next_bufsize(_M_orig_size);
	}
      else
	{
	  _M_avail = 0;
	  _M_next_bufsiz = _M_orig_size;
	}
    }

    memory_resource*
    upstream_resource() const noexcept
    __attribute__((__returns_nonnull__))
    { return _M_upstream; }

  protected:
    void*
    do_allocate(size_t __bytes, size_t __alignment) override
    {
      if (__builtin_expect(__bytes == 0, false))
	__bytes = 1; // Ensures we don't return the same pointer twice.

      void* __p = std::align(__alignment, __bytes, _M_current_buf, _M_avail);
      if (__builtin_expect(__p == nullptr, false))
	{
	  _M_new_buffer(__bytes, __alignment);
	  __p = _M_current_buf;
	}
      _M_current_buf = (char*)_M_current_buf + __bytes;
      _M_avail -= __bytes;
      return __p;
    }

    void
    do_deallocate(void*, size_t, size_t) override
    { }

    bool
    do_is_equal(const memory_resource& __other) const noexcept override
    { return this == &__other; }

  private:
    // Update _M_current_buf and _M_avail to refer to a new buffer with
    // at least the specified size and alignment, allocated from upstream.
    void
    _M_new_buffer(size_t __bytes, size_t __alignment);

    // Deallocate all buffers obtained from upstream.
    void
    _M_release_buffers() noexcept;

    static size_t
    _S_next_bufsize(size_t __buffer_size) noexcept
    {
      if (__builtin_expect(__buffer_size == 0, false))
	__buffer_size = 1;
      return __buffer_size * _S_growth_factor;
    }

    static constexpr size_t _S_init_bufsize = 128 * sizeof(void*);
    static constexpr float _S_growth_factor = 1.5;

    void*	_M_current_buf = nullptr;
    size_t	_M_avail = 0;
    size_t	_M_next_bufsiz = _S_init_bufsize;

    // Initial values set at construction and reused by release():
    memory_resource* const	_M_upstream;
    void* const			_M_orig_buf = nullptr;
    size_t const		_M_orig_size = _M_next_bufsiz;

    class _Chunk;
    _Chunk* _M_head = nullptr;
  };

} // namespace pmr
_GLIBCXX_END_NAMESPACE_VERSION
} // namespace std

#endif // C++17
#endif // _GLIBCXX_MEMORY_RESOURCE
