folly/folly/container/detail/F14Mask.h

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

#pragma once

#include <algorithm>
#include <cstdint>

#include <folly/Bits.h>
#include <folly/ConstexprMath.h>
#include <folly/Likely.h>
#include <folly/Portability.h>
#include <folly/container/detail/F14IntrinsicsAvailability.h>
#include <folly/lang/Assume.h>
#include <folly/lang/SafeAssert.h>

#if FOLLY_F14_VECTOR_INTRINSICS_AVAILABLE

namespace folly {
namespace f14 {
namespace detail {

template <typename T>
FOLLY_ALWAYS_INLINE static unsigned findFirstSetNonZero(T mask) { … }

#if FOLLY_NEON
MaskType;

constexpr unsigned kMaskSpacing = …;
#else // FOLLY_SSE >= 2 || FOLLY_RISCV64
using MaskType = uint32_t;

constexpr unsigned kMaskSpacing = 1;
#endif

template <unsigned BitCount>
struct FullMask { … };

template <>
struct FullMask<1> : std::integral_constant<MaskType, 1> { … };

#if FOLLY_ARM
// Mask iteration is different for ARM because that is the only platform
// for which the mask is bigger than a register.

// Iterates a mask, optimized for the case that only a few bits are set
class SparseMaskIter {
  static_assert(kMaskSpacing == 4, "");

  uint32_t interleavedMask_;

 public:
  explicit SparseMaskIter(MaskType mask)
      : interleavedMask_{static_cast<uint32_t>(((mask >> 32) << 2) | mask)} {}

  bool hasNext() { return interleavedMask_ != 0; }

  unsigned next() {
    FOLLY_SAFE_DCHECK(hasNext(), "");
    unsigned i = findFirstSetNonZero(interleavedMask_);
    interleavedMask_ &= (interleavedMask_ - 1);
    return ((i >> 2) | (i << 2)) & 0xf;
  }
};

// Iterates a mask, optimized for the case that most bits are set
class DenseMaskIter {
  static_assert(kMaskSpacing == 4, "");

  std::size_t count_;
  unsigned index_;
  uint8_t const* tags_;

 public:
  explicit DenseMaskIter(uint8_t const* tags, MaskType mask) {
    if (mask == 0) {
      count_ = 0;
    } else {
      count_ = popcount(static_cast<uint32_t>(((mask >> 32) << 2) | mask));
      if (FOLLY_LIKELY((mask & 1) != 0)) {
        index_ = 0;
      } else {
        index_ = findFirstSetNonZero(mask) / kMaskSpacing;
      }
      tags_ = tags;
    }
  }

  bool hasNext() { return count_ > 0; }

  unsigned next() {
    auto rv = index_;
    --count_;
    if (count_ > 0) {
      do {
        ++index_;
      } while ((tags_[index_] & 0x80) == 0);
    }
    FOLLY_SAFE_DCHECK(index_ < 16, "");
    return rv;
  }
};

#else
// Iterates a mask, optimized for the case that only a few bits are set
class SparseMaskIter { … };

// Iterates a mask, optimized for the case that most bits are set
class DenseMaskIter { … };
#endif

// Iterates a mask, returning pairs of [begin,end) index covering blocks
// of set bits
class MaskRangeIter { … };

// Holds the result of an index query that has an optional result,
// interpreting a mask of 0 to be the empty answer and the index of the
// last set bit to be the non-empty answer
class LastOccupiedInMask { … };

// Holds the result of an index query that has an optional result,
// interpreting a mask of 0 to be the empty answer and the index of the
// first set bit to be the non-empty answer
class FirstEmptyInMask { … };

} // namespace detail
} // namespace f14
} // namespace folly

#endif