//===- llvm/ADT/SmallVector.h - 'Normally small' vectors --------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file defines the SmallVector class.
//
//===----------------------------------------------------------------------===//

// ATen: modified from llvm::SmallVector.
// used std::is_trivially_{copy,move}_constructible
// replaced iterator_range constructor with inline Container&& constructor
// replaced LLVM_NODISCARD, LLVM_LIKELY, and LLVM_UNLIKELY with c10 equivalents
// removed LLVM_GSL_OWNER
// added SmallVector::at
// added operator<< for std::ostream
// added C10_API to export SmallVectorBase

#pragma once

#include <c10/macros/Macros.h>
#include <c10/util/AlignOf.h>

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <initializer_list>
#include <iterator>
#include <limits>
#include <memory>
#include <new>
#include <ostream>
#include <type_traits>
#include <utility>

C10_CLANG_DIAGNOSTIC_PUSH()
#if C10_CLANG_HAS_WARNING("-Wshorten-64-to-32")
C10_CLANG_DIAGNOSTIC_IGNORE("-Wshorten-64-to-32")
#endif

namespace c10 {

/// This is all the stuff common to all SmallVectors.
///
/// The template parameter specifies the type which should be used to hold the
/// Size and Capacity of the SmallVector, so it can be adjusted.
/// Using 32 bit size is desirable to shrink the size of the SmallVector.
/// Using 64 bit size is desirable for cases like SmallVector<char>, where a
/// 32 bit size would limit the vector to ~4GB. SmallVectors are used for
/// buffering bitcode output - which can exceed 4GB.
template <class Size_T>
class C10_API SmallVectorBase {
 protected:
  void* BeginX;
  Size_T Size = 0, Capacity;

  /// The maximum value of the Size_T used.
  static constexpr size_t SizeTypeMax() {
    return std::numeric_limits<Size_T>::max();
  }

  SmallVectorBase(void* FirstEl, size_t TotalCapacity)
      : BeginX(FirstEl), Capacity(TotalCapacity) {}

  /// This is a helper for \a grow() that's out of line to reduce code
  /// duplication.  This function will report a fatal error if it can't grow at
  /// least to \p MinSize.
  void* mallocForGrow(size_t MinSize, size_t TSize, size_t& NewCapacity);

  /// This is an implementation of the grow() method which only works
  /// on POD-like data types and is out of line to reduce code duplication.
  /// This function will report a fatal error if it cannot increase capacity.
  void grow_pod(void* FirstEl, size_t MinSize, size_t TSize);

 public:
  SmallVectorBase() = delete;
  size_t size() const {
    return Size;
  }
  size_t capacity() const {
    return Capacity;
  }

  C10_NODISCARD bool empty() const {
    return !Size;
  }

  /// Set the array size to \p N, which the current array must have enough
  /// capacity for.
  ///
  /// This does not construct or destroy any elements in the vector.
  ///
  /// Clients can use this in conjunction with capacity() to write past the end
  /// of the buffer when they know that more elements are available, and only
  /// update the size later. This avoids the cost of value initializing elements
  /// which will only be overwritten.
  void set_size(size_t N) {
    assert(N <= capacity());
    Size = N;
  }
};

template <class T>
using SmallVectorSizeType = typename std::
    conditional<sizeof(T) < 4 && sizeof(void*) >= 8, uint64_t, uint32_t>::type;

/// Figure out the offset of the first element.
template <class T, typename = void>
struct SmallVectorAlignmentAndSize {
  alignas(SmallVectorBase<SmallVectorSizeType<T>>) char Base[sizeof(
      SmallVectorBase<SmallVectorSizeType<T>>)];
  alignas(T) char FirstEl[sizeof(T)];
};

/// This is the part of SmallVectorTemplateBase which does not depend on whether
/// the type T is a POD. The extra dummy template argument is used by ArrayRef
/// to avoid unnecessarily requiring T to be complete.
template <typename T, typename = void>
class SmallVectorTemplateCommon
    : public SmallVectorBase<SmallVectorSizeType<T>> {
  using Base = SmallVectorBase<SmallVectorSizeType<T>>;

  /// Find the address of the first element.  For this pointer math to be valid
  /// with small-size of 0 for T with lots of alignment, it's important that
  /// SmallVectorStorage is properly-aligned even for small-size of 0.
  void* getFirstEl() const {
    return const_cast<void*>(reinterpret_cast<const void*>(
        reinterpret_cast<const char*>(this) +
        offsetof(SmallVectorAlignmentAndSize<T>, FirstEl)));
  }
  // Space after 'FirstEl' is clobbered, do not add any instance vars after it.

 protected:
  SmallVectorTemplateCommon(size_t Size) : Base(getFirstEl(), Size) {}

  void grow_pod(size_t MinSize, size_t TSize) {
    Base::grow_pod(getFirstEl(), MinSize, TSize);
  }

  /// Return true if this is a smallvector which has not had dynamic
  /// memory allocated for it.
  bool isSmall() const {
    return this->BeginX == getFirstEl();
  }

  /// Put this vector in a state of being small.
  void resetToSmall() {
    this->BeginX = getFirstEl();
    this->Size = this->Capacity = 0; // FIXME: Setting Capacity to 0 is suspect.
  }

  /// Return true if V is an internal reference to the given range.
  bool isReferenceToRange(const void* V, const void* First, const void* Last)
      const {
    // Use std::less to avoid UB.
    std::less<> LessThan;
    return !LessThan(V, First) && LessThan(V, Last);
  }

  /// Return true if V is an internal reference to this vector.
  bool isReferenceToStorage(const void* V) const {
    return isReferenceToRange(V, this->begin(), this->end());
  }

  /// Return true if First and Last form a valid (possibly empty) range in this
  /// vector's storage.
  bool isRangeInStorage(const void* First, const void* Last) const {
    // Use std::less to avoid UB.
    std::less<> LessThan;
    return !LessThan(First, this->begin()) && !LessThan(Last, First) &&
        !LessThan(this->end(), Last);
  }

  /// Return true unless Elt will be invalidated by resizing the vector to
  /// NewSize.
  bool isSafeToReferenceAfterResize(const void* Elt, size_t NewSize) {
    // Past the end.
    if (C10_LIKELY(!isReferenceToStorage(Elt)))
      return true;

    // Return false if Elt will be destroyed by shrinking.
    if (NewSize <= this->size())
      return Elt < this->begin() + NewSize;

    // Return false if we need to grow.
    return NewSize <= this->capacity();
  }

  /// Check whether Elt will be invalidated by resizing the vector to NewSize.
  void assertSafeToReferenceAfterResize(const void* Elt, size_t NewSize) {
    (void)Elt; // Suppress unused variable warning
    (void)NewSize; // Suppress unused variable warning
    assert(
        isSafeToReferenceAfterResize(Elt, NewSize) &&
        "Attempting to reference an element of the vector in an operation "
        "that invalidates it");
  }

  /// Check whether Elt will be invalidated by increasing the size of the
  /// vector by N.
  void assertSafeToAdd(const void* Elt, size_t N = 1) {
    this->assertSafeToReferenceAfterResize(Elt, this->size() + N);
  }

  /// Check whether any part of the range will be invalidated by clearing.
  void assertSafeToReferenceAfterClear(const T* From, const T* To) {
    if (From == To)
      return;
    this->assertSafeToReferenceAfterResize(From, 0);
    this->assertSafeToReferenceAfterResize(To - 1, 0);
  }
  template <
      class ItTy,
      std::enable_if_t<
          !std::is_same<std::remove_const_t<ItTy>, T*>::value,
          bool> = false>
  void assertSafeToReferenceAfterClear(ItTy, ItTy) {}

  /// Check whether any part of the range will be invalidated by growing.
  void assertSafeToAddRange(const T* From, const T* To) {
    if (From == To)
      return;
    this->assertSafeToAdd(From, To - From);
    this->assertSafeToAdd(To - 1, To - From);
  }
  template <
      class ItTy,
      std::enable_if_t<
          !std::is_same<std::remove_const_t<ItTy>, T*>::value,
          bool> = false>
  void assertSafeToAddRange(ItTy, ItTy) {}

  /// Reserve enough space to add one element, and return the updated element
  /// pointer in case it was a reference to the storage.
  template <class U>
  static const T* reserveForParamAndGetAddressImpl(
      U* This,
      const T& Elt,
      size_t N) {
    size_t NewSize = This->size() + N;
    if (C10_LIKELY(NewSize <= This->capacity()))
      return &Elt;

    bool ReferencesStorage = false;
    int64_t Index = -1;
    if (!U::TakesParamByValue) {
      if (C10_UNLIKELY(This->isReferenceToStorage(&Elt))) {
        ReferencesStorage = true;
        Index = &Elt - This->begin();
      }
    }
    This->grow(NewSize);
    return ReferencesStorage ? This->begin() + Index : &Elt;
  }

 public:
  using size_type = size_t;
  using difference_type = ptrdiff_t;
  using value_type = T;
  using iterator = T*;
  using const_iterator = const T*;

  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
  using reverse_iterator = std::reverse_iterator<iterator>;

  using reference = T&;
  using const_reference = const T&;
  using pointer = T*;
  using const_pointer = const T*;

  using Base::capacity;
  using Base::empty;
  using Base::size;

  // forward iterator creation methods.
  iterator begin() {
    return (iterator)this->BeginX;
  }
  const_iterator begin() const {
    return (const_iterator)this->BeginX;
  }
  iterator end() {
    return begin() + size();
  }
  const_iterator end() const {
    return begin() + size();
  }

  // reverse iterator creation methods.
  reverse_iterator rbegin() {
    return reverse_iterator(end());
  }
  const_reverse_iterator rbegin() const {
    return const_reverse_iterator(end());
  }
  reverse_iterator rend() {
    return reverse_iterator(begin());
  }
  const_reverse_iterator rend() const {
    return const_reverse_iterator(begin());
  }

  size_type size_in_bytes() const {
    return size() * sizeof(T);
  }
  size_type max_size() const {
    return std::min(this->SizeTypeMax(), size_type(-1) / sizeof(T));
  }

  size_t capacity_in_bytes() const {
    return capacity() * sizeof(T);
  }

  /// Return a pointer to the vector's buffer, even if empty().
  pointer data() {
    return pointer(begin());
  }
  /// Return a pointer to the vector's buffer, even if empty().
  const_pointer data() const {
    return const_pointer(begin());
  }

  // SmallVector::at is NOT from LLVM.
  reference at(size_type idx) {
    assert(idx < size());
    return begin()[idx];
  }
  const_reference at(size_type idx) const {
    assert(idx < size());
    return begin()[idx];
  }
  reference operator[](size_type idx) {
    assert(idx < size());
    return begin()[idx];
  }
  const_reference operator[](size_type idx) const {
    assert(idx < size());
    return begin()[idx];
  }