// Copyright 2022 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.

#include "tcmalloc/sizemap.h"

#include <algorithm>
#include <cstddef>
#include <cstdio>
#include <cstring>
#include <new>
#include <vector>

#include "absl/base/macros.h"
#include "absl/base/nullability.h"
#include "absl/types/span.h"
#include "tcmalloc/common.h"
#include "tcmalloc/huge_page_aware_allocator.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/parameter_accessors.h"
#include "tcmalloc/internal/sampled_allocation.h"
#include "tcmalloc/internal/size_class_info.h"
#include "tcmalloc/pages.h"
#include "tcmalloc/parameters.h"
#include "tcmalloc/span.h"
#include "tcmalloc/static_vars.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace tcmalloc_internal {

const SizeClasses& SizeMap::CurrentClasses() {
  switch (Static::size_class_configuration()) {
    case SizeClassConfiguration::kPow2Only:
      return kExperimentalPow2SizeClasses;
    case SizeClassConfiguration::kReuseRelaxedBelow64:
      return kReuseRelaxedBelow64SizeClasses;
    case SizeClassConfiguration::kReuse:
      // TODO(b/512895228): remove this opt out once we are done experimenting.
      return kSizeClasses;
    case SizeClassConfiguration::kLegacy:
      // TODO(b/242710633): remove this opt out.
      return kLegacySizeClasses;
  }
  TC_BUG("unreachable");
}

bool SizeMap::CheckAssumptions() {
  bool failed = false;
  auto a = CurrentClasses().assumptions;
  if (a.has_expanded_classes != kHasExpandedClasses) {
    fprintf(stderr, "kHasExpandedClasses: assumed %d, actual %d\n",
            a.has_expanded_classes, kHasExpandedClasses);
    failed |= true;
  }
#ifdef NDEBUG
  if (a.span_size != sizeof(Span)) {
    fprintf(stderr, "sizeof(Span): assumed %zu, actual %zu\n", a.span_size,
            sizeof(Span));
    failed |= true;
  }
#endif  // NDEBUG
  if (a.sampling_interval != kDefaultProfileSamplingInterval) {
    fprintf(stderr,
            "kDefaultProfileSamplingInterval: assumed %zu, actual %zu\n",
            a.sampling_interval, kDefaultProfileSamplingInterval);
    failed |= true;
  }
  if (a.large_size != SizeMap::kLargeSize) {
    fprintf(stderr, "SizeMap::kLargeSize: assumed %zu, actual %u\n",
            a.large_size, SizeMap::kLargeSize);
    failed |= true;
  }
  if (a.large_size_alignment != SizeMap::kLargeSizeAlignment) {
    fprintf(stderr, "SizeMap::kLargeSizeAlignment: assumed %zu, actual %u\n",
            a.large_size_alignment, SizeMap::kLargeSizeAlignment);
    failed |= true;
  }
  if (failed) {
    fprintf(stderr, "*************************************\n");
    fprintf(stderr, "* MISMATCHED SIZE CLASS ASSUMPTIONS *\n");
    fprintf(stderr, "*************************************\n");
  }
  return !failed;
}

bool SizeMap::IsValidSizeClass(size_t size, Length pages,
                               size_t num_objects_to_move) {
  if (size == 0) {
    TC_LOG("size class is 0");
    return false;
  }
  if (size > kMaxSize) {
    TC_LOG("size %v class too big %v", size, kMaxSize);
    return false;
  }
  // Check required alignment
  const size_t alignment = size > SizeMap::kLargeSize
                               ? kLargeSizeAlignment
                               : static_cast<size_t>(kAlignment);
  if ((size & (alignment - 1)) != 0) {
    TC_LOG("%v not aligned properly %v", size, alignment);
    return false;
  }
  if (pages == Length(0)) {
    TC_LOG("pages should not be 0");
    return false;
  }
  if (pages > Length(32)) {
    TC_LOG("pages %v limited to 32", pages);
    return false;
  }
  const size_t objects_per_span = pages.in_bytes() / size;
  if (objects_per_span < 1) {
    TC_LOG("each span must have at least one object");
    return false;
  }
  if (!Span::IsValidSizeClass(size, pages)) {
    TC_LOG("%v span size class assumptions are broken: pages=%v objs=%v", size,
           pages, objects_per_span);
    return false;
  }
  if (!HugePageAwareAllocator::IsValidSizeClass(size, pages)) {
    TC_LOG("%v hpaa size class assumptions are broken: pages=%v objs=%v", size,
           pages, objects_per_span);
    return false;
  }
  if (num_objects_to_move < 2) {
    TC_LOG("num objects to move %v too small (<2)", num_objects_to_move);
    return false;
  }
  if (num_objects_to_move > kMaxObjectsToMove) {
    TC_LOG("num objects to move %v too large %v", num_objects_to_move,
           kMaxObjectsToMove);
    return false;
  }

  return true;
}

bool SizeMap::SetSizeClasses(absl::Span<const SizeClassInfo> size_classes) {
  const int num_classes = size_classes.size();
  if (!ValidSizeClasses(size_classes)) {
    return false;
  }

  class_to_size_[0] = 0;
  class_to_pages_[0] = 0;
  num_objects_to_move_[0] = 0;

  int curr = 1;
  for (int c = 1; c < num_classes; c++) {
    class_to_size_[curr] = size_classes[c].size;
    class_to_pages_[curr] = BytesToLengthFloor(size_classes[c].bytes).raw_num();
    num_objects_to_move_[curr] = size_classes[c].num_to_move;
    ++curr;
  }

  // Fill any unspecified size classes with 0.
  for (int x = curr; x < kNumBaseClasses; x++) {
    class_to_size_[x] = 0;
    class_to_pages_[x] = 0;
    num_objects_to_move_[x] = 0;
  }

  // Copy selected size classes into the upper registers.
  for (int i = 1; i < (kNumClasses / kNumBaseClasses); i++) {
    std::copy(&class_to_size_[0], &class_to_size_[kNumBaseClasses],
              &class_to_size_[kNumBaseClasses * i]);
    std::copy(&class_to_pages_[0], &class_to_pages_[kNumBaseClasses],
              &class_to_pages_[kNumBaseClasses * i]);
    std::copy(&num_objects_to_move_[0], &num_objects_to_move_[kNumBaseClasses],
              &num_objects_to_move_[kNumBaseClasses * i]);
  }

  return true;
}

// Return true if all size classes meet the requirements for alignment
// ordering and min and max values.
bool SizeMap::ValidSizeClasses(absl::Span<const SizeClassInfo> size_classes) {
  if (size_classes.empty()) {
    return false;
  }
  int num_classes = size_classes.size();
  if (kHasExpandedClasses && num_classes > kNumBaseClasses) {
    num_classes = kNumBaseClasses;
  }

  if (size_classes[0].size != 0 || size_classes[0].bytes != Bytes(0) ||
      size_classes[0].num_to_move != 0) {
    return false;
  }

  for (int c = 1; c < num_classes; c++) {
    size_t class_size = size_classes[c].size;
    Length pages = BytesToLengthFloor(size_classes[c].bytes);
    if (Bytes(pages.in_bytes()) != size_classes[c].bytes) {
      TC_LOG("Non-page multiple span size %v: %v", c, size_classes[c].bytes);
      return false;
    }
    size_t num_objects_to_move = size_classes[c].num_to_move;
    // Each size class must be larger than the previous size class.
    if (class_size <= size_classes[c - 1].size) {
      TC_LOG("Non-increasing size class %v: prev=%v next=%v", c,
             size_classes[c - 1].size, class_size);
      return false;
    }
    if (!IsValidSizeClass(class_size, pages, num_objects_to_move)) {
      return false;
    }
  }
  // Last size class must be kMaxSize.  This is not strictly
  // class_to_size_[kNumBaseClasses - 1] because several size class
  // configurations populate fewer distinct size classes and fill the tail of
  // the array with zeroes.
  if (size_classes[num_classes - 1].size != kMaxSize) {
    TC_LOG("last class %v size %v doesn't cover kMaxSize %v", num_classes - 1,
           size_classes[num_classes - 1].size, kMaxSize);
    return false;
  }
  return true;
}

// Initialize the mapping arrays
bool SizeMap::Init(absl::Span<const SizeClassInfo> size_classes) {
  // Do some sanity checking on add_amount[]/shift_amount[]/class_array[]
  TC_CHECK_EQ(ClassIndex(0), 0);
  TC_CHECK_LT(ClassIndex(kMaxSize), sizeof(class_array_));
  static_assert(kAlignment <= std::align_val_t{16}, "kAlignment is too large");

  if (!SetSizeClasses(size_classes)) {
    return false;
  }

  int next_size = 0;
  for (int c = 1; c < kNumClasses; c++) {
    const int max_size_in_class = class_to_size_[c];

    for (int s = next_size; s <= max_size_in_class;
         s += static_cast<size_t>(kAlignment)) {
      class_array_[ClassIndex(s)] = c;
    }
    next_size = max_size_in_class + static_cast<size_t>(kAlignment);
    if (next_size > kMaxSize) {
      break;
    }
  }

  // Point all lookups in hot registers (Malloc P0, New P1, Malloc P1)
  // directly to New's P0. We only overwrite the lookups if heap
  // partitioning is active with the dedicated size classes.
  for (size_t i = 1; i < kHotRegisters; ++i) {
    std::copy(&class_array_[0], &class_array_[kClassArraySize],
              &class_array_[kClassArraySize * i]);
  }

  bool heap_partitioning_active = tc_globals.multiple_non_numa_partitions();
  if (kSecurityPartitions > 1 && heap_partitioning_active) {
    bool heap_partitioning_full =
        Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull;
    next_size = 0;
    for (int c = kNumBaseClasses + 1; c < kExpandedClassesStart; ++c) {
      const int max_size_in_class = class_to_size_[c];

      for (int s = next_size; s <= max_size_in_class;
           s += static_cast<size_t>(kAlignment)) {
        // Route Hot Malloc P1 to security partition P1.
        class_array_[ClassIndex(s) +
                     kClassArraySize * (kSecurityPartitions + 1)] = c;
        // Route Hot New P1 to security partition P1.
        class_array_[ClassIndex(s) + kClassArraySize] = c;
        if (heap_partitioning_full) continue;
        // In kLight mode, route Hot New P0 to P1.
        class_array_[ClassIndex(s)] = c;
      }
      next_size = max_size_in_class + static_cast<size_t>(kAlignment);
      if (next_size > kMaxSize) {
        break;
      }
    }
  }

  if (ColdFeatureActive()) {
    memset(cold_sizes_, 0, sizeof(cold_sizes_));
    cold_sizes_count_ = 0;

    // Point all lookups in the first or second upper register of class_array_
    // (allocations seeking cold memory, with hints for partition 0, i.e.,
    // pointerless allocations) to the lower size classes.  This gives us an
    // easy fallback for sizes that are too small for moving to cold memory (due
    // to intrusive span metadata).
    std::copy(&class_array_[0], &class_array_[kClassArraySize],
              &class_array_[kClassArraySize * kColdRegisterStride]);

    for (int c = kExpandedClassesStart; c < kNumClasses; c++) {
      size_t max_size_in_class = class_to_size_[c];
      if (max_size_in_class == 0) {
        next_size = max_size_in_class + static_cast<size_t>(kAlignment);
        continue;
      }

      cold_sizes_[cold_sizes_count_] = c;
      ++cold_sizes_count_;

      for (int s = next_size; s <= max_size_in_class;
           s += static_cast<size_t>(kAlignment)) {
        class_array_[ClassIndex(s) + kClassArraySize * kColdRegisterStride] = c;
      }
      next_size = max_size_in_class + static_cast<size_t>(kAlignment);
      if (next_size > kMaxSize) {
        break;
      }
    }
    if (kSecurityPartitions > 1) {
      if (Parameters::heap_partitioning_mode() == HeapPartitioningMode::kFull) {
        // Point all lookups in Cold New's P1 register to Hot New's P1.
        std::copy(
            &class_array_[kClassArraySize], &class_array_[kClassArraySize * 2],
            &class_array_[kClassArraySize * (2 * kSecurityPartitions + 1)]);
      } else {
        // Point all lookups in Cold New's P1 register to Cold New's P0.
        std::copy(
            &class_array_[kClassArraySize * (2 * kSecurityPartitions)],
            &class_array_[kClassArraySize * (2 * kSecurityPartitions + 1)],
            &class_array_[kClassArraySize * (2 * kSecurityPartitions + 1)]);
      }
    }
  }

  return true;
}

}  // namespace tcmalloc_internal
}  // namespace tcmalloc

extern "C" {

void TCMalloc_Internal_GetSizeClasses(
    std::vector<tcmalloc::tcmalloc_internal::TracerSizeClassInfo>* absl_nonnull
        size_classes) {
  tcmalloc::tcmalloc_internal::tc_globals.InitIfNecessary();
  size_t num_classes = tcmalloc::tcmalloc_internal::kNumClasses;
  size_classes->resize(num_classes);
  for (size_t i = 0; i < num_classes; ++i) {
    (*size_classes)[i].size =
        tcmalloc::tcmalloc_internal::tc_globals.sizemap().class_to_size(i);
    (*size_classes)[i].span_size_in_bytes =
        tcmalloc::tcmalloc_internal::tc_globals.sizemap()
            .class_to_pages(i)
            .in_bytes();
    (*size_classes)[i].num_objects_to_move =
        tcmalloc::tcmalloc_internal::tc_globals.sizemap().num_objects_to_move(
            i);
  }
}

}  // extern "C"

GOOGLE_MALLOC_SECTION_END
