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

#ifndef SRC_TRACE_PROCESSOR_CORE_UTIL_BIT_VECTOR_H_
#define SRC_TRACE_PROCESSOR_CORE_UTIL_BIT_VECTOR_H_

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <limits>
#include <utility>

#include "perfetto/base/build_config.h"
#include "perfetto/base/compiler.h"
#include "perfetto/base/logging.h"
#include "perfetto/ext/base/bits.h"
#include "perfetto/public/compiler.h"
#include "src/trace_processor/core/util/flex_vector.h"
#include "src/trace_processor/core/util/slab.h"

#if PERFETTO_BUILDFLAG(PERFETTO_X64_CPU_OPT)
#include <immintrin.h>
#endif

namespace perfetto::trace_processor::core {

// A space-efficient vector for storing and manipulating bit values.
//
// BitVector provides a memory-efficient alternative to vectors of boolean
// values, packing 64 boolean values into a single 64-bit word. This class
// offers efficient operations for common bit manipulation tasks including:
//
// - Setting, clearing, and testing individual bits
// - Efficient iteration over set/unset bits
// - Population counting within bit ranges
// - Filtering operations based on bit patterns
//
// Performance advantages over std::vector<bool> or other alternatives:
// - No branching in the hot path for most operations
// - Uses bitwise operations for maximum efficiency
// - Aligned storage for better memory access patterns
// - Specialized bit manipulation algorithms optimized for performance
struct BitVector {
 public:
  // Default constructor creates an empty bit vector.
  BitVector() = default;

  // Movable but not copyable.
  BitVector(BitVector&&) = default;
  BitVector& operator=(BitVector&&) = default;
  BitVector(const BitVector&) = delete;
  BitVector& operator=(const BitVector&) = delete;

  // Allocates a new BitVector with `size` unset bits.
  //
  // size: Size for how many bits to add to the BitVector.
  // Returns a BitVector with the given size and with all bits set to `value`.
  static BitVector CreateWithSize(uint64_t size, bool value = false) {
    if (size == 0) {
      return {};
    }
    auto words = FlexVector<uint64_t>::CreateWithSize((size + 63u) / 64u);
    memset(words.data(), value ? 0xFF : 0, words.size() * sizeof(uint64_t));
    // Clear any bits past `size`.
    if (size % 64u != 0u) {
      words.back() &= (1ull << (size % 64u)) - 1ull;
    }
    PERFETTO_DCHECK(size <= words.size() * 64u);
    return BitVector(std::move(words), size);
  }

  // Adds a bit to the end of the vector.
  //
  // bit: The boolean value to add to the end of the BitVector.
  PERFETTO_ALWAYS_INLINE void push_back(bool bit) {
    if (PERFETTO_UNLIKELY(size_ % 64ull == 0ull)) {
      words_.push_back(0ull);
    }
    words_[size_ / 64ull] |= static_cast<uint64_t>(bit) << (size_ % 64ull);
    ++size_;
  }

  // Adds n bits with the given value to the end of the vector.
  PERFETTO_ALWAYS_INLINE void push_back_multiple(bool bit, uint64_t count) {
    // TODO(lalitm): This is not the most efficient way to do this, but
    // it is simple and works for now. Consider optimizing this later.
    for (uint64_t i = 0; i < count; ++i) {
      push_back(bit);
    }
  }

  // Changes the value of a bit at the specified index.
  //
  // i: The index of the bit to change.
  // bit: The new boolean value.
  PERFETTO_ALWAYS_INLINE void change(uint64_t i, bool bit) {
    PERFETTO_DCHECK(i < size_);
    uint64_t n = i % 64ull;
    words_[i / 64ull] =
        (words_[i / 64ull] & ~(1ull << n)) | (static_cast<uint64_t>(bit) << n);
  }

  // Changes the value of a bit that is known to be currently unset.
  // This is more efficient than change() when the current bit value is known to
  // be 0.
  //
  // i: The index of the bit to change.
  // bit: The new boolean value.
  PERFETTO_ALWAYS_INLINE void change_assume_unset(size_t i, bool bit) {
    PERFETTO_DCHECK(i < size_);
    uint64_t n = i % 64ull;
    words_[i / 64ull] |= static_cast<uint64_t>(bit) << n;
  }

  // Sets the bit at the specified position to true.
  //
  // i: The index of the bit to set.
  PERFETTO_ALWAYS_INLINE void set(uint64_t i) {
    PERFETTO_DCHECK(i < size_);
    words_[i / 64ull] |= 1ull << (i % 64ull);
  }

  // Sets the bit at the specified position to false.
  //
  // i: The index of the bit to clear.
  PERFETTO_ALWAYS_INLINE void clear(uint64_t i) {
    PERFETTO_DCHECK(i < size_);
    words_[i / 64ull] &= ~(1ull << (i % 64ull));
  }

  // Checks if the bit at the specified position is set.
  //
  // i: The index of the bit to check.
  // Returns true if the bit is set, false otherwise.
  PERFETTO_ALWAYS_INLINE bool is_set(uint64_t i) const {
    PERFETTO_DCHECK(i < size_);
    return (words_[i / 64ull] >> (i % 64ull)) & 1ull;
  }

  // Counts how many bits are set in the same word up to a given position.
  //
  // i: The index position to check up to.
  // Returns the number of set bits in the same 64-bit word as the bit at
  // position i up to position in word (i % 64).
  PERFETTO_ALWAYS_INLINE uint64_t
  count_set_bits_until_in_word(uint64_t i) const {
    PERFETTO_DCHECK(i < size_);
    return static_cast<uint64_t>(
        PERFETTO_POPCOUNT(words_[i / 64ull] & ((1ull << (i % 64ull)) - 1ull)));
  }

  // Counts how many bits are set in the same word as the bit position `i`.
  //
  // i: The index position of the word of interest.
  // Returns the number of set bits in the same 64-bit word as the bit at
  // position i.
  PERFETTO_ALWAYS_INLINE uint64_t count_set_bits_in_word(uint64_t i) const {
    return static_cast<uint64_t>(PERFETTO_POPCOUNT(words_[i / 64ull]));
  }

  // Filters a sequence by keeping only elements whose bit is set (or not set).
  //
  // This function takes a source array and copies elements to a target array
  // only if the corresponding bit in the BitVector matches the desired state.
  //
  // invert: If true, copies elements where the bit is NOT set.
  // source_begin: Pointer to the start of the source array.
  // source_end: Pointer to the end of the source array.
  // target: Pointer to the target array where filtered items will be stored.
  // Returns a pointer to the end of the written data in the target array.
  template <bool invert = false>
  [[nodiscard]] PERFETTO_ALWAYS_INLINE uint32_t* PackLeft(
      const uint32_t* source_begin,
      const uint32_t* source_end,
      uint32_t* target) const {
    uint32_t* out = target;
    for (const uint32_t* s = source_begin; s != source_end; ++s) {
      bool res = (words_[*s / 64ull] >> (*s % 64ull)) & 1ull;
      if constexpr (invert) {
        res = !res;
      }
      *out = *s;
      out += res;
    }
    return out;
  }

  // Computes the prefix sum of set bits for each 64-bit word.
  //
  // This creates an array where each element contains the count of set bits
  // in all preceding 64-bit words, useful for various bit-manipulation
  // algorithms.
  //
  // Returns a Slab containing the prefix sum counts.
  PERFETTO_ALWAYS_INLINE Slab<uint32_t> PrefixPopcount() const {
    Slab<uint32_t> res = Slab<uint32_t>::Alloc((size_ + 63ull) / 64ull);
    ComputePrefixPopcount(res.data());
    return res;
  }

  // Same as PrefixPopcount but returns a FlexVector instead of a Slab.
  PERFETTO_ALWAYS_INLINE FlexVector<uint32_t> PrefixPopcountFlexVector() const {
    auto res = FlexVector<uint32_t>::CreateWithSize((size_ + 63ull) / 64ull);
    ComputePrefixPopcount(res.data());
    return res;
  }

  // Creates a deep copy of this BitVector using word-level memcpy.
  BitVector Clone() const {
    if (size_ == 0) {
      return {};
    }
    auto words = FlexVector<uint64_t>::CreateWithSize((size_ + 63u) / 64u);
    memcpy(words.data(), words_.data(), words.size() * sizeof(uint64_t));
    return BitVector(std::move(words), size_);
  }

  // Gathers only the bits at positions where |keep| is set into a dense
  // result of size popcount(keep). Equivalent to:
  //   result[j] = this[i]  for each i where keep[i] is set (j = 0,1,2,...)
  //
  // Both vectors must have the same size. Uses PEXT per word with bit
  // packing across word boundaries.
  //
  // The && qualifier allows reusing the existing words_ buffer in-place.
  BitVector Compact(const BitVector& keep) && { return CompactInPlace(keep); }

  // Clears all bits in the vector, resetting it to an empty state.
  void clear() {
    words_.clear();
    size_ = 0;
  }

  // Sets all bits to zero without changing the size.
  void ClearAllBits() {
    memset(words_.data(), 0, words_.size() * sizeof(uint64_t));
  }

  // Resizes the vector to the specified size. If shrinking, bits past the new
  // size are cleared. If growing, new bits are set to the given value.
  void resize(uint64_t new_size, bool value = false) {
    uint64_t new_words = (new_size + 63) / 64;
    words_.resize(new_words);
    if (new_size > size_) {
      // Set new bits to value.
      uint64_t fill = value ? ~0ull : 0ull;
      for (uint64_t i = (size_ + 63) / 64; i < new_words; ++i) {
        words_[i] = fill;
      }
      // Handle partial word at old size boundary.
      if (size_ % 64 != 0) {
        uint64_t mask = ~((1ull << (size_ % 64)) - 1);
        if (value) {
          words_[size_ / 64] |= mask;
        } else {
          words_[size_ / 64] &= ~mask;
        }
      }
    }
    // Clear bits past new_size in the last word.
    if (new_size % 64 != 0 && new_words > 0) {
      words_[new_words - 1] &= (1ull << (new_size % 64)) - 1;
    }
    size_ = new_size;
  }

  // Shrinks the memory allocated by the vector to be as small as possible while
  // still maintaining the invariants of the class.
  void shrink_to_fit() { words_.shrink_to_fit(); }

  // Returns the number of bits in the vector.
  PERFETTO_ALWAYS_INLINE uint64_t size() const { return size_; }

 private:
  // Software emulation of the x64 PEXT instruction. Extracts bits from |word|
  // at positions where |mask| has set bits, packing them into the low bits.
  // See https://www.felixcloutier.com/x86/pext for details.
  //
  // Scales with the number of set bits in |mask| rather than being constant
  // time: should be avoided where real instructions are available.
  PERFETTO_ALWAYS_INLINE static uint64_t PextSlow(uint64_t word,
                                                  uint64_t mask) {
    if (word == 0 || mask == std::numeric_limits<uint64_t>::max()) {
      return word;
    }
    uint64_t result = 0;
    for (uint64_t bb = 1; mask; bb += bb) {
      if (word & mask & (0ull - mask)) {
        result |= bb;
      }
      mask &= mask - 1;
    }
    return result;
  }

  // See |PextSlow| for information on PEXT.
  PERFETTO_ALWAYS_INLINE static uint64_t Pext(uint64_t word, uint64_t mask) {
#if PERFETTO_BUILDFLAG(PERFETTO_X64_CPU_OPT)
    base::ignore_result(PextSlow);
    return _pext_u64(word, mask);
#else
    return PextSlow(word, mask);
#endif
  }

  // Packs bits in-place using PEXT, reusing the existing words_ buffer.
  BitVector CompactInPlace(const BitVector& keep) {
    PERFETTO_DCHECK(size_ == keep.size_);

    const uint64_t* cur_word = words_.data();
    const uint64_t* end_word = words_.data() + words_.size();
    const uint64_t* cur_mask = keep.words_.data();

    uint32_t out_word_bits = 0;
    uint64_t* out_word = words_.data();
    uint32_t total_kept = 0;
    for (; cur_word != end_word; ++cur_word, ++cur_mask) {
      uint64_t ext = Pext(*cur_word, *cur_mask);

      *out_word = out_word_bits == 0 ? ext : *out_word | (ext << out_word_bits);

      auto popcount = static_cast<uint32_t>(PERFETTO_POPCOUNT(*cur_mask));
      out_word_bits += popcount;
      total_kept += popcount;

      bool spillover = out_word_bits > 64;
      out_word += out_word_bits >= 64;
      out_word_bits %= 64;

      if (spillover) {
        *out_word = ext >> (popcount - out_word_bits);
      }
    }
    resize(total_kept);
    return std::move(*this);
  }

  // Constructor used by Alloc.
  explicit BitVector(FlexVector<uint64_t> data, uint64_t size)
      : words_(std::move(data)), size_(size) {}

  // Helper to compute prefix popcount into an output buffer.
  PERFETTO_ALWAYS_INLINE void ComputePrefixPopcount(uint32_t* out) const {
    uint32_t accum = 0;
    for (uint32_t i = 0; i < (size_ + 63ull) / 64ull; ++i) {
      out[i] = accum;
      accum += static_cast<uint32_t>(PERFETTO_POPCOUNT(words_[i]));
    }
  }

  // The underlying storage as 64-bit words.
  FlexVector<uint64_t> words_;

  // Number of bits in the vector.
  uint64_t size_ = 0;
};

}  // namespace perfetto::trace_processor::core

#endif  // SRC_TRACE_PROCESSOR_CORE_UTIL_BIT_VECTOR_H_
