/*
 * Copyright (C) 2025 The Android Open Source Project
 *
 * 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
 *
 *      http://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 "src/trace_processor/core/util/bit_vector.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>

#include "test/gtest_and_gmock.h"

namespace perfetto::trace_processor::core {
namespace {

// Test that the default constructor creates an empty bit vector
TEST(BitVectorTest, DefaultConstructor) {
  BitVector bits;
  EXPECT_EQ(bits.size(), 0u);
}

TEST(BitVectorTest, CreateWithSize) {
  {
    auto bits = BitVector::CreateWithSize(31);
    EXPECT_EQ(bits.size(), 31u);
    for (size_t i = 0; i < 31; ++i) {
      EXPECT_FALSE(bits.is_set(i)) << "Bit " << i << " should be unset";
    }
  }
  {
    auto bits = BitVector::CreateWithSize(65, true);
    EXPECT_EQ(bits.size(), 65u);
    for (size_t i = 0; i < 65; ++i) {
      EXPECT_TRUE(bits.is_set(i)) << "Bit " << i << " should be set";
    }
  }
}

// Test push_back functionality
TEST(BitVectorTest, PushBack) {
  BitVector bits;

  // Size should start at 0
  EXPECT_EQ(bits.size(), 0u);

  // Add some bits
  bits.push_back(true);   // bit 0
  bits.push_back(false);  // bit 1
  bits.push_back(true);   // bit 2

  EXPECT_EQ(bits.size(), 3u);

  EXPECT_TRUE(bits.is_set(0));
  EXPECT_FALSE(bits.is_set(1));
  EXPECT_TRUE(bits.is_set(2));

  // Add more bits to cross word boundaries (assuming 64-bit words)
  for (size_t i = 3; i < 64; ++i) {
    bits.push_back(false);
  }

  // Add bits that go into a second word
  bits.push_back(true);   // bit 64
  bits.push_back(false);  // bit 65
  bits.push_back(true);   // bit 66

  EXPECT_EQ(bits.size(), 67u);

  EXPECT_TRUE(bits.is_set(64));
  EXPECT_FALSE(bits.is_set(65));
  EXPECT_TRUE(bits.is_set(66));
}

// Test set, clear, and is_set methods
TEST(BitVectorTest, SetClearAndIsSet) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 200; ++i) {
    bits.push_back(false);
  }

  // Initially all bits should be unset (we pushed all false values)
  for (size_t i = 0; i < 200; ++i) {
    EXPECT_FALSE(bits.is_set(i));
  }

  // Set specific bits including across word boundaries
  bits.set(5);
  bits.set(63);   // Last bit of first word
  bits.set(64);   // First bit of second word
  bits.set(127);  // Last bit of second word
  bits.set(128);  // First bit of third word
  bits.set(198);

  // Check set bits are set
  EXPECT_TRUE(bits.is_set(5));
  EXPECT_TRUE(bits.is_set(63));
  EXPECT_TRUE(bits.is_set(64));
  EXPECT_TRUE(bits.is_set(127));
  EXPECT_TRUE(bits.is_set(128));
  EXPECT_TRUE(bits.is_set(198));

  // Check some unset bits remain unset
  EXPECT_FALSE(bits.is_set(0));
  EXPECT_FALSE(bits.is_set(6));
  EXPECT_FALSE(bits.is_set(65));
  EXPECT_FALSE(bits.is_set(199));

  // Clear some of the set bits
  bits.clear(5);
  bits.clear(64);
  bits.clear(198);

  // Verify cleared bits are now unset
  EXPECT_FALSE(bits.is_set(5));
  EXPECT_TRUE(bits.is_set(63));
  EXPECT_FALSE(bits.is_set(64));
  EXPECT_TRUE(bits.is_set(127));
  EXPECT_TRUE(bits.is_set(128));
  EXPECT_FALSE(bits.is_set(198));
}

// Test the change method
TEST(BitVectorTest, ChangeMethod) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 100; ++i) {
    bits.push_back(false);
  }

  // Set some bits first
  bits.set(10);
  bits.set(20);
  bits.set(30);

  // Change them with different transitions
  bits.change(10, false);  // true -> false
  bits.change(20, true);   // true -> true (no change)
  bits.change(30, false);  // true -> false
  bits.change(40, true);   // false -> true

  // Verify changes took effect
  EXPECT_FALSE(bits.is_set(10));
  EXPECT_TRUE(bits.is_set(20));
  EXPECT_FALSE(bits.is_set(30));
  EXPECT_TRUE(bits.is_set(40));
}

// Test the change_assume_unset method
TEST(BitVectorTest, ChangeAssumeUnsetMethod) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 100; ++i) {
    bits.push_back(false);
  }

  // Use change_assume_unset on unset bits
  bits.change_assume_unset(15, true);
  bits.change_assume_unset(25, false);  // No change since value is false
  bits.change_assume_unset(35, true);

  // Verify changes
  EXPECT_TRUE(bits.is_set(15));
  EXPECT_FALSE(bits.is_set(25));
  EXPECT_TRUE(bits.is_set(35));
}

// Test the set_bits_until_in_word method
TEST(BitVectorTest, SetBitsUntilInWord) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 130; ++i) {  // Only need 130 for the test
    bits.push_back(false);
  }

  // Set bits in first word
  bits.set(1);
  bits.set(3);
  bits.set(7);
  bits.set(20);

  // Check counts in first word
  EXPECT_EQ(bits.count_set_bits_until_in_word(0),
            0u);  // No bits before position 0
  EXPECT_EQ(bits.count_set_bits_until_in_word(1),
            0u);  // The bit at position 1 itself isn't counted
  EXPECT_EQ(bits.count_set_bits_until_in_word(2), 1u);  // Only bit 1 is counted
  EXPECT_EQ(bits.count_set_bits_until_in_word(4), 2u);  // Bits 1 and 3 counted
  EXPECT_EQ(bits.count_set_bits_until_in_word(10),
            3u);  // Bits 1, 3, and 7 counted
  EXPECT_EQ(bits.count_set_bits_until_in_word(21),
            4u);  // Bits 1, 3, 7, and 20 counted

  // Set bits in second word
  bits.set(64);  // First bit of second word
  bits.set(70);  // Another bit in second word

  // Check counts in second word - this should start from 0 again
  EXPECT_EQ(bits.count_set_bits_until_in_word(64),
            0u);  // No bits before position 64 in this word
  EXPECT_EQ(bits.count_set_bits_until_in_word(65),
            1u);  // Just bit 64 counted in this word
  EXPECT_EQ(bits.count_set_bits_until_in_word(71),
            2u);  // Bits 64 and 70 counted in this word
}

// Test the PackLeft method
TEST(BitVectorTest, PackLeftMethod) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 20; ++i) {
    bits.push_back(false);
  }

  // Set every other bit
  for (size_t i = 0; i < 20; i += 2) {
    bits.set(i);
  }

  // Create source array with indices 0-19
  uint32_t source[20];
  for (uint32_t i = 0; i < 20; ++i) {
    source[i] = i;
  }

  // Create target array
  uint32_t target[20] = {0};

  // Filter with default behavior (keep set bits)
  uint32_t* end = bits.PackLeft(source, source + 20, target);

  // Should keep 10 elements (those with even indices where bits are set)
  EXPECT_EQ(end - target, 10);

  // Check the filtered values
  for (size_t i = 0; i < 10; ++i) {
    EXPECT_EQ(target[i], i * 2);
  }

  // Reset target
  std::fill_n(target, 20, 0);

  // Filter with inverted behavior (keep unset bits)
  end = bits.PackLeft<true>(source, source + 20, target);

  // Should keep 10 elements (those with odd indices where bits are not set)
  EXPECT_EQ(end - target, 10);

  // Check the filtered values
  for (size_t i = 0; i < 10; ++i) {
    EXPECT_EQ(target[i], i * 2 + 1);
  }
}

// Test the PrefixPopcount method
TEST(BitVectorTest, PrefixPopcountMethod) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 200; ++i) {
    bits.push_back(false);
  }

  // First word (0-63): set 3 bits
  bits.set(5);
  bits.set(20);
  bits.set(40);

  // Second word (64-127): set 2 bits
  bits.set(70);
  bits.set(100);

  // Third word (128-191): set 4 bits
  bits.set(130);
  bits.set(140);
  bits.set(150);
  bits.set(160);

  // Get the prefix popcounts
  auto prefixes = bits.PrefixPopcount();

  // Check size - should be one entry per word
  EXPECT_EQ(prefixes.size(), (bits.size() + 63) / 64);

  // Check values
  EXPECT_EQ(prefixes[0], 0u);  // No words before the first word
  EXPECT_EQ(prefixes[1], 3u);  // First word had 3 set bits
  EXPECT_EQ(prefixes[2], 5u);  // First+Second word had 3+2=5 set bits
  EXPECT_EQ(prefixes[3], 9u);  // First+Second+Third had 3+2+4=9 set bits
}

// Test across word boundaries
TEST(BitVectorTest, WordBoundaries) {
  BitVector bits;

  // First initialize bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 256; ++i) {
    bits.push_back(false);
  }

  // Set bits at word boundaries
  bits.set(0);    // First bit of first word
  bits.set(63);   // Last bit of first word
  bits.set(64);   // First bit of second word
  bits.set(127);  // Last bit of second word
  bits.set(128);  // First bit of third word
  bits.set(191);  // Last bit of third word
  bits.set(192);  // First bit of fourth word
  bits.set(255);  // Last bit of fourth word

  // Check all bits are correctly set
  EXPECT_TRUE(bits.is_set(0));
  EXPECT_TRUE(bits.is_set(63));
  EXPECT_TRUE(bits.is_set(64));
  EXPECT_TRUE(bits.is_set(127));
  EXPECT_TRUE(bits.is_set(128));
  EXPECT_TRUE(bits.is_set(191));
  EXPECT_TRUE(bits.is_set(192));
  EXPECT_TRUE(bits.is_set(255));

  // Clear some boundary bits
  bits.clear(0);
  bits.clear(64);
  bits.clear(128);
  bits.clear(192);

  // Check cleared bits are unset and others remained set
  EXPECT_FALSE(bits.is_set(0));
  EXPECT_TRUE(bits.is_set(63));
  EXPECT_FALSE(bits.is_set(64));
  EXPECT_TRUE(bits.is_set(127));
  EXPECT_FALSE(bits.is_set(128));
  EXPECT_TRUE(bits.is_set(191));
  EXPECT_FALSE(bits.is_set(192));
  EXPECT_TRUE(bits.is_set(255));
}

// Test with many bits
TEST(BitVectorTest, LargeVector) {
  // Create a bit vector with capacity for 8192 bits (power of two bits)
  BitVector bits;
  EXPECT_EQ(bits.size(), 0u);  // Size starts at 0

  // Initialize bits
  for (size_t i = 0; i < 8192; ++i) {
    bits.push_back(false);
  }

  EXPECT_EQ(bits.size(), 8192u);

  // Set every 1000th bit
  for (size_t i = 0; i < 8192; i += 1000) {
    bits.set(i);
  }

  // Verify the correct bits are set
  for (size_t i = 0; i < 8192; ++i) {
    if (i % 1000 == 0) {
      EXPECT_TRUE(bits.is_set(i)) << "Bit " << i << " should be set";
    } else {
      EXPECT_FALSE(bits.is_set(i)) << "Bit " << i << " should be unset";
    }
  }
}

// Test a mixed sequence of operations
TEST(BitVectorTest, MixedOperations) {
  BitVector bits;

  EXPECT_EQ(bits.size(), 0u);  // Initial size should be 0

  // First initialize 100 bits (since CreateWithCapacity only sets capacity)
  for (size_t i = 0; i < 100; ++i) {
    bits.push_back(false);
  }

  // Add some more bits
  bits.push_back(true);
  bits.push_back(false);
  bits.push_back(true);

  // Current size should be 103
  EXPECT_EQ(bits.size(), 103u);

  // Set, clear, and modify bits
  bits.set(5);
  bits.set(50);
  bits.clear(50);
  bits.change(5, false);
  bits.change(10, true);
  bits.change_assume_unset(20, true);

  // Check final state
  EXPECT_FALSE(bits.is_set(5));
  EXPECT_FALSE(bits.is_set(50));
  EXPECT_TRUE(bits.is_set(10));
  EXPECT_TRUE(bits.is_set(20));
  EXPECT_TRUE(bits.is_set(100));
  EXPECT_FALSE(bits.is_set(101));
  EXPECT_TRUE(bits.is_set(102));
}

// Test resize to shrink
TEST(BitVectorTest, ResizeShrink) {
  auto bits = BitVector::CreateWithSize(100, false);
  bits.set(10);
  bits.set(50);
  bits.set(90);

  bits.resize(60);
  EXPECT_EQ(bits.size(), 60u);

  // Bits within new size should be preserved
  EXPECT_TRUE(bits.is_set(10));
  EXPECT_TRUE(bits.is_set(50));
}

// Test resize to grow with default value (false)
TEST(BitVectorTest, ResizeGrowFalse) {
  auto bits = BitVector::CreateWithSize(50, false);
  bits.set(10);
  bits.set(40);

  bits.resize(100, false);
  EXPECT_EQ(bits.size(), 100u);

  // Original bits preserved
  EXPECT_TRUE(bits.is_set(10));
  EXPECT_TRUE(bits.is_set(40));

  // New bits should be unset
  for (size_t i = 50; i < 100; ++i) {
    EXPECT_FALSE(bits.is_set(i)) << "Bit " << i << " should be unset";
  }
}

// Test resize to grow with value true
TEST(BitVectorTest, ResizeGrowTrue) {
  auto bits = BitVector::CreateWithSize(50, false);
  bits.set(10);

  bits.resize(100, true);
  EXPECT_EQ(bits.size(), 100u);

  // Original bits preserved
  EXPECT_TRUE(bits.is_set(10));
  EXPECT_FALSE(bits.is_set(20));

  // New bits should be set
  for (size_t i = 50; i < 100; ++i) {
    EXPECT_TRUE(bits.is_set(i)) << "Bit " << i << " should be set";
  }
}

// Test resize across word boundaries
TEST(BitVectorTest, ResizeAcrossWordBoundary) {
  // Start with size that doesn't align to 64
  auto bits = BitVector::CreateWithSize(70, false);
  bits.set(60);
  bits.set(65);

  // Shrink to within first word
  bits.resize(30);
  EXPECT_EQ(bits.size(), 30u);

  // Grow back across boundary with true
  bits.resize(80, true);
  EXPECT_EQ(bits.size(), 80u);

  // Bits 30-79 should be set
  for (size_t i = 30; i < 80; ++i) {
    EXPECT_TRUE(bits.is_set(i)) << "Bit " << i << " should be set";
  }
}

// Helper: build a BitVector from a string of '0' and '1' characters.
BitVector BvFromStr(const char* s) {
  BitVector bv;
  for (; *s; ++s) {
    bv.push_back(*s == '1');
  }
  return bv;
}

// Helper: convert a BitVector to a string of '0' and '1' characters.
std::string BvToStr(const BitVector& bv) {
  std::string s;
  for (uint64_t i = 0; i < bv.size(); ++i) {
    s += bv.is_set(i) ? '1' : '0';
  }
  return s;
}

TEST(BitVectorTest, CompactEmpty) {
  auto bv = BvFromStr("");
  auto keep = BvFromStr("");
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 0u);
}

TEST(BitVectorTest, CompactKeepAll) {
  auto bv = BvFromStr("10110");
  auto keep = BvFromStr("11111");
  EXPECT_EQ(BvToStr(std::move(bv).Compact(keep)), "10110");
}

TEST(BitVectorTest, CompactKeepNone) {
  auto bv = BvFromStr("10110");
  auto keep = BvFromStr("00000");
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 0u);
}

TEST(BitVectorTest, CompactKeepSome) {
  // bv:   1 0 1 1 0
  // keep: 1 0 0 1 1
  // kept positions: 0, 3, 4 → values: 1, 1, 0
  auto bv = BvFromStr("10110");
  auto keep = BvFromStr("10011");
  EXPECT_EQ(BvToStr(std::move(bv).Compact(keep)), "110");
}

TEST(BitVectorTest, CompactKeepFirst) {
  auto bv = BvFromStr("10110");
  auto keep = BvFromStr("10000");
  EXPECT_EQ(BvToStr(std::move(bv).Compact(keep)), "1");
}

TEST(BitVectorTest, CompactKeepLast) {
  auto bv = BvFromStr("10111");
  auto keep = BvFromStr("00001");
  EXPECT_EQ(BvToStr(std::move(bv).Compact(keep)), "1");
}

TEST(BitVectorTest, CompactSingleWord) {
  // 64-bit word, keep every other bit.
  auto bv = BitVector::CreateWithSize(64, false);
  auto keep = BitVector::CreateWithSize(64, false);
  // bv: bits 0,2,4,...,62 are set (even positions)
  // keep: bits 0,1,2,...,31 are set (first half)
  for (uint32_t i = 0; i < 64; i += 2) {
    bv.set(i);
  }
  for (uint32_t i = 0; i < 32; ++i) {
    keep.set(i);
  }
  // Kept positions 0..31. At even positions bv=1, at odd bv=0.
  // Result should be 16 bits: alternating 1,0,1,0,...
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 32u);
  for (uint32_t i = 0; i < 32; ++i) {
    EXPECT_EQ(result.is_set(i), (i % 2 == 0)) << "bit " << i;
  }
}

TEST(BitVectorTest, CompactCrossWordBoundary) {
  // 128 bits: set bits at positions 63 and 64 (word boundary).
  // Keep only those two positions.
  auto bv = BitVector::CreateWithSize(128, false);
  auto keep = BitVector::CreateWithSize(128, false);
  bv.set(63);
  bv.set(64);
  keep.set(63);
  keep.set(64);
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 2u);
  EXPECT_TRUE(result.is_set(0));
  EXPECT_TRUE(result.is_set(1));
}

TEST(BitVectorTest, CompactWordBoundarySpillover) {
  // Create a scenario where PEXT output from one word spills into the next
  // output word. 128 bits, keep all 64 bits of first word (fills output word
  // completely) plus some from second word.
  auto bv = BitVector::CreateWithSize(128, true);  // all 1s
  auto keep = BitVector::CreateWithSize(128, false);
  // Keep all 64 bits of first word.
  for (uint32_t i = 0; i < 64; ++i) {
    keep.set(i);
  }
  // Keep 10 bits of second word.
  for (uint32_t i = 64; i < 74; ++i) {
    keep.set(i);
  }
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 74u);
  // All kept bits were 1, so result should be all 1s.
  for (uint32_t i = 0; i < 74; ++i) {
    EXPECT_TRUE(result.is_set(i)) << "bit " << i;
  }
}

TEST(BitVectorTest, CompactLargeAlternating) {
  // 256 bits: bv has all bits set, keep has alternating bits.
  // Result should be 128 bits, all set.
  auto bv = BitVector::CreateWithSize(256, true);
  auto keep = BitVector::CreateWithSize(256, false);
  for (uint32_t i = 0; i < 256; i += 2) {
    keep.set(i);
  }
  auto result = std::move(bv).Compact(keep);
  EXPECT_EQ(result.size(), 128u);
  for (uint32_t i = 0; i < 128; ++i) {
    EXPECT_TRUE(result.is_set(i)) << "bit " << i;
  }
}

TEST(BitVectorTest, CompactLargeSelectPattern) {
  // 200 bits. bv = alternating 1010..., keep = every 3rd bit.
  // Verify against naive reference implementation.
  auto bv = BitVector::CreateWithSize(200, false);
  auto keep = BitVector::CreateWithSize(200, false);
  for (uint32_t i = 0; i < 200; i += 2) {
    bv.set(i);
  }
  for (uint32_t i = 0; i < 200; i += 3) {
    keep.set(i);
  }

  // Naive reference: collect bv[i] for each i where keep[i] is set.
  std::vector<bool> expected;
  for (uint32_t i = 0; i < 200; ++i) {
    if (keep.is_set(i)) {
      expected.push_back(bv.is_set(i));
    }
  }

  auto result = std::move(bv).Compact(keep);
  ASSERT_EQ(result.size(), static_cast<uint64_t>(expected.size()));
  for (uint32_t i = 0; i < expected.size(); ++i) {
    EXPECT_EQ(result.is_set(i), expected[i]) << "bit " << i;
  }
}

TEST(BitVectorTest, CompactNonAlignedSize) {
  // Size that doesn't align to 64: 100 bits.
  auto bv = BitVector::CreateWithSize(100, false);
  auto keep = BitVector::CreateWithSize(100, false);
  // Set bits 0, 50, 99 in bv.
  bv.set(0);
  bv.set(50);
  bv.set(99);
  // Keep bits 0, 50, 60, 99.
  keep.set(0);
  keep.set(50);
  keep.set(60);
  keep.set(99);
  // Expected: bv[0]=1, bv[50]=1, bv[60]=0, bv[99]=1 → "1101"
  EXPECT_EQ(BvToStr(std::move(bv).Compact(keep)), "1101");
}

}  // namespace
}  // namespace perfetto::trace_processor::core
