// Copyright 2019 The TCMalloc 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.

#ifndef TCMALLOC_INTERNAL_RANGE_TRACKER_H_
#define TCMALLOC_INTERNAL_RANGE_TRACKER_H_

#include <stddef.h>
#include <stdint.h>
#include <sys/types.h>

#include <algorithm>
#include <limits>
#include <type_traits>

#include "absl/numeric/bits.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/optimization.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace tcmalloc_internal {

// Keeps a bitmap of some fixed size (N bits).
template <size_t N>
class Bitmap {
 public:
  constexpr Bitmap() : bits_{} {}

  size_t size() const { return N; }
  bool GetBit(size_t i) const;

  // Sets the bit at index `i`.  Returns the previous value of that bit.
  bool SetBit(size_t i);
  // Clears the bit at index `i`.  Returns the previous value of that bit.
  bool ClearBit(size_t i);

  // Returns the number of set bits [index, ..., index + n - 1].
  size_t CountBits(size_t index, size_t n) const;
  // Returns the total number of set bits.
  size_t CountBits() const;

  // Returns whether the bitmap is entirely zero or not.
  bool IsZero() const;

  // Equivalent to SetBit on bits [index, index + 1, ... index + n - 1].
  void SetRange(size_t index, size_t n);
  void ClearRange(size_t index, size_t n);

  // Clears the lowest set bit. Special case is faster than more flexible code.
  void ClearLowestBit();

  template <typename Visitor>
  size_t PopBatch(Visitor visitor, size_t limit) {
    size_t count = 0;
    for (size_t i = 0; i < kWords; ++i) {
      if (count >= limit) break;
      size_t word = bits_[i];
      while (word != 0 && count < limit) {
        size_t off = absl::countr_zero(word);
        visitor(i * kWordSize + off);
        count++;
        word &= word - 1;
      }
      bits_[i] = word;
    }
    return count;
  }

  // If there is at least one free range at or after <start>,
  // put it in *index, *length and return true; else return false.
  bool NextFreeRange(size_t start, size_t* index, size_t* length) const;

  // Returns index of the first {true, false} bit >= index, or N if none.
  size_t FindSet(size_t index) const;
  size_t FindClear(size_t index) const;

  // Returns index of the first {set, clear} bit in [index, 0] or -1 if none.
  ssize_t FindSetBackwards(size_t index) const;
  ssize_t FindClearBackwards(size_t index) const;

  void Clear();

  // Bitwise operators.
  Bitmap<N> operator~() const;
  Bitmap<N> operator&(const Bitmap<N>& other) const;
  Bitmap<N> operator|(const Bitmap<N>& other) const;
  Bitmap<N> operator^(const Bitmap<N>& other) const;

 private:
  static constexpr size_t kWordSize = sizeof(size_t) * 8;
  static constexpr size_t kWords = (N + kWordSize - 1) / kWordSize;
  static constexpr size_t kDeadBits = kWordSize * kWords - N;

  size_t bits_[kWords];

  size_t CountWordBits(size_t i, size_t from, size_t to) const;

  template <bool Value>
  void SetWordBits(size_t i, size_t from, size_t to);
  template <bool Value>
  void SetRangeValue(size_t index, size_t n);

  template <bool Goal>
  size_t FindValue(size_t index) const;
  template <bool Goal>
  ssize_t FindValueBackwards(size_t index) const;
};

constexpr size_t kMaxResidencyBits = 512;
using ResidencyBitmap = Bitmap<kMaxResidencyBits>;
// Tracks allocations in a range of items of fixed size.  Supports
// finding an unset range of a given length, while keeping track of
// the largest remaining unmarked length.
template <size_t N>
class RangeTracker {
 public:
  constexpr RangeTracker()
      : bits_{}, longest_free_(N), nused_(0), nallocs_(0) {}

  size_t size() const;
  // Number of bits marked
  size_t used() const;
  // Longest contiguous range of clear bits.
  size_t longest_free() const;
  // Count of live allocations.
  size_t allocs() const;

  // REQUIRES: there is a free range of at least n bits
  // (i.e. n <= longest_free()).
  //
  // Finds and marks n free bits, returning index of the first bit.  Chooses by
  // best fit.
  size_t FindAndMark(size_t n);

  // REQUIRES: the range [index, index + n) is fully unmarked.
  void Mark(size_t index, size_t n);

  // REQUIRES: the range [index, index + n) is fully marked, and
  // was the returned value from a call to FindAndMark.
  // Unmarks it.
  void Unmark(size_t index, size_t n);
  // If there is at least one free range at or after <start>,
  // put it in *index, *length and return true; else return false.
  bool NextFreeRange(size_t start, size_t* index, size_t* length) const;

  void Clear();

  Bitmap<N> bits() const;

 private:
  Bitmap<N> bits_;

  // Computes the smallest unsigned type that can hold the constant N.
  class UnsignedTypeFittingSize {
   private:
    static_assert(N <= std::numeric_limits<uint64_t>::max(),
                  "size_t more than 64 bits??");
    template <typename T>
    static constexpr bool Fit() {
      return N <= std::numeric_limits<T>::max();
    }
    struct U32 {
      using type = std::conditional_t<Fit<uint32_t>(), uint32_t, uint64_t>;
    };

    struct U16 {
      using type =
          std::conditional_t<Fit<uint16_t>(), uint16_t, typename U32::type>;
    };

    struct U8 {
      using type =
          std::conditional_t<Fit<uint8_t>(), uint8_t, typename U16::type>;
    };

   public:
    using type = typename U8::type;
  };

  // we keep various stats in the range [0, N]; make them as small as possible.
  using Count = typename UnsignedTypeFittingSize::type;

  Count longest_free_;
  Count nused_;
  Count nallocs_;
};

template <size_t N>
inline size_t RangeTracker<N>::size() const {
  return bits_.size();
}

template <size_t N>
inline size_t RangeTracker<N>::used() const {
  return nused_;
}

template <size_t N>
inline size_t RangeTracker<N>::longest_free() const {
  return longest_free_;
}

template <size_t N>
inline size_t RangeTracker<N>::allocs() const {
  return nallocs_;
}

template <size_t N>
inline size_t RangeTracker<N>::FindAndMark(size_t n) {
  TC_ASSERT_GT(n, 0);

  // We keep the two longest ranges in the bitmap since we might allocate
  // from one.
  size_t longest_len = 0;
  size_t second_len = 0;

  // the best (shortest) range we could use
  // TODO(b/134691947): shortest? lowest-addressed?
  size_t best_index = N;
  size_t best_len = 2 * N;
  // Iterate over free ranges:
  size_t index = 0, len;

  while (bits_.NextFreeRange(index, &index, &len)) {
    if (len > longest_len) {
      second_len = longest_len;
      longest_len = len;
    } else if (len > second_len) {
      second_len = len;
    }

    if (len >= n && len < best_len) {
      best_index = index;
      best_len = len;
    }

    index += len;
  }

  TC_CHECK_LT(best_index, N);
  bits_.SetRange(best_index, n);

  if (best_len == longest_len) {
    longest_len -= n;
    if (longest_len < second_len) longest_len = second_len;
  }

  longest_free_ = longest_len;
  nused_ += n;
  nallocs_++;
  return best_index;
}

// REQUIRES: the range [index, index + n) is fully unmarked.
// Marks it.
template <size_t N>
inline void RangeTracker<N>::Mark(size_t index, size_t n) {
  TC_ASSERT_GE(bits_.FindSet(index), index + n);
  bits_.SetRange(index, n);
  nused_ += n;
  nallocs_++;

  size_t longest_len = 0;
  size_t scan_index = 0, scan_len;

  // We just marked a range as used. This might change the longest free range
  // recorded in longest_free_. Recompute.
  while (bits_.NextFreeRange(scan_index, &scan_index, &scan_len)) {
    if (scan_len > longest_len) {
      longest_len = scan_len;
    }

    scan_index += scan_len;
  }

  longest_free_ = longest_len;
}

// REQUIRES: the range [index, index + n) is fully marked.
// Unmarks it.
template <size_t N>
inline void RangeTracker<N>::Unmark(size_t index, size_t n) {
  TC_ASSERT(bits_.FindClear(index) >= index + n);
  bits_.ClearRange(index, n);
  nused_ -= n;
  nallocs_--;

  // We just opened up a new free range--it might be the longest.
  size_t lim = bits_.FindSet(index + n - 1);
  index = bits_.FindSetBackwards(index) + 1;
  n = lim - index;
  if (n > longest_free()) {
    longest_free_ = n;
  }
}

template <size_t N>
inline Bitmap<N> RangeTracker<N>::bits() const {
  return bits_;
}

// If there is at least one free range at or after <start>,
// put it in *index, *length and return true; else return false.
template <size_t N>
inline bool RangeTracker<N>::NextFreeRange(size_t start, size_t* index,
                                           size_t* length) const {
  return bits_.NextFreeRange(start, index, length);
}

template <size_t N>
inline void RangeTracker<N>::Clear() {
  bits_.Clear();
  nallocs_ = 0;
  nused_ = 0;
  longest_free_ = N;
}

// Count the set bits [from, to) in the i-th word to Value.
template <size_t N>
inline size_t Bitmap<N>::CountWordBits(size_t i, size_t from, size_t to) const {
  TC_ASSERT_LT(from, kWordSize);
  TC_ASSERT_LE(to, kWordSize);
  const size_t all_ones = ~static_cast<size_t>(0);
  // how many bits are we setting?
  const size_t n = to - from;
  TC_ASSERT(0 < n && n <= kWordSize);
  const size_t mask = (all_ones >> (kWordSize - n)) << from;

  ASSUME(i < kWords);
  return absl::popcount(bits_[i] & mask);
}

// Set the bits [from, to) in the i-th word to Value.
template <size_t N>
template <bool Value>
inline void Bitmap<N>::SetWordBits(size_t i, size_t from, size_t to) {
  TC_ASSERT_LT(from, kWordSize);
  TC_ASSERT_LE(to, kWordSize);
  const size_t all_ones = ~static_cast<size_t>(0);
  // how many bits are we setting?
  const size_t n = to - from;
  TC_ASSERT(n > 0 && n <= kWordSize);
  const size_t mask = (all_ones >> (kWordSize - n)) << from;
  ASSUME(i < kWords);
  if (Value) {
    bits_[i] |= mask;
  } else {
    bits_[i] &= ~mask;
  }
}

template <size_t N>
inline bool Bitmap<N>::GetBit(size_t i) const {
  TC_ASSERT_LT(i, N);
  size_t word = i / kWordSize;
  size_t offset = i % kWordSize;
  ASSUME(word < kWords);
  return bits_[word] & (size_t{1} << offset);
}

template <size_t N>
inline bool Bitmap<N>::SetBit(size_t i) {
  TC_ASSERT_LT(i, N);
  size_t word = i / kWordSize;
  size_t offset = i % kWordSize;
  ASSUME(word < kWords);
  size_t mask = size_t{1} << offset;
  bool old = bits_[word] & mask;
  bits_[word] |= mask;
  return old;
}

template <size_t N>
inline bool Bitmap<N>::ClearBit(size_t i) {
  TC_ASSERT_LT(i, N);
  size_t word = i / kWordSize;
  size_t offset = i % kWordSize;
  ASSUME(word < kWords);
  size_t mask = size_t{1} << offset;
  bool old = bits_[word] & mask;
  bits_[word] &= ~mask;
  return old;
}

template <size_t N>
inline size_t Bitmap<N>::CountBits() const {
  return CountBits(0, N);
}

template <size_t N>
inline size_t Bitmap<N>::CountBits(size_t index, size_t n) const {
  ASSUME(index + n <= N);
  size_t count = 0;
  if (n == 0) {
    return count;
  }

  size_t word = index / kWordSize;
  size_t offset = index % kWordSize;
  size_t k = std::min(offset + n, kWordSize);
  count += CountWordBits(word, offset, k);
  n -= k - offset;
  while (n > 0) {
    word++;
    k = std::min(n, kWordSize);
    count += CountWordBits(word, 0, k);
    n -= k;
  }

  return count;
}

template <size_t N>
inline bool Bitmap<N>::IsZero() const {
  for (int i = 0; i < kWords; ++i) {
    if (bits_[i] != 0) {
      return false;
    }
  }
  return true;
}

template <size_t N>
inline void Bitmap<N>::SetRange(size_t index, size_t n) {
  SetRangeValue<true>(index, n);
}

template <size_t N>
inline void Bitmap<N>::ClearRange(size_t index, size_t n) {
  SetRangeValue<false>(index, n);
}

template <size_t N>
inline void Bitmap<N>::ClearLowestBit() {
  for (int i = 0; i < kWords; ++i) {
    if (bits_[i] != 0) {
      bits_[i] &= bits_[i] - 1;
      break;
    }
  }
}

template <size_t N>
template <bool Value>
inline void Bitmap<N>::SetRangeValue(size_t index, size_t n) {
  TC_ASSERT_LE(index + n, N);
  size_t word = index / kWordSize;
  size_t offset = index % kWordSize;
  size_t k = offset + n;
  if (k > kWordSize) k = kWordSize;
  SetWordBits<Value>(word, offset, k);
  n -= k - offset;
  while (n > 0) {
    word++;
    k = n;
    if (k > kWordSize) k = kWordSize;
    SetWordBits<Value>(word, 0, k);
    n -= k;
  }
}

template <size_t N>
inline bool Bitmap<N>::NextFreeRange(size_t start, size_t* index,
                                     size_t* length) const {
  if (start >= N) return false;
  size_t i = FindClear(start);
  if (i == N) return false;
  size_t j = FindSet(i);
  *index = i;
  *length = j - i;
  return true;
}

template <size_t N>
inline size_t Bitmap<N>::FindSet(size_t index) const {
  return FindValue<true>(index);
}

template <size_t N>
inline size_t Bitmap<N>::FindClear(size_t index) const {
  return FindValue<false>(index);
}

template <size_t N>
inline ssize_t Bitmap<N>::FindSetBackwards(size_t index) const {
  return FindValueBackwards<true>(index);
}

template <size_t N>
inline ssize_t Bitmap<N>::FindClearBackwards(size_t index) const {
  return FindValueBackwards<false>(index);
}

template <size_t N>
inline void Bitmap<N>::Clear() {
  for (int i = 0; i < kWords; ++i) {
    bits_[i] = 0;
  }
}

// Bitwise ~ requires special handling of the dead bits.
template <size_t N>
inline Bitmap<N> Bitmap<N>::operator~() const {
  Bitmap<N> result;
  for (size_t i = 0; i < kWords; ++i) {
    result.bits_[i] = ~bits_[i];
  }
  if constexpr (kDeadBits > 0) {
    result.bits_[kWords - 1] &= (~static_cast<size_t>(0)) >> kDeadBits;
  }
  return result;
}

// Bitwise (&, |, ^) do not require special handling of the dead bits, since
// 0 & 0 = 0, 0 | 0 = 0, and 0 ^ 0 = 0.
template <size_t N>
inline Bitmap<N> Bitmap<N>::operator&(const Bitmap<N>& other) const {
  Bitmap<N> result;
  for (size_t i = 0; i < kWords; ++i) {
    result.bits_[i] = bits_[i] & other.bits_[i];
  }
  return result;
}

template <size_t N>
inline Bitmap<N> Bitmap<N>::operator|(const Bitmap<N>& other) const {
  Bitmap<N> result;
  for (size_t i = 0; i < kWords; ++i) {
    result.bits_[i] = bits_[i] | other.bits_[i];
  }
  return result;
}

template <size_t N>
inline Bitmap<N> Bitmap<N>::operator^(const Bitmap<N>& other) const {
  Bitmap<N> result;
  for (size_t i = 0; i < kWords; ++i) {
    result.bits_[i] = bits_[i] ^ other.bits_[i];
  }
  return result;
}

template <size_t N>
template <bool Goal>
inline size_t Bitmap<N>::FindValue(size_t index) const {
  TC_ASSERT_LT(index, N);
  size_t offset = index % kWordSize;
  size_t word = index / kWordSize;
  ASSUME(word < kWords);
  size_t here = bits_[word];
  if (!Goal) here = ~here;
  size_t mask = ~static_cast<size_t>(0) << offset;
  here &= mask;
  while (here == 0) {
    ++word;
    if (word >= kWords) {
      return N;
    }
    here = bits_[word];
    if (!Goal) here = ~here;
  }

  word *= kWordSize;
  ASSUME(here != 0);
  size_t ret = absl::countr_zero(here) + word;
  if constexpr (kDeadBits > 0) {
    if constexpr (Goal
#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING
                  && false
#endif  // TCMALLOC_INTERNAL_LEGACY_LOCKING
    ) {
      // We did not early return above, so a set bit was found.  Dead bits are
      // never set.
      ASSUME(ret < N);
    } else {
      // FindClear inverts the word, so we need to clamp.
      if (ret > N) ret = N;
    }
  }
  return ret;
}

template <size_t N>
template <bool Goal>
inline ssize_t Bitmap<N>::FindValueBackwards(size_t index) const {
  TC_ASSERT_LT(index, N);
  size_t offset = index % kWordSize;
  ssize_t word = index / kWordSize;
  ASSUME(word < kWords);
  size_t here = bits_[word];
  if (!Goal) here = ~here;
  size_t mask = (static_cast<size_t>(2) << offset) - 1;
  here &= mask;
  while (here == 0) {
    --word;
    if (word < 0) {
      return -1;
    }
    here = bits_[word];
    if (!Goal) here = ~here;
  }

  word *= kWordSize;
  ASSUME(here != 0);
  size_t ret = absl::bit_width(here) - 1 + word;
  return ret;
}

// Reduction operations for Scale(). In the case of contracting bitmaps,
// we need to reduce multiple (n) bits to a single bit.
//
// - kAny: If any of the n bits are set, the resulting bit is set.
// - kAll: If all of the n bits are set, the resulting bit is set.
enum class ReductionOp {
  kAny,
  kAll,
};

// Scales `src` bitmap from `src_len` bits to `M` bits. If it's a contraction,
// applies `op` to reduce multiple bits to one. If it's an expansion (or
// identity), `op` is a no-op.
template <size_t M, size_t N>
Bitmap<M> Scale(const Bitmap<N>& src, size_t src_len, ReductionOp op) {
  TC_ASSERT_LE(src_len, N);
  TC_ASSERT(src_len % M == 0 || M % src_len == 0);

  Bitmap<M> res;
  const size_t src_per_dst = std::max<size_t>(1, src_len / M);
  const size_t dst_per_src = std::max<size_t>(1, M / src_len);

  const size_t required_count = (op == ReductionOp::kAll) ? src_per_dst : 1;

  size_t dst_idx = 0;
  for (size_t src_idx = 0; src_idx < src_len; src_idx += src_per_dst) {
    size_t count = src.CountBits(src_idx, src_per_dst);
    if (count >= required_count) {
      res.SetRange(dst_idx, dst_per_src);
    }
    dst_idx += dst_per_src;
  }
  return res;
}

}  // namespace tcmalloc_internal
}  // namespace tcmalloc
GOOGLE_MALLOC_SECTION_END

#endif  // TCMALLOC_INTERNAL_RANGE_TRACKER_H_
