/* Copyright 2025 Joaquin M Lopez Munoz.
 * Distributed under the Boost Software License, Version 1.0.
 * (See accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * See https://www.boost.org/libs/bloom for library home page.
 */

#ifndef BOOST_BLOOM_BOBW_MULTIBLOCK64_AVX2_HPP
#define BOOST_BLOOM_BOBW_MULTIBLOCK64_AVX2_HPP

#include <boost/bloom/detail/avx2.hpp>
#include <boost/bloom/detail/block_base.hpp>
#include <boost/bloom/detail/fast_multiblock64_avx2.hpp>
#include <boost/bloom/detail/multiblock_fpr_base.hpp>
#include <boost/bloom/detail/mulx64.hpp>
#include <boost/config.hpp>
#include <cstddef>
#include <cstdint>

namespace boost{
namespace bloom{

#if defined(BOOST_MSVC)
#pragma warning(push)
#pragma warning(disable:4714) /* marked as __forceinline not inlined */
#endif

template<std::size_t K>
struct bobw_multiblock64_avx2:
  public detail::multiblock_fpr_base<K>,
  private detail::block_base<std::uint64_t,K%8?K%8:1>
{
  static constexpr std::size_t k=K;
  using value_type=detail::m256ix2[(k+7)/8];
  static constexpr std::size_t used_value_size=sizeof(std::uint64_t)*k;

  static BOOST_FORCEINLINE void mark(value_type& x,std::uint64_t hash)
  {
    for(int i=0;i<k/8;++i){
      mark_m256ix2(x[i],hash,8);
      hash=detail::mulx64(hash);
    }
    if(k%8){
      std::uint64_t* y=reinterpret_cast<std::uint64_t*>(&x[k/8]);
      std::size_t    j=0;
      loop(hash,[&](std::uint64_t h){y[j++]|=std::uint64_t(1)<<(h&mask);});
    }
  }

  static BOOST_FORCEINLINE bool check(const value_type& x,std::uint64_t hash)
  {
    for(int i=0;i<k/8;++i){
      if(!check_m256ix2(x[i],hash,8))return false;
      hash=detail::mulx64(hash);
    }
    if(k%8){
      const std::uint64_t* y=reinterpret_cast<const std::uint64_t*>(&x[k/8]);
      std::uint64_t        res=1;
      std::size_t          j=0;
      loop(hash,[&](std::uint64_t h){res&=(y[j++]>>(h&mask));});
      return res;
    }
    else return true;
  }

private:
  using super=detail::block_base<std::uint64_t,K%8?K%8:1>;
  using super::mask;
  using super::loop;

  static BOOST_FORCEINLINE detail::m256ix2 make_m256ix2(
    std::uint64_t hash,std::size_t kp)
  {
    const detail::m256ix2 ones[8]={
      {_mm256_set_epi64x(0,0,0,1),_mm256_set_epi64x(0,0,0,0)},
      {_mm256_set_epi64x(0,0,1,1),_mm256_set_epi64x(0,0,0,0)},
      {_mm256_set_epi64x(0,1,1,1),_mm256_set_epi64x(0,0,0,0)},
      {_mm256_set_epi64x(1,1,1,1),_mm256_set_epi64x(0,0,0,0)},
      {_mm256_set_epi64x(1,1,1,1),_mm256_set_epi64x(0,0,0,1)},
      {_mm256_set_epi64x(1,1,1,1),_mm256_set_epi64x(0,0,1,1)},
      {_mm256_set_epi64x(1,1,1,1),_mm256_set_epi64x(0,1,1,1)},
      {_mm256_set_epi64x(1,1,1,1),_mm256_set_epi64x(1,1,1,1)},
    };

    __m256i h=_mm256_set1_epi64x(hash);
    h=_mm256_sllv_epi64(h,_mm256_set_epi64x(18,12,6,0));
    h=_mm256_srli_epi32(h,32-6);
    return {
      _mm256_sllv_epi64(
        ones[kp-1].lo,_mm256_cvtepu32_epi64(_mm256_extracti128_si256(h,0))),
      kp<=4?
      _mm256_set1_epi64x(0):
      _mm256_sllv_epi64(
        ones[kp-1].hi,_mm256_cvtepu32_epi64(_mm256_extracti128_si256(h,1)))
    };
  }

  static BOOST_FORCEINLINE void mark_m256ix2(
    detail::m256ix2& x,std::uint64_t hash,std::size_t kp)
  {
    detail::m256ix2 h=make_m256ix2(hash,kp);
    x.lo=_mm256_or_si256(x.lo,h.lo);
    if(kp>4)x.hi=_mm256_or_si256(x.hi,h.hi);
  }

  static BOOST_FORCEINLINE bool check_m256ix2(
    const detail::m256ix2& x,std::uint64_t hash,std::size_t kp)
  {
    detail::m256ix2 h=make_m256ix2(hash,kp);
    auto res=_mm256_testc_si256(x.lo,h.lo);
    if(kp>4)res&=_mm256_testc_si256(x.hi,h.hi);
    return res;
  }
};

#if defined(BOOST_MSVC)
#pragma warning(pop) /* C4714 */
#endif

} /* namespace bloom */
} /* namespace boost */

#endif