// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// 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 HIGHWAY_HWY_BIT_SET_H_
#define HIGHWAY_HWY_BIT_SET_H_

// Various BitSet for 64, up to 4096, or any number of bits.

#include <stddef.h>

#include <atomic>

#include "third_party/highway/hwy/base.h"

namespace hwy {

// 64-bit specialization of `std::bitset`, which lacks `Foreach`.
class BitSet64 {
 public:
  constexpr size_t MaxSize() const { return 64; }

  // No harm if `i` is already set.
  void Set(size_t i) {
    HWY_DASSERT(i < 64);
    bits_ |= (1ULL << i);
    HWY_DASSERT(Get(i));
  }

  // Equivalent to Set(i) for i in [0, 64) where (bits >> i) & 1. This does
  // not clear any existing bits.
  void SetNonzeroBitsFrom64(uint64_t bits) { bits_ |= bits; }

  void Clear(size_t i) {
    HWY_DASSERT(i < 64);
    bits_ &= ~(1ULL << i);
  }

  bool Get(size_t i) const {
    HWY_DASSERT(i < 64);
    return (bits_ & (1ULL << i)) != 0;
  }

  // Returns true if Get(i) would return true for any i in [0, 64).
  bool Any() const { return bits_ != 0; }

  // Returns true if Get(i) would return true for all i in [0, 64).
  bool All() const { return bits_ == ~uint64_t{0}; }

  // Returns lowest i such that `Get(i)`. Caller must first ensure `Any()`!
  size_t First() const {
    HWY_DASSERT(Any());
    return Num0BitsBelowLS1Bit_Nonzero64(bits_);
  }

  // Returns lowest i such that `!Get(i)`. Caller must first ensure `!All()`!
  size_t First0() const {
    HWY_DASSERT(!All());
    return Num0BitsBelowLS1Bit_Nonzero64(~bits_);
  }

  // Returns uint64_t(Get(i)) << i for i in [0, 64).
  uint64_t Get64() const { return bits_; }

  // Calls `func(i)` for each `i` in the set. It is safe for `func` to modify
  // the set, but the current Foreach call is unaffected.
  template <class Func>
  void Foreach(const Func& func) const {
    uint64_t remaining_bits = bits_;
    while (remaining_bits != 0) {
      const size_t i = Num0BitsBelowLS1Bit_Nonzero64(remaining_bits);
      remaining_bits &= remaining_bits - 1;  // clear LSB
      func(i);
    }
  }

  size_t Count() const { return PopCount(bits_); }

 private:
  uint64_t bits_ = 0;
};

// Any number of bits, flat array.
template <size_t kMaxSize>
class BitSet {
  static_assert(kMaxSize != 0, "BitSet requires non-zero size");

 public:
  constexpr size_t MaxSize() const { return kMaxSize; }

  // No harm if `i` is already set.
  void Set(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].Set(mod);
  }

  void Clear(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].Clear(mod);
    HWY_DASSERT(!Get(i));
  }

  bool Get(size_t i) const {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    return bits_[idx].Get(mod);
  }

  // Returns true if Get(i) would return true for any i in [0, kMaxSize).
  bool Any() const {
    for (const BitSet64& bits : bits_) {
      if (bits.Any()) return true;
    }
    return false;
  }

  // Returns true if Get(i) would return true for all i in [0, kMaxSize).
  bool All() const {
    for (size_t idx = 0; idx < kNum64 - 1; ++idx) {
      if (!bits_[idx].All()) return false;
    }

    constexpr size_t kRemainder = kMaxSize % 64;
    if (kRemainder == 0) {
      return bits_[kNum64 - 1].All();
    }
    return bits_[kNum64 - 1].Count() == kRemainder;
  }

  // Returns lowest i such that `Get(i)`. Caller must first ensure `Any()`!
  size_t First() const {
    HWY_DASSERT(Any());
    for (size_t idx = 0;; ++idx) {
      HWY_DASSERT(idx < kNum64);
      if (bits_[idx].Any()) return idx * 64 + bits_[idx].First();
    }
  }

  // Returns lowest i such that `!Get(i)`. Caller must first ensure `All()`!
  size_t First0() const {
    HWY_DASSERT(!All());
    for (size_t idx = 0;; ++idx) {
      HWY_DASSERT(idx < kNum64);
      if (!bits_[idx].All()) {
        const size_t first0 = idx * 64 + bits_[idx].First0();
        HWY_DASSERT(first0 < kMaxSize);
        return first0;
      }
    }
  }

  // Calls `func(i)` for each `i` in the set. It is safe for `func` to modify
  // the set, but the current Foreach call is only affected if changing one of
  // the not yet visited BitSet64.
  template <class Func>
  void Foreach(const Func& func) const {
    for (size_t idx = 0; idx < kNum64; ++idx) {
      bits_[idx].Foreach([idx, &func](size_t mod) { func(idx * 64 + mod); });
    }
  }

  size_t Count() const {
    size_t total = 0;
    for (const BitSet64& bits : bits_) {
      total += bits.Count();
    }
    return total;
  }

 private:
  static constexpr size_t kNum64 = DivCeil(kMaxSize, size_t{64});
  BitSet64 bits_[kNum64];
};

// Any number of bits, flat array, atomic updates to the u64.
template <size_t kMaxSize>
class AtomicBitSet {
  static_assert(kMaxSize != 0, "AtomicBitSet requires non-zero size");

  // Bits may signal something to other threads, hence relaxed is insufficient.
  // Acq/Rel ensures a happens-before relationship.
  static constexpr auto kAcq = std::memory_order_acquire;
  static constexpr auto kRel = std::memory_order_release;

 public:
  constexpr size_t MaxSize() const { return kMaxSize; }

  // No harm if `i` is already set.
  void Set(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].fetch_or(1ULL << mod, kRel);
  }

  void Clear(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].fetch_and(~(1ULL << mod), kRel);
    HWY_DASSERT(!Get(i));
  }

  bool Get(size_t i) const {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    return ((bits_[idx].load(kAcq) & (1ULL << mod))) != 0;
  }

  // Returns true if Get(i) would return true for any i in [0, kMaxSize).
  bool Any() const {
    for (const std::atomic<uint64_t>& bits : bits_) {
      if (bits.load(kAcq)) return true;
    }
    return false;
  }

  // Returns true if Get(i) would return true for all i in [0, kMaxSize).
  bool All() const {
    for (size_t idx = 0; idx < kNum64 - 1; ++idx) {
      if (bits_[idx].load(kAcq) != ~uint64_t{0}) return false;
    }

    constexpr size_t kRemainder = kMaxSize % 64;
    const uint64_t last_bits = bits_[kNum64 - 1].load(kAcq);
    if (kRemainder == 0) {
      return last_bits == ~uint64_t{0};
    }
    return PopCount(last_bits) == kRemainder;
  }

  // Returns lowest i such that `Get(i)`. Caller must first ensure `Any()`!
  size_t First() const {
    HWY_DASSERT(Any());
    for (size_t idx = 0;; ++idx) {
      HWY_DASSERT(idx < kNum64);
      const uint64_t bits = bits_[idx].load(kAcq);
      if (bits != 0) {
        return idx * 64 + Num0BitsBelowLS1Bit_Nonzero64(bits);
      }
    }
  }

  // Returns lowest i such that `!Get(i)`. Caller must first ensure `!All()`!
  size_t First0() const {
    HWY_DASSERT(!All());
    for (size_t idx = 0;; ++idx) {
      HWY_DASSERT(idx < kNum64);
      const uint64_t inv_bits = ~bits_[idx].load(kAcq);
      if (inv_bits != 0) {
        const size_t first0 =
            idx * 64 + Num0BitsBelowLS1Bit_Nonzero64(inv_bits);
        HWY_DASSERT(first0 < kMaxSize);
        return first0;
      }
    }
  }

  // Calls `func(i)` for each `i` in the set. It is safe for `func` to modify
  // the set, but the current Foreach call is only affected if changing one of
  // the not yet visited uint64_t.
  template <class Func>
  void Foreach(const Func& func) const {
    for (size_t idx = 0; idx < kNum64; ++idx) {
      uint64_t remaining_bits = bits_[idx].load(kAcq);
      while (remaining_bits != 0) {
        const size_t i = Num0BitsBelowLS1Bit_Nonzero64(remaining_bits);
        remaining_bits &= remaining_bits - 1;  // clear LSB
        func(idx * 64 + i);
      }
    }
  }

  size_t Count() const {
    size_t total = 0;
    for (const std::atomic<uint64_t>& bits : bits_) {
      total += PopCount(bits.load(kAcq));
    }
    return total;
  }

 private:
  static constexpr size_t kNum64 = DivCeil(kMaxSize, size_t{64});
  std::atomic<uint64_t> bits_[kNum64] = {};
};

// Two-level bitset for up to `kMaxSize` <= 4096 values. The iterators
// (`Any/First/Foreach/Count`) are more efficient than `BitSet` for sparse sets.
// This comes at the cost of slightly slower mutators (`Set/Clear`).
template <size_t kMaxSize = 4096>
class BitSet4096 {
  static_assert(kMaxSize != 0, "BitSet4096 requires non-zero size");

 public:
  constexpr size_t MaxSize() const { return kMaxSize; }

  // No harm if `i` is already set.
  void Set(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].Set(mod);
    nonzero_.Set(idx);
    HWY_DASSERT(Get(i));
  }

  // Equivalent to Set(i) for i in [0, 64) where (bits >> i) & 1. This does
  // not clear any existing bits.
  void SetNonzeroBitsFrom64(uint64_t bits) {
    bits_[0].SetNonzeroBitsFrom64(bits);
    if (bits) nonzero_.Set(0);
  }

  void Clear(size_t i) {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    bits_[idx].Clear(mod);
    if (!bits_[idx].Any()) {
      nonzero_.Clear(idx);
    }
    HWY_DASSERT(!Get(i));
  }

  bool Get(size_t i) const {
    HWY_DASSERT(i < kMaxSize);
    const size_t idx = i / 64;
    const size_t mod = i % 64;
    return bits_[idx].Get(mod);
  }

  // Returns true if `Get(i)` would return true for any i in [0, kMaxSize).
  bool Any() const { return nonzero_.Any(); }

  // Returns true if `Get(i)` would return true for all i in [0, kMaxSize).
  bool All() const {
    // Do not check `nonzero_.All()` - that only works if `kMaxSize` is 4096.
    if (nonzero_.Count() != kNum64) return false;
    return Count() == kMaxSize;
  }

  // Returns lowest i such that `Get(i)`. Caller must first ensure `Any()`!
  size_t First() const {
    HWY_DASSERT(Any());
    const size_t idx = nonzero_.First();
    return idx * 64 + bits_[idx].First();
  }

  // Returns lowest i such that `!Get(i)`. Caller must first ensure `!All()`!
  size_t First0() const {
    HWY_DASSERT(!All());
    // It is likely not worthwhile to have a separate `BitSet64` for `not_all_`,
    // hence iterate over all u64.
    for (size_t idx = 0;; ++idx) {
      HWY_DASSERT(idx < kNum64);
      if (!bits_[idx].All()) {
        const size_t first0 = idx * 64 + bits_[idx].First0();
        HWY_DASSERT(first0 < kMaxSize);
        return first0;
      }
    }
  }

  // Returns uint64_t(Get(i)) << i for i in [0, 64).
  uint64_t Get64() const { return bits_[0].Get64(); }

  // Calls `func(i)` for each `i` in the set. It is safe for `func` to modify
  // the set, but the current Foreach call is only affected if changing one of
  // the not yet visited BitSet64 for which Any() is true.
  template <class Func>
  void Foreach(const Func& func) const {
    nonzero_.Foreach([&func, this](size_t idx) {
      bits_[idx].Foreach([idx, &func](size_t mod) { func(idx * 64 + mod); });
    });
  }

  size_t Count() const {
    size_t total = 0;
    nonzero_.Foreach(
        [&total, this](size_t idx) { total += bits_[idx].Count(); });
    return total;
  }

 private:
  static_assert(kMaxSize <= 64 * 64, "One BitSet64 insufficient");
  static constexpr size_t kNum64 = DivCeil(kMaxSize, size_t{64});
  BitSet64 nonzero_;
  BitSet64 bits_[kNum64];
};

}  // namespace hwy

#endif  // HIGHWAY_HWY_BIT_SET_H_