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

#include "tcmalloc/internal/range_tracker.h"

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

#include <cstdint>
#include <utility>
#include <vector>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/container/fixed_array.h"
#include "absl/random/random.h"

namespace tcmalloc {
namespace tcmalloc_internal {
namespace {

using testing::ElementsAre;
using testing::Pair;

class BitmapTest : public testing::Test {
 protected:
  template <size_t N>
  std::vector<size_t> FindSetResults(const Bitmap<N>& map) {
    return FindResults<N, true>(map);
  }

  template <size_t N>
  std::vector<size_t> FindClearResults(const Bitmap<N>& map) {
    return FindResults<N, false>(map);
  }

  template <size_t N, bool Value>
  std::vector<size_t> FindResults(const Bitmap<N>& map) {
    std::vector<size_t> results;
    ssize_t last = -1;
    for (size_t i = 0; i < N; ++i) {
      ssize_t j = Value ? map.FindSet(i) : map.FindClear(i);
      EXPECT_LE(last, j) << i;
      EXPECT_LE(i, j) << i;
      EXPECT_GE(N, j) << i;
      if (last != j) {
        results.push_back(j);
        last = j;
      }
    }

    return results;
  }

  template <size_t N>
  std::vector<size_t> FindSetResultsBackwards(const Bitmap<N>& map) {
    return FindResultsBackwards<N, true>(map);
  }

  template <size_t N>
  std::vector<size_t> FindClearResultsBackwards(const Bitmap<N>& map) {
    return FindResultsBackwards<N, false>(map);
  }

  template <size_t N, bool Value>
  std::vector<size_t> FindResultsBackwards(const Bitmap<N>& map) {
    std::vector<size_t> results;
    ssize_t last = N;
    for (ssize_t i = N - 1; i >= 0; --i) {
      ssize_t j = Value ? map.FindSetBackwards(i) : map.FindClearBackwards(i);
      EXPECT_GE(last, j) << i;
      EXPECT_GE(i, j) << i;
      EXPECT_LE(-1, j) << i;
      if (last != j) {
        results.push_back(j);
        last = j;
      }
    }

    return results;
  }
};

TEST_F(BitmapTest, GetBitEmpty) {
  Bitmap<253> map;
  for (size_t i = 0; i < map.size(); ++i) {
    EXPECT_EQ(map.GetBit(i), 0);
  }
}

TEST_F(BitmapTest, CheckIsZero) {
  Bitmap<253> map;
  EXPECT_EQ(map.IsZero(), true);
  for (size_t i = 0; i < map.size(); ++i) {
    map.Clear();
    EXPECT_EQ(map.IsZero(), true);
    EXPECT_EQ(map.SetBit(i), false);
    EXPECT_EQ(map.IsZero(), false);
    EXPECT_EQ(map.SetBit(i), true);
    EXPECT_TRUE(map.ClearBit(i));
    EXPECT_FALSE(map.ClearBit(i));
  }
}

TEST_F(BitmapTest, CheckClearLowestBit) {
  Bitmap<253> map;
  for (size_t i = 0; i < map.size(); ++i) {
    EXPECT_EQ(map.SetBit(i), false);
  }
  for (size_t i = 0; i < map.size(); ++i) {
    size_t index = map.FindSet(0);
    EXPECT_EQ(index, i);
    map.ClearLowestBit();
  }
}

TEST_F(BitmapTest, GetBitOneSet) {
  const size_t N = 251;
  for (size_t s = 0; s < N; s++) {
    Bitmap<N> map;
    EXPECT_EQ(map.SetBit(s), false);
    for (size_t i = 0; i < map.size(); ++i) {
      EXPECT_EQ(map.GetBit(i), i == s ? 1 : 0);
    }
  }
}

TEST_F(BitmapTest, FindSet) {
  Bitmap<253> map;
  EXPECT_THAT(FindSetResults(map), ElementsAre(253));
  EXPECT_THAT(FindSetResultsBackwards(map), ElementsAre(-1));
  map.SetBit(7);
  map.SetBit(14);
  map.SetBit(15);
  map.SetBit(63);
  map.SetBit(128);
  EXPECT_THAT(FindSetResults(map), ElementsAre(7, 14, 15, 63, 128, 253));
  EXPECT_THAT(FindSetResultsBackwards(map),
              ElementsAre(128, 63, 15, 14, 7, -1));
  map.SetBit(195);
  map.SetBit(196);
  map.SetBit(251);
  map.SetBit(252);
  EXPECT_THAT(FindSetResults(map),
              ElementsAre(7, 14, 15, 63, 128, 195, 196, 251, 252));
  EXPECT_THAT(FindSetResultsBackwards(map),
              ElementsAre(252, 251, 196, 195, 128, 63, 15, 14, 7, -1));
  map.SetBit(0);
  EXPECT_THAT(FindSetResultsBackwards(map),
              ElementsAre(252, 251, 196, 195, 128, 63, 15, 14, 7, 0));
}

TEST_F(BitmapTest, FindClear) {
  Bitmap<253> map;
  map.SetRange(0, 253);
  EXPECT_THAT(FindClearResults(map), ElementsAre(253));
  EXPECT_THAT(FindClearResultsBackwards(map), ElementsAre(-1));

  EXPECT_TRUE(map.ClearBit(7));
  EXPECT_TRUE(map.ClearBit(14));
  EXPECT_TRUE(map.ClearBit(15));
  EXPECT_TRUE(map.ClearBit(63));
  EXPECT_TRUE(map.ClearBit(128));
  EXPECT_THAT(FindClearResults(map), ElementsAre(7, 14, 15, 63, 128, 253));
  EXPECT_THAT(FindClearResultsBackwards(map),
              ElementsAre(128, 63, 15, 14, 7, -1));
  EXPECT_TRUE(map.ClearBit(195));
  EXPECT_TRUE(map.ClearBit(196));
  EXPECT_TRUE(map.ClearBit(251));
  EXPECT_TRUE(map.ClearBit(252));
  EXPECT_THAT(FindClearResults(map),
              ElementsAre(7, 14, 15, 63, 128, 195, 196, 251, 252));
  EXPECT_THAT(FindClearResultsBackwards(map),
              ElementsAre(252, 251, 196, 195, 128, 63, 15, 14, 7, -1));
  map.ClearBit(0);
  EXPECT_THAT(FindClearResultsBackwards(map),
              ElementsAre(252, 251, 196, 195, 128, 63, 15, 14, 7, 0));
}

TEST_F(BitmapTest, CountBits) {
  Bitmap<253> map;
  map.SetRange(0, 253);
  EXPECT_EQ(map.CountBits(), 253);
  EXPECT_EQ(map.CountBits(0, 253), 253);
  EXPECT_EQ(map.CountBits(8, 245), 245);
  EXPECT_EQ(map.CountBits(0, 250), 250);

  map.ClearBit(7);
  map.ClearBit(14);
  map.ClearBit(15);
  map.ClearBit(63);
  map.ClearBit(128);

  EXPECT_EQ(map.CountBits(), 248);
  EXPECT_EQ(map.CountBits(0, 253), 248);
  EXPECT_EQ(map.CountBits(8, 245), 241);
  EXPECT_EQ(map.CountBits(0, 250), 245);

  map.ClearBit(195);
  map.ClearBit(196);
  map.ClearBit(251);
  map.ClearBit(252);

  EXPECT_EQ(map.CountBits(), 244);
  EXPECT_EQ(map.CountBits(0, 253), 244);
  EXPECT_EQ(map.CountBits(8, 245), 237);
  EXPECT_EQ(map.CountBits(0, 250), 243);

  map.ClearBit(0);

  EXPECT_EQ(map.CountBits(), 243);
  EXPECT_EQ(map.CountBits(0, 253), 243);
  EXPECT_EQ(map.CountBits(8, 245), 237);
  EXPECT_EQ(map.CountBits(0, 250), 242);
}

TEST_F(BitmapTest, BitwiseOperators) {
  // Test with dead bits (N=253 has 3 dead bits in last word)
  {
    Bitmap<253> map1;
    map1.SetBit(1);
    map1.SetBit(63);
    map1.SetBit(64);
    map1.SetBit(252);

    Bitmap<253> map2;
    map2.SetBit(1);
    map2.SetBit(62);
    map2.SetBit(64);
    map2.SetBit(251);

    // Test AND
    Bitmap<253> and_map = map1 & map2;
    EXPECT_THAT(FindSetResults(and_map), ElementsAre(1, 64, 253));

    // Test OR
    Bitmap<253> or_map = map1 | map2;
    EXPECT_THAT(FindSetResults(or_map), ElementsAre(1, 62, 63, 64, 251, 252));

    // Test XOR
    Bitmap<253> xor_map = map1 ^ map2;
    EXPECT_THAT(FindSetResults(xor_map), ElementsAre(62, 63, 251, 252));

    // Test NOT
    Bitmap<253> not_map = ~map1;
    EXPECT_THAT(FindClearResults(not_map), ElementsAre(1, 63, 64, 252));
    EXPECT_EQ(not_map.CountBits(), 253 - 4);

    // Test that dead bits are not mapped.
    Bitmap<253> all_ones;
    all_ones.SetRange(0, 253);
    Bitmap<253> not_map_dead_bits = ~all_ones;
    EXPECT_TRUE(not_map_dead_bits.IsZero());
  }

  // Test without dead bits (N=64)
  {
    Bitmap<64> map1;
    map1.SetBit(0);
    map1.SetBit(63);

    Bitmap<64> map2;
    map2.SetBit(0);
    map2.SetBit(1);

    Bitmap<64> and_map = map1 & map2;
    EXPECT_THAT(FindSetResults(and_map), ElementsAre(0, 64));

    Bitmap<64> or_map = map1 | map2;
    EXPECT_THAT(FindSetResults(or_map), ElementsAre(0, 1, 63));

    Bitmap<64> xor_map = map1 ^ map2;
    EXPECT_THAT(FindSetResults(xor_map), ElementsAre(1, 63));

    Bitmap<64> not_map = ~map1;
    EXPECT_THAT(FindClearResults(not_map), ElementsAre(0, 63));
    EXPECT_EQ(not_map.CountBits(), 62);
  }
}

TEST_F(BitmapTest, PopBatch) {
  Bitmap<253> map;
  map.SetBit(7);
  map.SetBit(14);
  map.SetBit(15);
  map.SetBit(63);
  map.SetBit(128);

  std::vector<size_t> dest;
  size_t popped = map.PopBatch([&](size_t v) { dest.push_back(v); }, 3);
  EXPECT_EQ(dest.size(), popped);
  EXPECT_THAT(dest, ElementsAre(7, 14, 15));
  EXPECT_FALSE(map.GetBit(7));
  EXPECT_FALSE(map.GetBit(14));
  EXPECT_FALSE(map.GetBit(15));
  EXPECT_TRUE(map.GetBit(63));
  EXPECT_TRUE(map.GetBit(128));

  dest.clear();
  popped = map.PopBatch([&](size_t v) { dest.push_back(v); }, 10);
  EXPECT_EQ(popped, dest.size());
  EXPECT_THAT(dest, ElementsAre(63, 128));
  EXPECT_TRUE(map.IsZero());
}

class RangeTrackerTest : public ::testing::Test {
 protected:
  std::vector<std::pair<size_t, size_t>> FreeRanges() {
    std::vector<std::pair<size_t, size_t>> ret;
    size_t index = 0, len;
    while (range_.NextFreeRange(index, &index, &len)) {
      ret.push_back({index, len});
      index += len;
    }
    return ret;
  }
  static constexpr size_t kBits = 1017;
  RangeTracker<kBits> range_;
};

TEST_F(RangeTrackerTest, Trivial) {
  EXPECT_EQ(kBits, range_.size());
  EXPECT_EQ(0, range_.used());
  EXPECT_EQ(kBits, range_.longest_free());
  EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, kBits)));
  ASSERT_EQ(0, range_.FindAndMark(kBits));
  EXPECT_EQ(0, range_.longest_free());
  EXPECT_EQ(kBits, range_.used());
  EXPECT_THAT(FreeRanges(), ElementsAre());
  range_.Unmark(0, 100);
  EXPECT_EQ(100, range_.longest_free());
  EXPECT_EQ(kBits - 100, range_.used());
  EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100)));
  // non-contiguous - shouldn't increase longest
  range_.Unmark(200, 100);
  EXPECT_EQ(100, range_.longest_free());
  EXPECT_EQ(kBits - 200, range_.used());
  EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100), Pair(200, 100)));
  range_.Unmark(100, 100);
  EXPECT_EQ(300, range_.longest_free());
  EXPECT_EQ(kBits - 300, range_.used());
  EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 300)));
}

TEST_F(RangeTrackerTest, Mark) {
  EXPECT_EQ(range_.longest_free(), kBits);
  range_.Mark(100, 100);
  EXPECT_EQ(range_.used(), 100);
  EXPECT_EQ(range_.longest_free(), kBits - 200);
  EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100), Pair(200, kBits - 200)));
  range_.Unmark(100, 100);
  EXPECT_EQ(range_.used(), 0);
  EXPECT_EQ(range_.longest_free(), kBits);
}

TEST(BitmapScaleTest, ScaleAssertionFailures) {
#ifdef NDEBUG
  GTEST_SKIP() << "Requires debug mode";
#endif
  Bitmap<64> map1;
  // 1. src_len (64) doesn't divide M (10) and vice versa
  EXPECT_DEATH(Scale<10>(map1, 64, ReductionOp::kAll), "");
  // 2. src_len > N (64)
  EXPECT_DEATH(Scale<256>(map1, 128, ReductionOp::kAll), "");
}

TEST(BitmapScaleTest, Scale) {
  // Test case 1: N == M (identity)
  {
    Bitmap<64> map1;
    map1.SetBit(3);
    map1.SetBit(17);
    map1.SetBit(63);

    auto map_all = Scale<64>(map1, 64, ReductionOp::kAll);
    auto map_any = Scale<64>(map1, 64, ReductionOp::kAny);
    for (size_t i = 0; i < 64; ++i) {
      EXPECT_EQ(map1.GetBit(i), map_all.GetBit(i));
      EXPECT_EQ(map1.GetBit(i), map_any.GetBit(i));
    }
  }

  // Test case 2: N > M (contracting)
  {
    Bitmap<64> map1;
    // We group by 2 bits.
    // Group 0: bits 0, 1 -> both must be set
    map1.SetBit(0);
    map1.SetBit(1);

    // Group 1: bits 2, 3 -> only one set
    // Should not be set for kAll, should be set for kAny
    map1.SetBit(2);

    // Group 7: bits 14, 15 -> both set
    map1.SetBit(14);
    map1.SetBit(15);

    auto map_all = Scale<8>(map1, 16, ReductionOp::kAll);
    auto map_any = Scale<8>(map1, 16, ReductionOp::kAny);

    EXPECT_TRUE(map_all.GetBit(0));
    EXPECT_TRUE(map_any.GetBit(0));

    EXPECT_FALSE(map_all.GetBit(1));
    EXPECT_TRUE(map_any.GetBit(1));

    for (size_t i = 2; i <= 6; ++i) {
      EXPECT_FALSE(map_all.GetBit(i));
      EXPECT_FALSE(map_any.GetBit(i));
    }

    EXPECT_TRUE(map_all.GetBit(7));
    EXPECT_TRUE(map_any.GetBit(7));
  }

  // Test case 3: N < M (expanding)
  {
    Bitmap<64> map1;
    // We expand by 2 bits.
    map1.SetBit(0);
    map1.SetBit(7);

    auto map_all = Scale<16>(map1, 8, ReductionOp::kAll);
    auto map_any = Scale<16>(map1, 8, ReductionOp::kAny);

    // Both operations behave identically for expanding.
    EXPECT_TRUE(map_all.GetBit(0));
    EXPECT_TRUE(map_all.GetBit(1));
    EXPECT_TRUE(map_any.GetBit(0));
    EXPECT_TRUE(map_any.GetBit(1));

    EXPECT_FALSE(map_all.GetBit(2));
    EXPECT_FALSE(map_all.GetBit(3));
    EXPECT_FALSE(map_any.GetBit(2));
    EXPECT_FALSE(map_any.GetBit(3));

    EXPECT_TRUE(map_all.GetBit(14));
    EXPECT_TRUE(map_all.GetBit(15));
    EXPECT_TRUE(map_any.GetBit(14));
    EXPECT_TRUE(map_any.GetBit(15));
  }
}

}  // namespace
}  // namespace tcmalloc_internal
}  // namespace tcmalloc
