/// @file dpf/packed_array.hpp
/// @brief Dynamic array of `dpf::twobit` or `dpf::nyble` lanes.
/// @details Storage is the same bit layout as a packed leaf, so
///          `store_leaf_bytes` can memcpy a node at lane `index`.
///          Iterators hold a pointer to that storage and stay valid when
///          the owning buffer is moved.

#ifndef LIBDPF_INCLUDE_DPF_PACKED_ARRAY_HPP__
#define LIBDPF_INCLUDE_DPF_PACKED_ARRAY_HPP__

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <new>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/aligned_allocator.hpp"
#include "dpf/nyble.hpp"
#include "dpf/secret_share.hpp"
#include "dpf/twobit.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

template <typename LaneT>
class dynamic_packed_array
{
    static constexpr std::size_t lane_bits = utils::packed_lane_bits_v<LaneT>;
    static_assert(lane_bits == 2 || lane_bits == 4,
        "dynamic_packed_array lanes are 2 or 4 bits");
    static constexpr std::size_t lanes_per_word = 64u / lane_bits;
    static constexpr unsigned lane_mask = (1u << lane_bits) - 1u;

  public:
    using value_type = LaneT;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using word_type = std::uint64_t;

  private:
    using allocator = aligned_allocator<word_type, utils::max_align_v>;
    using unique_ptr = typename allocator::unique_ptr;

    class lane_ref
    {
      public:
        HEDLEY_NO_THROW
        lane_ref(word_type * word, unsigned shift) noexcept
          : word_{word}, shift_{shift} {}

        HEDLEY_NO_THROW
        HEDLEY_ALWAYS_INLINE
        operator LaneT() const noexcept
        {
            return static_cast<LaneT>((*word_ >> shift_) & lane_mask);
        }

        HEDLEY_NO_THROW
        HEDLEY_ALWAYS_INLINE
        lane_ref & operator=(LaneT value) noexcept
        {
            const auto val = static_cast<word_type>(
                static_cast<unsigned>(value) & lane_mask);
            *word_ = (*word_ & ~(static_cast<word_type>(lane_mask) << shift_))
                | (val << shift_);
            return *this;
        }

        HEDLEY_NO_THROW
        lane_ref & operator=(const lane_ref & other) noexcept
        {
            return (*this = static_cast<LaneT>(other));
        }

        HEDLEY_NO_THROW
        friend bool operator==(lane_ref lhs, LaneT rhs) noexcept
        {
            return static_cast<LaneT>(lhs) == rhs;
        }
        HEDLEY_NO_THROW
        friend bool operator==(LaneT lhs, lane_ref rhs) noexcept
        {
            return rhs == lhs;
        }
        HEDLEY_NO_THROW
        friend bool operator!=(lane_ref lhs, LaneT rhs) noexcept
        {
            return !(lhs == rhs);
        }
        HEDLEY_NO_THROW
        friend bool operator!=(LaneT lhs, lane_ref rhs) noexcept
        {
            return !(rhs == lhs);
        }

      private:
        word_type * word_;
        unsigned shift_;
    };

  public:
    using reference = lane_ref;

    class iterator
    {
      public:
        using iterator_category = std::random_access_iterator_tag;
        using value_type = LaneT;
        using difference_type = std::ptrdiff_t;
        using pointer = void;
        using reference = lane_ref;

        HEDLEY_NO_THROW
        iterator() noexcept = default;
        HEDLEY_NO_THROW
        iterator(word_type * data, size_type index) noexcept
          : data_{data}, index_{index} {}

        HEDLEY_NO_THROW
        lane_ref operator*() const noexcept { return ref_at(index_); }
        HEDLEY_NO_THROW
        lane_ref operator[](difference_type n) const noexcept
        {
            return ref_at(static_cast<size_type>(
                static_cast<difference_type>(index_) + n));
        }

        HEDLEY_NO_THROW
        iterator & operator++() noexcept { ++index_; return *this; }
        HEDLEY_NO_THROW
        iterator operator++(int) noexcept
        {
            iterator prev = *this;
            ++*this;
            return prev;
        }
        HEDLEY_NO_THROW
        iterator & operator--() noexcept { --index_; return *this; }
        HEDLEY_NO_THROW
        iterator operator--(int) noexcept
        {
            iterator prev = *this;
            --*this;
            return prev;
        }

        HEDLEY_NO_THROW
        iterator & operator+=(difference_type n) noexcept
        {
            index_ = static_cast<size_type>(
                static_cast<difference_type>(index_) + n);
            return *this;
        }
        HEDLEY_NO_THROW
        iterator & operator-=(difference_type n) noexcept
        {
            return *this += -n;
        }
        HEDLEY_NO_THROW
        friend iterator operator+(iterator it, difference_type n) noexcept
        {
            it += n;
            return it;
        }
        HEDLEY_NO_THROW
        friend iterator operator+(difference_type n, iterator it) noexcept
        {
            return it + n;
        }
        HEDLEY_NO_THROW
        friend iterator operator-(iterator it, difference_type n) noexcept
        {
            it -= n;
            return it;
        }
        HEDLEY_NO_THROW
        friend difference_type operator-(iterator a, iterator b) noexcept
        {
            return static_cast<difference_type>(a.index_)
                - static_cast<difference_type>(b.index_);
        }

        HEDLEY_NO_THROW
        friend bool operator==(iterator a, iterator b) noexcept
        {
            return a.index_ == b.index_;
        }
        HEDLEY_NO_THROW
        friend bool operator!=(iterator a, iterator b) noexcept
        {
            return !(a == b);
        }
        HEDLEY_NO_THROW
        friend bool operator<(iterator a, iterator b) noexcept
        {
            return a.index_ < b.index_;
        }
        HEDLEY_NO_THROW
        friend bool operator>(iterator a, iterator b) noexcept { return b < a; }
        HEDLEY_NO_THROW
        friend bool operator<=(iterator a, iterator b) noexcept { return !(b < a); }
        HEDLEY_NO_THROW
        friend bool operator>=(iterator a, iterator b) noexcept { return !(a < b); }

      private:
        HEDLEY_NO_THROW
        lane_ref ref_at(size_type index) const noexcept
        {
            const size_type bit = index * lane_bits;
            return lane_ref(data_ + (bit / 64u),
                static_cast<unsigned>(bit % 64u));
        }

        word_type * data_ = nullptr;
        size_type index_ = 0;
    };

    using const_iterator = iterator;

    explicit dynamic_packed_array(size_type nlanes)
      : nlanes_{nlanes},
        nwords_{utils::quotient_ceiling(nlanes, lanes_per_word)}
    {
        if (nwords_ == 0)
        {
            return;
        }
        data_ = allocator{}.allocate_unique_ptr(nwords_);
        if (HEDLEY_UNLIKELY(data_ == nullptr))
        {
            throw std::bad_alloc{};
        }
        std::fill_n(data_.get(), nwords_, word_type{0});
    }

    dynamic_packed_array(const dynamic_packed_array &) = delete;
    dynamic_packed_array & operator=(const dynamic_packed_array &) = delete;

    HEDLEY_NO_THROW
    dynamic_packed_array(dynamic_packed_array && other) noexcept
      : nlanes_{std::exchange(other.nlanes_, 0)},
        nwords_{std::exchange(other.nwords_, 0)},
        data_{std::move(other.data_)}
    {}

    HEDLEY_NO_THROW
    dynamic_packed_array & operator=(dynamic_packed_array && other) noexcept
    {
        if (this != &other)
        {
            nlanes_ = std::exchange(other.nlanes_, 0);
            nwords_ = std::exchange(other.nwords_, 0);
            data_ = std::move(other.data_);
        }
        return *this;
    }

    ~dynamic_packed_array() = default;

    HEDLEY_NO_THROW
    size_type size() const noexcept { return nlanes_; }
    HEDLEY_NO_THROW
    bool empty() const noexcept { return nlanes_ == 0; }
    HEDLEY_NO_THROW
    size_type data_length() const noexcept { return nwords_; }

    HEDLEY_NO_THROW
    word_type * data() noexcept { return data_.get(); }
    HEDLEY_NO_THROW
    const word_type * data() const noexcept { return data_.get(); }

    HEDLEY_NO_THROW
    LaneT operator[](size_type i) const noexcept
    {
        assert(i < nlanes_);
        const size_type bit = i * lane_bits;
        const unsigned shift = static_cast<unsigned>(bit % 64u);
        return static_cast<LaneT>((data_[bit / 64u] >> shift) & lane_mask);
    }

    HEDLEY_NO_THROW
    lane_ref operator[](size_type i) noexcept
    {
        assert(i < nlanes_);
        const size_type bit = i * lane_bits;
        return lane_ref(data_.get() + (bit / 64u),
            static_cast<unsigned>(bit % 64u));
    }

    HEDLEY_NO_THROW
    iterator begin() noexcept { return iterator{data(), 0}; }
    HEDLEY_NO_THROW
    iterator end() noexcept { return iterator{data(), nlanes_}; }
    HEDLEY_NO_THROW
    iterator begin() const noexcept { return iterator{data_.get(), 0}; }
    HEDLEY_NO_THROW
    iterator end() const noexcept { return iterator{data_.get(), nlanes_}; }
    HEDLEY_NO_THROW
    iterator cbegin() const noexcept { return begin(); }
    HEDLEY_NO_THROW
    iterator cend() const noexcept { return end(); }

  private:
    size_type nlanes_ = 0;
    size_type nwords_ = 0;
    unique_ptr data_{};
};

/// Packed lane storage whose iterators yield `subtractive_share<LaneT, Party>`.
/// The bytes are the leaf image (`store_leaf_bytes`); each lane is one share.
template <typename LaneT, std::size_t Party>
class packed_share_output : public dynamic_packed_array<LaneT>
{
    using lanes = dynamic_packed_array<LaneT>;
    using share_type = subtractive_share<LaneT, Party>;

  public:
    using value_type = share_type;
    using size_type = typename lanes::size_type;
    using difference_type = typename lanes::difference_type;

    class reference
    {
      public:
        HEDLEY_NO_THROW
        explicit reference(typename lanes::reference lane) noexcept : lane_{lane} {}

        HEDLEY_NO_THROW
        operator share_type() const noexcept
        {
            return share_type::from_raw(static_cast<LaneT>(lane_));
        }

        HEDLEY_NO_THROW
        reference & operator=(const share_type & share) noexcept
        {
            lane_ = share.raw();
            return *this;
        }

        HEDLEY_NO_THROW
        reference & operator=(LaneT value) noexcept
        {
            lane_ = value;
            return *this;
        }

        HEDLEY_NO_THROW
        reference & operator=(const reference & other) noexcept
        {
            return (*this = static_cast<share_type>(other));
        }

      private:
        typename lanes::reference lane_;
    };

    class iterator
    {
      public:
        using iterator_category = std::random_access_iterator_tag;
        using value_type = share_type;
        using difference_type = std::ptrdiff_t;
        using pointer = void;
        using reference = share_type;

        HEDLEY_NO_THROW
        iterator() noexcept = default;
        HEDLEY_NO_THROW
        explicit iterator(typename lanes::iterator it) noexcept : it_{it} {}

        HEDLEY_NO_THROW
        share_type operator*() const noexcept
        {
            return share_type::from_raw(static_cast<LaneT>(*it_));
        }
        HEDLEY_NO_THROW
        share_type operator[](difference_type n) const noexcept
        {
            return share_type::from_raw(static_cast<LaneT>(it_[n]));
        }

        HEDLEY_NO_THROW
        iterator & operator++() noexcept { ++it_; return *this; }
        HEDLEY_NO_THROW
        iterator operator++(int) noexcept { iterator p = *this; ++*this; return p; }
        HEDLEY_NO_THROW
        iterator & operator--() noexcept { --it_; return *this; }
        HEDLEY_NO_THROW
        iterator operator--(int) noexcept { iterator p = *this; --*this; return p; }
        HEDLEY_NO_THROW
        iterator & operator+=(difference_type n) noexcept { it_ += n; return *this; }
        HEDLEY_NO_THROW
        iterator & operator-=(difference_type n) noexcept { it_ -= n; return *this; }

        HEDLEY_NO_THROW
        friend iterator operator+(iterator it, difference_type n) noexcept
        {
            it += n;
            return it;
        }
        HEDLEY_NO_THROW
        friend iterator operator+(difference_type n, iterator it) noexcept
        {
            return it + n;
        }
        HEDLEY_NO_THROW
        friend iterator operator-(iterator it, difference_type n) noexcept
        {
            it -= n;
            return it;
        }
        HEDLEY_NO_THROW
        friend difference_type operator-(iterator a, iterator b) noexcept
        {
            return a.it_ - b.it_;
        }
        HEDLEY_NO_THROW
        friend bool operator==(iterator a, iterator b) noexcept { return a.it_ == b.it_; }
        HEDLEY_NO_THROW
        friend bool operator!=(iterator a, iterator b) noexcept { return !(a == b); }
        HEDLEY_NO_THROW
        friend bool operator<(iterator a, iterator b) noexcept { return a.it_ < b.it_; }
        HEDLEY_NO_THROW
        friend bool operator>(iterator a, iterator b) noexcept { return b < a; }
        HEDLEY_NO_THROW
        friend bool operator<=(iterator a, iterator b) noexcept { return !(b < a); }
        HEDLEY_NO_THROW
        friend bool operator>=(iterator a, iterator b) noexcept { return !(a < b); }

      private:
        typename lanes::iterator it_{};
    };

    using const_iterator = iterator;

    explicit packed_share_output(size_type nlanes) : lanes(nlanes) {}

    packed_share_output(const packed_share_output &) = delete;
    packed_share_output & operator=(const packed_share_output &) = delete;
    HEDLEY_NO_THROW
    packed_share_output(packed_share_output &&) noexcept = default;
    HEDLEY_NO_THROW
    packed_share_output & operator=(packed_share_output &&) noexcept = default;
    ~packed_share_output() = default;

    using lanes::data;
    using lanes::empty;
    using lanes::size;

    HEDLEY_NO_THROW
    reference operator[](size_type i) noexcept
    {
        return reference{lanes::operator[](i)};
    }
    HEDLEY_NO_THROW
    share_type operator[](size_type i) const noexcept
    {
        return share_type::from_raw(lanes::operator[](i));
    }

    HEDLEY_NO_THROW
    iterator begin() noexcept { return iterator{lanes::begin()}; }
    HEDLEY_NO_THROW
    iterator end() noexcept { return iterator{lanes::end()}; }
    HEDLEY_NO_THROW
    iterator begin() const noexcept
    {
        return iterator{typename lanes::iterator{this->data(), 0}};
    }
    HEDLEY_NO_THROW
    iterator end() const noexcept
    {
        return iterator{typename lanes::iterator{this->data(), this->size()}};
    }
    HEDLEY_NO_THROW
    iterator cbegin() const noexcept { return begin(); }
    HEDLEY_NO_THROW
    iterator cend() const noexcept { return end(); }
};

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PACKED_ARRAY_HPP__