Added BitArrayIterator and range-for over set bits on BitArrayViewBase. (#467)

include/bx/bitarray.h

@@ -12,6 +12,66 @@
 
 namespace bx
 {
+	class BitArrayViewBase;
+
+	/// Walks set bits of a mask, lowest first, by shifting the mask down to the next set bit.
+	///
+	/// @tparam MaskT Unsigned integer type of the mask.
+	///
+	template<typename MaskT>
+	struct BitMaskToIndexIteratorT
+	{
+		///
+		BitMaskToIndexIteratorT(MaskT _mask);
+
+		/// Advances to the next set bit.
+		void next();
+
+		/// Returns true when no set bits remain.
+		bool isDone() const;
+
+		/// Returns index of the current set bit.
+		uint32_t getIndex() const;
+
+		/// Returns remaining mask, shifted so that the current set bit is bit 0.
+		MaskT getMask() const;
+
+	private:
+		MaskT    m_mask;
+		uint32_t m_idx;
+	};
+
+	/// Range-for iterator over set bits of a bit array, in index order. Dereferences to the bit
+	/// index. The current word is cached and walked with BitMaskToIndexIteratorT; the next word is
+	/// looked up with findSet only when the cached one runs out. Clearing the bit under the
+	/// iterator keeps it valid. Other bits changed during iteration are seen only once the walk
+	/// reaches their word.
+	struct BitArrayIterator
+	{
+		///
+		BitArrayIterator(const BitArrayViewBase& _view, uint32_t _idx);
+
+		///
+		uint32_t operator*() const;
+
+		///
+		bool operator!=(const BitArrayIterator& _rhs) const;
+
+		///
+		const BitArrayIterator& operator++();
+
+		/// Advances to the next set bit. Returns false at the end.
+		bool next();
+
+	private:
+		void load(uint32_t _idx);
+
+		const BitArrayViewBase& m_view;
+		BitMaskToIndexIteratorT<uint64_t> m_bits;
+		uint32_t m_word;
+		uint32_t m_idx;
+	};
+
 	/// Base class for immutable bit array views providing read-only access to a contiguous bit set.
 	class BitArrayViewBase
 	{
@@ -62,6 +122,12 @@ namespace bx
 		/// Returns the maximum number of bits the bit set can hold.
 		uint32_t getCapacity() const;
 
+		/// Returns iterator to the first set bit.
+		BitArrayIterator begin() const;
+
+		/// Returns iterator past the last set bit.
+		BitArrayIterator end() const;
+
 		/// Finds a contiguous run of zero bits.
 		///
 		/// @param[in] _numBits Number of contiguous zero bits to find.

include/bx/inline/bitarray.inl

@@ -53,6 +53,118 @@ namespace bx
 		return m_sizeInBytes * 8;
 	}
 
+	inline BitArrayIterator BitArrayViewBase::begin() const
+	{
+		return BitArrayIterator(*this, findSet(0, m_count) );
+	}
+
+	inline BitArrayIterator BitArrayViewBase::end() const
+	{
+		return BitArrayIterator(*this, kInvalid);
+	}
+
+	template<typename MaskT>
+	inline BitMaskToIndexIteratorT<MaskT>::BitMaskToIndexIteratorT(MaskT _mask)
+		: m_mask(0)
+		, m_idx(0)
+	{
+		const uint8_t ntz = countTrailingZeros<MaskT>(_mask);
+		m_mask = 0 == _mask ? MaskT(0) : MaskT(_mask >> ntz);
+		m_idx  = ntz;
+	}
+
+	template<typename MaskT>
+	inline void BitMaskToIndexIteratorT<MaskT>::next()
+	{
+		m_mask  = MaskT(m_mask >> 1);
+		m_idx  += 1;
+
+		const uint8_t ntz = countTrailingZeros<MaskT>(m_mask);
+		m_mask = 0 == m_mask ? MaskT(0) : MaskT(m_mask >> ntz);
+		m_idx += ntz;
+	}
+
+	template<typename MaskT>
+	inline bool BitMaskToIndexIteratorT<MaskT>::isDone() const
+	{
+		return 0 == m_mask;
+	}
+
+	template<typename MaskT>
+	inline uint32_t BitMaskToIndexIteratorT<MaskT>::getIndex() const
+	{
+		return m_idx;
+	}
+
+	template<typename MaskT>
+	inline MaskT BitMaskToIndexIteratorT<MaskT>::getMask() const
+	{
+		return m_mask;
+	}
+
+	inline BitArrayIterator::BitArrayIterator(const BitArrayViewBase& _view, uint32_t _idx)
+		: m_view(_view)
+		, m_bits(0)
+		, m_word(0)
+		, m_idx(kInvalid)
+	{
+		load(_idx);
+	}
+
+	inline uint32_t BitArrayIterator::operator*() const
+	{
+		return m_idx;
+	}
+
+	inline bool BitArrayIterator::operator!=(const BitArrayIterator& _rhs) const
+	{
+		return m_idx != _rhs.m_idx;
+	}
+
+	inline const BitArrayIterator& BitArrayIterator::operator++()
+	{
+		next();
+		return *this;
+	}
+
+	inline void BitArrayIterator::load(uint32_t _idx)
+	{
+		m_idx = _idx;
+
+		if (kInvalid != _idx)
+		{
+			m_word = _idx/64;
+
+			const uint64_t word = reinterpret_cast<const uint64_t*>(m_view.getPtr() )[m_word];
+			m_bits = BitMaskToIndexIteratorT<uint64_t>(word & (~UINT64_C(0) << (_idx & 63) ) );
+		}
+	}
+
+	inline bool BitArrayIterator::next()
+	{
+		if (kInvalid == m_idx)
+		{
+			return false;
+		}
+
+		const uint32_t count = m_view.getCount();
+
+		m_bits.next();
+
+		if (!m_bits.isDone() )
+		{
+			const uint32_t idx = m_word*64 + m_bits.getIndex();
+			m_idx = idx < count ? idx : kInvalid;
+		}
+		else
+		{
+			const uint32_t begin = (m_word + 1)*64;
+			load(begin < count ? m_view.findSet(begin, count) : kInvalid);
+		}
+
+		return kInvalid != m_idx;
+	}
+
 	inline uint32_t BitArrayViewBase::find(uint32_t _numBits) const
 	{
 		return find(0, m_count, _numBits);

tests/bitarray_test.cpp

@@ -1512,3 +1512,143 @@ TEST_CASE("MippedBitArray-random", "[container]")
 		REQUIRE(expected == array.findClear(findBegin, kNumBits) );
 	}
 }
+
+TEST_CASE("BitArray-iterator", "[container]")
+{
+	constexpr uint32_t kNumWords = 40;
+	constexpr uint32_t kNumBits  = kNumWords*64;
+
+	uint64_t bits[kNumWords];
+	bx::MutableBitArrayView bs(bits, sizeof(bits) );
+
+	bs.clear();
+
+	{
+		uint32_t numVisited = 0;
+
+		for (const uint32_t idx : bs)
+		{
+			BX_UNUSED(idx);
+			++numVisited;
+		}
+
+		REQUIRE(0 == numVisited);
+	}
+
+	bs.set(0, true);
+	bs.set(63, true);
+	bs.set(64, true);
+	bs.set(1000, true);
+	bs.set(kNumBits - 1, true);
+
+	{
+		const uint32_t expected[] = { 0, 63, 64, 1000, kNumBits - 1 };
+
+		uint32_t numVisited = 0;
+
+		for (const uint32_t idx : bs)
+		{
+			REQUIRE(numVisited < BX_COUNTOF(expected) );
+			REQUIRE(expected[numVisited] == idx);
+			++numVisited;
+		}
+
+		REQUIRE(BX_COUNTOF(expected) == numVisited);
+	}
+
+	bx::RngMwc rng;
+
+	for (uint32_t iter = 0; iter < 1000; ++iter)
+	{
+		fillRandom(rng, bits, kNumWords);
+
+		const uint32_t count = 1 + rng.gen() % kNumBits;
+		bx::MutableBitArrayView view(bits, sizeof(bits), count);
+
+		uint32_t expected = view.findSet(0, count);
+
+		for (const uint32_t idx : view)
+		{
+			REQUIRE(expected == idx);
+			REQUIRE(view.get(idx) );
+			expected = view.findSet(idx + 1, count);
+		}
+
+		REQUIRE(bx::kInvalid == expected);
+	}
+
+	fillRandom(rng, bits, kNumWords);
+
+	{
+		const uint32_t numSet = bs.countBits();
+
+		uint32_t numVisited = 0;
+
+		for (const uint32_t idx : bs)
+		{
+			++numVisited;
+
+			if (0 == (idx & 1) )
+			{
+				bs.set(idx, false);
+			}
+		}
+
+		REQUIRE(numSet == numVisited);
+
+		for (const uint32_t idx : bs)
+		{
+			REQUIRE(1 == (idx & 1) );
+			bs.set(idx, false);
+		}
+
+		REQUIRE(bs.testNone() );
+	}
+}
+
+template<typename MaskT>
+static void testBitMaskToIndexIterator(MaskT _mask)
+{
+	uint32_t numVisited = 0;
+	uint32_t prevIdx    = 0;
+
+	for (bx::BitMaskToIndexIteratorT<MaskT> it(_mask); !it.isDone(); it.next() )
+	{
+		const uint32_t idx = it.getIndex();
+
+		REQUIRE(idx < sizeof(MaskT)*8);
+		REQUIRE( (0 == numVisited || prevIdx < idx) );
+		REQUIRE(0 != ( (_mask >> idx) & 1) );
+		REQUIRE(1 == (it.getMask() & 1) );
+
+		prevIdx = idx;
+		++numVisited;
+	}
+
+	REQUIRE(bx::countBits<uint64_t>(uint64_t(_mask) ) == numVisited);
+}
+
+TEST_CASE("BitMaskToIndexIteratorT", "[container]")
+{
+	testBitMaskToIndexIterator<uint8_t>(0);
+	testBitMaskToIndexIterator<uint8_t>(1);
+	testBitMaskToIndexIterator<uint8_t>(0x80);
+	testBitMaskToIndexIterator<uint8_t>(0xff);
+	testBitMaskToIndexIterator<uint8_t>(0xa5);
+	testBitMaskToIndexIterator<uint16_t>(0x8001);
+	testBitMaskToIndexIterator<uint32_t>(0);
+	testBitMaskToIndexIterator<uint32_t>(0x80000001);
+	testBitMaskToIndexIterator<uint32_t>(UINT32_MAX);
+	testBitMaskToIndexIterator<uint64_t>(0);
+	testBitMaskToIndexIterator<uint64_t>(1);
+	testBitMaskToIndexIterator<uint64_t>(UINT64_C(0x8000000000000000) );
+	testBitMaskToIndexIterator<uint64_t>(UINT64_MAX);
+
+	bx::RngMwc rng;
+
+	for (uint32_t ii = 0; ii < 1000; ++ii)
+	{
+		testBitMaskToIndexIterator<uint32_t>(rng.gen() );
+		testBitMaskToIndexIterator<uint64_t>( (uint64_t(rng.gen() ) << 32) | rng.gen() );
+	}
+}