// Copyright (c) 2012 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef BASE_CONTAINERS_STACK_CONTAINER_H_
#define BASE_CONTAINERS_STACK_CONTAINER_H_

#include <stddef.h>
#include <memory>
#include <vector>

#include "base/compiler_specific.h"
#include "base/memory/raw_ptr_exclusion.h"
#include "build/build_config.h"

namespace base {

// This allocator can be used with STL containers to provide a stack buffer
// from which to allocate memory and overflows onto the heap. This stack buffer
// would be allocated on the stack and allows us to avoid heap operations in
// some situations.
//
// STL likes to make copies of allocators, so the allocator itself can't hold
// the data. Instead, we make the creator responsible for creating a
// StackAllocator::Source which contains the data. Copying the allocator
// merely copies the pointer to this shared source, so all allocators created
// based on our allocator will share the same stack buffer.
//
// This stack buffer implementation is very simple. The first allocation that
// fits in the stack buffer will use the stack buffer. Any subsequent
// allocations will not use the stack buffer, even if there is unused room.
// This makes it appropriate for array-like containers, but the caller should
// be sure to reserve() in the container up to the stack buffer size. Otherwise
// the container will allocate a small array which will "use up" the stack
// buffer.
template <typename T, size_t stack_capacity, typename FallbackAllocator>
class StackAllocator : public FallbackAllocator {
 public:
  using pointer = typename std::allocator_traits<FallbackAllocator>::pointer;
  using size_type =
      typename std::allocator_traits<FallbackAllocator>::size_type;

  // Backing store for the allocator. The container owner is responsible for
  // maintaining this for as long as any containers using this allocator are
  // live.
  struct Source {
    Source() : used_stack_buffer_(false) {
    }

    // Casts the buffer in its right type.
    NO_SANITIZE("cfi-unrelated-cast")
    T* stack_buffer() { return reinterpret_cast<T*>(stack_buffer_); }
    NO_SANITIZE("cfi-unrelated-cast")
    const T* stack_buffer() const {
      return reinterpret_cast<const T*>(&stack_buffer_);
    }

    // The buffer itself. It is not of type T because we don't want the
    // constructors and destructors to be automatically called. Define a POD
    // buffer of the right size instead.
    alignas(T) char stack_buffer_[sizeof(T[stack_capacity])];
#if defined(__GNUC__) && !defined(ARCH_CPU_X86_FAMILY)
    static_assert(alignof(T) <= 16, "http://crbug.com/115612");
#endif

    // Set when the stack buffer is used for an allocation. We do not track
    // how much of the buffer is used, only that somebody is using it.
    bool used_stack_buffer_;
  };

  // Used by containers when they want to refer to an allocator of type U.
  template<typename U>
  struct rebind {
    typedef StackAllocator<U, stack_capacity, FallbackAllocator> other;
  };

  // For the straight up copy c-tor, we can share storage.
  StackAllocator(
      const StackAllocator<T, stack_capacity, FallbackAllocator>& rhs)
      : source_(rhs.source_) {}

  // ISO C++ requires the following constructor to be defined,
  // and std::vector in VC++2008SP1 Release fails with an error
  // in the class _Container_base_aux_alloc_real (from <xutility>)
  // if the constructor does not exist.
  // For this constructor, we cannot share storage; there's
  // no guarantee that the Source buffer of Ts is large enough
  // for Us.
  // TODO: If we were fancy pants, perhaps we could share storage
  // iff sizeof(T) == sizeof(U).
  template <typename U, size_t other_capacity, typename FA>
  StackAllocator(const StackAllocator<U, other_capacity, FA>& other)
      : source_(nullptr) {}

  // This constructor must exist. It creates a default allocator that doesn't
  // actually have a stack buffer. glibc's std::string() will compare the
  // current allocator against the default-constructed allocator, so this
  // should be fast.
  StackAllocator() : source_(nullptr) {}

  explicit StackAllocator(Source* source) : source_(source) {
  }

  // Actually do the allocation. Use the stack buffer if nobody has used it yet
  // and the size requested fits. Otherwise, fall through to the standard
  // allocator.
  pointer allocate(size_type n) {
    if (source_ && !source_->used_stack_buffer_ && n <= stack_capacity) {
      source_->used_stack_buffer_ = true;
      return source_->stack_buffer();
    } else {
      return std::allocator_traits<FallbackAllocator>::allocate(*this, n);
    }
  }

  // Free: when trying to free the stack buffer, just mark it as free. For
  // non-stack-buffer pointers, just fall though to the standard allocator.
  void deallocate(pointer p, size_type n) {
    if (source_ && p == source_->stack_buffer())
      source_->used_stack_buffer_ = false;
    else
      std::allocator_traits<FallbackAllocator>::deallocate(*this, p, n);
  }

 private:
  // `source_` is not a raw_ptr<T> for performance reasons: on-stack pointee.
  RAW_PTR_EXCLUSION Source* source_;
};

// A wrapper around STL containers that maintains a stack-sized buffer that the
// initial capacity of the vector is based on. Growing the container beyond the
// stack capacity will transparently overflow onto the heap. The container must
// support reserve().
//
// This will not work with std::string since some implementations allocate
// more bytes than requested in calls to reserve(), forcing the allocation onto
// the heap.  http://crbug.com/709273
//
// WATCH OUT: the ContainerType MUST use the proper StackAllocator for this
// type. This object is really intended to be used only internally. You'll want
// to use the wrappers below for different types.
template <typename TContainerType, int stack_capacity>
class StackContainer {
 public:
  using ContainerType = TContainerType;
  using ContainedType = typename ContainerType::value_type;
  using Allocator = typename ContainerType::allocator_type;

  // Allocator must be constructed before the container!
  StackContainer() : allocator_(&stack_data_), container_(allocator_) {
    // Make the container use the stack allocation by reserving our buffer size
    // before doing anything else.
    container_.reserve(stack_capacity);
  }
  StackContainer(const StackContainer&) = delete;
  StackContainer& operator=(const StackContainer&) = delete;

  // Getters for the actual container.
  //
  // Danger: any copies of this made using the copy constructor must have
  // shorter lifetimes than the source. The copy will share the same allocator
  // and therefore the same stack buffer as the original. Use std::copy to
  // copy into a "real" container for longer-lived objects.
  ContainerType& container() { return container_; }
  const ContainerType& container() const { return container_; }

  // Support operator-> to get to the container. This allows nicer syntax like:
  //   StackContainer<...> foo;
  //   std::sort(foo->begin(), foo->end());
  ContainerType* operator->() { return &container_; }
  const ContainerType* operator->() const { return &container_; }

#ifdef UNIT_TEST
  // Retrieves the stack source so that that unit tests can verify that the
  // buffer is being used properly.
  const typename Allocator::Source& stack_data() const {
    return stack_data_;
  }
#endif

 protected:
  typename Allocator::Source stack_data_;
  NO_UNIQUE_ADDRESS Allocator allocator_;
  ContainerType container_;
};

// Range-based iteration support for StackContainer.
template <typename TContainerType, int stack_capacity>
auto begin(
    const StackContainer<TContainerType, stack_capacity>& stack_container)
    -> decltype(begin(stack_container.container())) {
  return begin(stack_container.container());
}

template <typename TContainerType, int stack_capacity>
auto begin(StackContainer<TContainerType, stack_capacity>& stack_container)
    -> decltype(begin(stack_container.container())) {
  return begin(stack_container.container());
}

template <typename TContainerType, int stack_capacity>
auto end(StackContainer<TContainerType, stack_capacity>& stack_container)
    -> decltype(end(stack_container.container())) {
  return end(stack_container.container());
}

template <typename TContainerType, int stack_capacity>
auto end(const StackContainer<TContainerType, stack_capacity>& stack_container)
    -> decltype(end(stack_container.container())) {
  return end(stack_container.container());
}

// StackVector -----------------------------------------------------------------

// Example:
//   StackVector<int, 16> foo;
//   foo->push_back(22);  // we have overloaded operator->
//   foo[0] = 10;         // as well as operator[]
template <typename T,
          size_t stack_capacity,
          typename FallbackAllocator = std::allocator<T>>
class StackVector
    : public StackContainer<
          std::vector<T, StackAllocator<T, stack_capacity, FallbackAllocator>>,
          stack_capacity> {
 public:
  StackVector()
      : StackContainer<
            std::vector<T,
                        StackAllocator<T, stack_capacity, FallbackAllocator>>,
            stack_capacity>() {}

  // We need to put this in STL containers sometimes, which requires a copy
  // constructor. We can't call the regular copy constructor because that will
  // take the stack buffer from the original. Here, we create an empty object
  // and make a stack buffer of its own.
  StackVector(const StackVector<T, stack_capacity, FallbackAllocator>& other)
      : StackContainer<
            std::vector<T,
                        StackAllocator<T, stack_capacity, FallbackAllocator>>,
            stack_capacity>() {
    this->container().assign(other->begin(), other->end());
  }

  StackVector<T, stack_capacity, FallbackAllocator>& operator=(
      const StackVector<T, stack_capacity, FallbackAllocator>& other) {
    this->container().assign(other->begin(), other->end());
    return *this;
  }

  // Vectors are commonly indexed, which isn't very convenient even with
  // operator-> (using "->at()" does exception stuff we don't want).
  T& operator[](size_t i) { return this->container().operator[](i); }
  const T& operator[](size_t i) const {
    return this->container().operator[](i);
  }
};

}  // namespace base

#endif  // BASE_CONTAINERS_STACK_CONTAINER_H_