/*
  Copyright (C) 2018 David Robillard <d@drobilla.net>
  Copyright (C) 2006-2007 Chris Hamilton <chamilton@cs.dal.ca>

  This program is free software: you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free Software
  Foundation, either version 2 of the License, or (at your option) any later
  version.

  This program is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
  details.

  You should have received a copy of the GNU General Public License along with
  this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#ifndef CHILBERT_SMALLBITVEC_HPP
#define CHILBERT_SMALLBITVEC_HPP

#include "chilbert/detail/operations.hpp"

#include <cassert>
#include <climits>
#include <cstddef>
#include <cstdint>
#include <iostream>

namespace chilbert {

/** A bit vector small enough to fit in a single word. */
class SmallBitVec
{
public:
	using Rack = uintptr_t;

	static constexpr size_t bits_per_rack = sizeof(Rack) * CHAR_BIT;

	/** Mask for a bit that can be incremented like an index.
	 *
	 * This enables fast iteration while setting or resetting bits since it is
	 * internally stored as a mask which avoids repeated index math.
	 */
	class Mask
	{
	public:
		void operator++() { m_mask <<= 1; }
		void operator--() { m_mask >>= 1; }

		bool operator==(const Mask& mask) const
		{
			return m_mask == mask.m_mask;
		}

		bool operator!=(const Mask& mask) const
		{
			return m_mask == mask.m_mask;
		}

	private:
		friend class SmallBitVec;

		Mask(const size_t index)
		  : m_mask{index < bits_per_rack ? Rack{1} << index : 0}
		{
		}

		Rack m_mask;
	};

	explicit SmallBitVec(const size_t bits)
	  : m_rack{0}
	  , m_size{bits}
	{
		assert(bits <= bits_per_rack);
	}

	SmallBitVec(const size_t bits, const Rack value)
	  : m_rack{value}
	  , m_size{bits}
	{
		assert(bits <= bits_per_rack);
	}

	/// Return the size in bits
	size_t size() const { return m_size; }

	/// Set all bits to one
	SmallBitVec& set()
	{
		m_rack = ~Rack{0} >> (bits_per_rack - m_size);
		return *this;
	}

	/// Set all bits to zero
	SmallBitVec& reset()
	{
		m_rack = 0;
		return *this;
	}

	/// Return the value of the bit covered by `mask`
	bool test(const Mask mask) const { return m_rack & mask.m_mask; }

	/// Return the value of the `index`th bit
	bool test(const size_t index) const { return test(mask(index)); }

	/// Set the bit covered by `mask` to 1
	SmallBitVec& set(const Mask mask)
	{
		m_rack |= mask.m_mask;
		return *this;
	}

	/// Set the `index`th bit to 1
	SmallBitVec& set(const size_t index) { return set(mask(index)); }

	/// Reset the bit covered by `mask` to 0
	SmallBitVec& reset(const Mask mask)
	{
		m_rack &= ~mask.m_mask;
		return *this;
	}

	/// Reset the `index`th bit to 0
	SmallBitVec& reset(const size_t index) { return reset(mask(index)); }

	/// Set the bit covered by `mask` to `value`
	SmallBitVec& set(const Mask mask, const bool value)
	{
		m_rack ^= (-Rack{value} ^ m_rack) & mask.m_mask;
		return *this;
	}

	/// Set the `index`th bit to `value`
	SmallBitVec& set(const size_t index, const bool value)
	{
		return set(mask(index), value);
	}

	/// Flip the value of the bit covered by `mask`
	SmallBitVec& flip(const Mask mask)
	{
		m_rack ^= mask.m_mask;
		return *this;
	}

	/// Flip the value of the `index`th bit
	SmallBitVec& flip(const size_t index) { return flip(mask(index)); }

	bool operator==(const SmallBitVec& vec) const
	{
		return m_rack == vec.m_rack;
	}

	bool operator!=(const SmallBitVec& vec) const
	{
		return m_rack != vec.m_rack;
	}

	bool operator<(const SmallBitVec& vec) const { return m_rack < vec.m_rack; }

	SmallBitVec& operator=(const Rack i)
	{
		m_rack = i;
		return *this;
	}

	SmallBitVec& operator&=(const SmallBitVec& vec)
	{
		m_rack &= vec.m_rack;
		return *this;
	}

	SmallBitVec& operator|=(const SmallBitVec& vec)
	{
		m_rack |= vec.m_rack;
		return *this;
	}

	SmallBitVec& operator^=(const SmallBitVec& vec)
	{
		m_rack ^= vec.m_rack;
		return *this;
	}

	SmallBitVec& operator^=(const Rack i)
	{
		m_rack ^= i;
		return *this;
	}

	SmallBitVec& operator<<=(const size_t bits)
	{
		assert(bits < size());
		m_rack <<= bits;
		return *this;
	}

	SmallBitVec& operator>>=(const size_t bits)
	{
		assert(bits < size());
		m_rack >>= bits;
		return *this;
	}

	/// Right-rotate by `bits` positions
	SmallBitVec& rotr(const size_t bits)
	{
		if (bits > 0 && bits < size()) {
			assert(bits <= bits_per_rack);
			m_rack = (m_rack >> bits) | (m_rack << (size() - bits));
			m_rack &= (~Rack{0} >> (bits_per_rack - size()));
		}
		return *this;
	}

	/// Left-rotate by `bits` positions
	SmallBitVec& rotl(const size_t bits)
	{
		if (bits > 0 && bits < size()) {
			assert(bits <= bits_per_rack);
			m_rack = (m_rack << bits) | (m_rack >> (size() - bits));
			m_rack &= (~Rack{0} >> (bits_per_rack - size()));
		}
		return *this;
	}

	/// Return true iff all bits are zero
	bool none() const { return m_rack == 0; }

	/// Return 1 + the index of the first set bit, or 0 if there are none
	size_t find_first() const
	{
		return static_cast<size_t>(detail::find_first(m_rack));
	}

	/// Return the number of set bits
	size_t count() const
	{
		return static_cast<size_t>(detail::pop_count(m_rack));
	}

	/// Flip all bits (one's complement)
	SmallBitVec& flip()
	{
		m_rack = ~m_rack;
		return *this;
	}

	/// Return the first rack
	Rack& rack() { return m_rack; }
	Rack  rack() const { return m_rack; }

	/// Return a raw pointer to the racks
	Rack*       data() { return &m_rack; }
	const Rack* data() const { return &m_rack; }

	/// Return the number of racks
	size_t num_racks() const { return 1; }

	template <class BitVec>
	class iterator_base : public Mask
	{
	public:
		iterator_base& operator++()
		{
			Mask::operator++();
			return *this;
		}

		iterator_base& operator--()
		{
			Mask::operator--();
			return *this;
		}

		bool operator==(const iterator_base& rhs) const
		{
			return m_vec == rhs.m_vec && Mask::operator==(rhs);
		}

		bool operator!=(const iterator_base& rhs) const
		{
			return !operator==(rhs);
		}

		bool operator*() const { return m_vec->test(*this); }

	protected:
		iterator_base(BitVec& vec, const size_t index)
		  : Mask{index}
		  , m_vec{&vec}
		{
		}

		BitVec* m_vec;
	};

	class iterator : public iterator_base<SmallBitVec>
	{
	public:
		void set() { m_vec->set(*this); }
		void reset() { m_vec->reset(*this); }

	private:
		friend class SmallBitVec;

		iterator(SmallBitVec& vec, const size_t index)
		  : iterator_base{vec, index}
		{
		}
	};

	class const_iterator : public iterator_base<const SmallBitVec>
	{
	private:
		friend class SmallBitVec;

		const_iterator(const SmallBitVec& vec, const size_t index)
		  : iterator_base{vec, index}
		{
		}
	};

	Mask mask(const size_t i = 0) const
	{
		assert(i <= size());
		return Mask{i};
	}

	iterator begin(const size_t i = 0) { return iterator(*this, i); }

	iterator end() { return iterator(*this, size()); }

	const_iterator begin(const size_t i = 0) const
	{
		return const_iterator(*this, i);
	}

	const_iterator end() const { return const_iterator(*this, size()); }

private:
	static_assert(8 * sizeof(Rack) == bits_per_rack, "");
	static_assert((sizeof(Rack) == 4) || (sizeof(Rack) == 8), "");

	Rack   m_rack{};
	size_t m_size{};
};

namespace detail {

template <>
struct is_bitvec<SmallBitVec>
{
	constexpr static bool value = true;
};

template <>
void
gray_code(SmallBitVec& value)
{
	value.rack() ^= (value.rack() >> 1);
}

template <>
void
gray_code_inv(SmallBitVec& value)
{
	gray_code_inv<SmallBitVec::Rack>(value.rack());
}

} // namespace detail

} // namespace chilbert

#endif