/*
 * 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 <array>
#include <cstddef>
#include <cstdint>

#include <folly/Portability.h>
#include <folly/Utility.h>
#include <folly/algorithm/simd/Movemask.h>
#include <folly/container/SparseByteSet.h>
#include <folly/container/span.h>
#include <folly/lang/Align.h>
#include <folly/lang/Bits.h>
#include <folly/lang/Hint.h>

#if FOLLY_SSE
#include <immintrin.h>
#endif

#if FOLLY_NEON
#include <arm_neon.h>
#endif

#if FOLLY_ARM_FEATURE_SVE
#include <arm_sve.h>
#if __has_include(<arm_neon_sve_bridge.h>)
#include <arm_neon_sve_bridge.h> // @manual
#endif
#endif

namespace folly::simd {

namespace detail {

/// stdfind_scalar_finder_first_of
///
/// A find-first-of finder which simply wraps std::find.
template <typename CharT>
class stdfind_scalar_finder_first_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;

  alignas(sizeof(view)) view const alphabet_;

 public:
  constexpr explicit stdfind_scalar_finder_first_of(
      view const alphabet) noexcept
      : alphabet_{alphabet} {}

  size_t operator()(view const input, size_t const pos = 0) const noexcept {
    auto const r = std::find_first_of(
        input.subspan(pos).begin(),
        input.end(),
        alphabet_.begin(),
        alphabet_.end());
    return r - input.begin();
  }
};

/// default_scalar_finder_first_of
///
/// A find-first-of finder which, for each element of the input, iterates the
/// search alphabet. Has complexity O(MN), with M the length of the alphabet and
/// with N the length of the input.
///
/// Requires no precomputation or storage.
template <typename CharT, bool Eq>
class default_scalar_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;

  alignas(sizeof(view)) view alphabet_;

  bool match(value_type const c) const noexcept {
    bool ret = !Eq;
    for (auto const a : alphabet_) {
      auto const v = a == c;
      ret = Eq ? ret || v : ret && !v;
    }
    return ret;
  }

 public:
  constexpr explicit default_scalar_finder_first_op_of(
      view const alphabet) noexcept
      : alphabet_{alphabet} {}

  size_t operator()(view const input, size_t const pos = 0) const noexcept {
    for (size_t i = pos; i < input.size(); ++i) {
      if (match(input[i])) {
        return i;
      }
    }
    return input.size();
  }
};

/// ltindex_scalar_finder_first_of
///
/// A find-first-of finder which, for each element of the input, looks up that
/// element in a lookup table. Has complexity O(N), with N the length of the
/// input.
///
/// Precomputes and stores a 256-byte lookup table. Precomputation has
/// complexity O(M), with M the length of the alphabet.
///
/// Restricted to elements which are 1 byte wide.
template <typename CharT, bool Eq>
class ltindex_scalar_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;
  using index = std::array<bool, 256>;

  static_assert(sizeof(value_type) == 1);

  alignas(hardware_destructive_interference_size) index const ltindex_;

  static constexpr index make_index(view const alphabet) noexcept {
    index ltindex{};
    for (auto const a : alphabet) {
      ltindex[static_cast<uint8_t>(a)] = true;
    }
    return ltindex;
  }

  bool match(value_type const c) const noexcept {
    return Eq == ltindex_[static_cast<uint8_t>(c)];
  }

 public:
  constexpr explicit ltindex_scalar_finder_first_op_of(
      view const alphabet) noexcept
      : ltindex_{make_index(alphabet)} {}

  size_t operator()(view const input, size_t const pos = 0) const noexcept {
    for (size_t i = pos; i < input.size(); ++i) {
      if (match(input[i])) {
        return i;
      }
    }
    return input.size();
  }
};

/// ltsparse_scalar_finder_first_of
///
/// A find-first-of finder which, for each element of the input, looks up that
/// element in a lookup table. Has complexity O(M+N), with M the length of the
/// alphabet and with N the length of the input.
///
/// Similar to ltindex_scalar_finder_first_of, but where the precomputation is
/// instead done at the beginning of each search using an alternative set type.
/// This alternative set type has lower setup cost but higher lookup cost as
/// compared with the set type in ltindex_scalar_finder_first_of, making this
/// implementation more suitable for unpredictable alphabets.
///
/// Requires no precomputation or storage.
///
/// Restricted to elements which are 1 byte wide.
template <typename CharT, bool Eq>
class ltsparse_scalar_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;

  static_assert(sizeof(value_type) == 1);

  alignas(sizeof(view)) view alphabet_;

  void prep(SparseByteSet& set) const noexcept {
    for (auto const a : alphabet_) {
      set.add(static_cast<uint8_t>(a));
    }
  }

  bool match(value_type const c, SparseByteSet const& set) const noexcept {
    return Eq == set.contains(static_cast<uint8_t>(c));
  }

 public:
  constexpr explicit ltsparse_scalar_finder_first_op_of(
      view const alphabet) noexcept
      : alphabet_{alphabet} {}

  size_t operator()(view const input, size_t const pos = 0) const noexcept {
    [[FOLLY_ATTR_CLANG_UNINITIALIZED]] SparseByteSet set;
    prep(set);
    for (size_t i = pos; i < input.size(); ++i) {
      if (match(input[i], set)) {
        return i;
      }
    }
    return input.size();
  }
};

/// default_vector_finder_first_of
///
/// A find-first-of finder which, for each element of the input, iterates the
/// search alphabet. Has complexity O(MN), with M the length of the alphabet and
/// with N the length of the input.
///
/// Like default_scalar_finder_first_of, but accelerated with simd instructions
/// to search up to 16 elements of the input at a time.
///
/// Requires no precomputation or storage.
///
/// Restricted to elements which are 1 byte wide.
///
/// Implemented for x86-64 and aarch64 architectures.
///
/// Requires a fallback scalar finder for not-implemented-architecture and for
/// near-end-of-input.
template <typename CharT, bool Eq>
class default_vector_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;

  static_assert(sizeof(value_type) == 1);

  alignas(sizeof(view)) view const alphabet_;

 public:
  constexpr explicit default_vector_finder_first_op_of(
      view const alphabet) noexcept
      : alphabet_{alphabet} {}

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      view const input,
      size_t const pos = 0) const noexcept {
    return operator()(scalar, true, input, pos);
  }

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      bool const vector,
      view const input,
      size_t const pos = 0) const noexcept {
    size_t size = pos;
    if (vector) {
#if (FOLLY_SSE >= 2 || (FOLLY_NEON && FOLLY_AARCH64))
      while (input.size() >= size + 16) {
#if FOLLY_SSE
        auto const vhaystack = _mm_loadu_si128(
            reinterpret_cast<__m128i const*>(input.data() + size));
        auto vmask = _mm_set1_epi8(Eq ? 0 : -1);
        for (auto const a : alphabet_) {
          auto const veq = _mm_cmpeq_epi8(vhaystack, _mm_set1_epi8(a));
          vmask = Eq ? _mm_or_si128(veq, vmask) : _mm_andnot_si128(veq, vmask);
        }
#elif FOLLY_NEON
        auto const vhaystack =
            vld1q_u8(reinterpret_cast<uint8_t const*>(input.data() + size));
        auto vmask = vdupq_n_u8(Eq ? 0 : -1);
        for (auto const a : alphabet_) {
          auto const veq = vhaystack == vdupq_n_u8(a);
          vmask = Eq ? veq | vmask : ~veq & vmask;
        }
#endif
        if (auto const [word, bits] = movemask<CharT>(vmask); word) {
          return size + to_signed((findFirstSet(word) - 1) / bits);
        }
        size += 16;
      }
      if (input.size() < size) {
        compiler_may_unsafely_assume_unreachable();
      }
#endif
    }
    return scalar(input, size);
  }
};

/// shuffle_vector_finder_first_of
///
/// A find-first-of finder which, for each element of the input, looks up that
/// element in a lookup table. Has complexity O(MN), with M the length of the
/// alphabet after deduplication and with N the length of the input.
///
/// Precomputes and stores a 256-byte lookup table. Precomputation has
/// complexity O(M), with M the length of the alphabet.
///
/// Like ltindex_scalar_finder_first_of, but accelerated with simd instructions
/// to search up to 16 elements of the input at a time and decelerated by
/// splitting the lookup table into a sequence of lookup tables of length O(M).
///
/// Restricted to elements which are 1 byte wide.
///
/// Implemented for x86-64 and aarch64 architectures.
///
/// Requires a fallback scalar finder for not-implemented-architecture and for
/// near-end-of-input.
template <typename CharT, bool Eq>
class shuffle_vector_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;
  using shufvec = std::array<value_type, 256>;

  struct shuffle {
    shufvec table;
    size_t rounds;
  };

  static_assert(sizeof(value_type) == 1);

  //  invariant: (a in alphabet) <=> (exists k : shufvec[k * 16 + a % 16] = a)
  alignas(hardware_destructive_interference_size) shuffle const shuffle_;

  //  mimic: std::exchange (constexpr), C++20
  template <typename T, typename U = T>
  static constexpr T exchange(T& obj, U&& val) noexcept {
    auto ret = std::move(obj);
    obj = std::forward<U>(val);
    return ret;
  }

  static constexpr shuffle make_shuffle(view const alphabet) noexcept {
    //  init requires: forall k, a : result[k * 16 + a % 16] != a
    shufvec table{1}; // 1, 0, 0, ...
    size_t maxk{};

    std::array<bool, 256> seen{};
    std::array<size_t, 16> lo_seen{};

    for (auto const a : alphabet) {
      auto const v = static_cast<uint8_t>(a);
      if (!exchange(seen[v], true)) {
        auto const k = lo_seen[v % 16]++;
        maxk = maxk < k ? k : maxk;
        table[k * 16 + v % 16] = v;
      }
    }

    return {table, maxk + 1};
  }

 public:
  constexpr explicit shuffle_vector_finder_first_op_of(
      view const alphabet) noexcept
      : shuffle_{make_shuffle(alphabet)} {}

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      view const input,
      size_t const pos = 0) const noexcept {
    return operator()(scalar, true, input, pos);
  }

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      bool const vector,
      view const input,
      size_t const pos = 0) const noexcept {
    size_t size = pos;
    if (vector) {
#if ((FOLLY_SSE >= 2 && FOLLY_SSSE >= 3) || (FOLLY_NEON && FOLLY_AARCH64))
      auto const table = shuffle_.table.data();
      while (input.size() >= size + 16) {
#if FOLLY_SSE
        auto const vtable = reinterpret_cast<__m128i const*>(table);
        auto const vhaystack = _mm_loadu_si128(
            reinterpret_cast<__m128i const*>(input.data() + size));
        auto const vhaystackm = _mm_and_si128(vhaystack, _mm_set1_epi8(15));
        auto vmask = _mm_set1_epi8(Eq ? 0 : -1);
        for (size_t i = 0; i < shuffle_.rounds; ++i) {
          auto const vshuffle = _mm_shuffle_epi8(vtable[i], vhaystackm);
          auto const veq = _mm_cmpeq_epi8(vshuffle, vhaystack);
          vmask = Eq ? _mm_or_si128(veq, vmask) : _mm_andnot_si128(veq, vmask);
        }
#elif FOLLY_NEON
        auto const vtable = reinterpret_cast<uint8x16_t const*>(table);
        auto const vhaystack =
            vld1q_u8(reinterpret_cast<uint8_t const*>(input.data() + size));
        auto vmask = vdupq_n_u8(Eq ? 0 : -1);
        for (size_t i = 0; i < shuffle_.rounds; ++i) {
          auto const veq = vqtbl1q_u8(vtable[i], vhaystack & 15) == vhaystack;
          vmask = Eq ? veq | vmask : ~veq & vmask;
        }
#endif
        if (auto const [word, bits] = movemask<CharT>(vmask); word) {
          return size + to_signed((findFirstSet(word) - 1) / bits);
        }
        size += 16;
      }
      if (input.size() < size) {
        compiler_may_unsafely_assume_unreachable();
      }
#endif
    }
    return scalar(input, size);
  }
};

/// azmatch_vector_finder_first_of
///
/// A find-first-of finder which, for each element of the input, looks up that
/// element in a lookup table. Has complexity O(MN), with M the length of the
/// alphabet after deduplication and with N the length of the input.
///
/// Precomputes and stores a 256-byte lookup table. Precomputation has
/// complexity O(M), with M the length of the alphabet.
///
/// Like ltindex_scalar_finder_first_of, but accelerated with simd instructions
/// to search up to 16 elements of the input at a time and decelerated by
/// splitting the lookup table into a sequence of lookup tables of length O(M).
///
/// Like ltindex_vector_finder_first_of, but with a different technique.
///
/// Restricted to elements which are 1 byte wide.
///
/// Implemented for aarch64 architectures with sve.
///
/// Requires a fallback scalar finder for not-implemented-architecture and for
/// near-end-of-input.
template <typename CharT, bool Eq>
class azmatch_vector_finder_first_op_of {
 private:
  using value_type = CharT;
  using view = span<CharT const>;
  using matchvec = std::array<value_type, 256>;

  struct meta {
    matchvec table;
    size_t rounds;
  };

  static_assert(sizeof(value_type) == 1);

  alignas(hardware_destructive_interference_size) meta const meta_;

  //  mimic: std::exchange (constexpr), C++20
  template <typename T, typename U = T>
  static constexpr T exchange(T& obj, U&& val) noexcept {
    auto ret = std::move(obj);
    obj = std::forward<U>(val);
    return ret;
  }

  static constexpr size_t next_segment(
      view& alphabet, span<value_type, 16> out, span<bool, 256> seen) noexcept {
    if (!alphabet.size()) {
      return 0;
    }
    for (size_t i = 0; i < 16; ++i) {
      out[i] = alphabet[0];
    }
    size_t items = 0;
    while (items < 16 && alphabet.size()) {
      auto const v = static_cast<uint8_t>(alphabet[0]);
      alphabet = alphabet.subspan(1);
      if (!exchange(seen[v], true)) {
        out[items++] = v;
      }
    }
    return items;
  }

  static constexpr meta make_meta(view alphabet) noexcept {
    std::array<bool, 256> seen{};
    size_t rounds = 0;
    matchvec vec{};
    while (true) {
      auto segment = span<value_type, 16>{vec.data() + 16 * rounds, 16};
      auto segsize = next_segment(alphabet, segment, seen);
      if (!segsize) {
        break;
      }
      ++rounds;
    }
    return meta{vec, rounds};
  }

#if FOLLY_ARM_FEATURE_SVE
  static auto svld1_u8_nopred_16(uint8_t const* p) noexcept {
#if __has_include(<arm_neon_sve_bridge.h>)
    return svset_neonq_u8(svundef_u8(), vld1q_u8(p));
#else
    return svld1_u8(svptrue_pat_b8(SV_VL16), p);
#endif
  }
#endif

 public:
  constexpr explicit azmatch_vector_finder_first_op_of(
      view const alphabet) noexcept
      : meta_{make_meta(alphabet)} {}

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      view const input,
      size_t const pos = 0) const noexcept {
    return operator()(scalar, true, input, pos);
  }

  template <typename Scalar>
  size_t operator()(
      Scalar const& scalar,
      bool const vector,
      view const input,
      size_t const pos = 0) const noexcept {
    size_t size = pos;
    if (vector) {
#if FOLLY_ARM_FEATURE_SVE
      auto const table = reinterpret_cast<uint8_t const*>(meta_.table.data());
      while (input.size() >= size + 16) {
        auto const pred = svptrue_b8();
        auto const vhaystack = svld1_u8_nopred_16(
            reinterpret_cast<uint8_t const*>(input.data() + size));
        auto vmask = Eq ? svpfalse_b() : pred;
        for (size_t i = 0; i < meta_.rounds; ++i) {
          auto const vsegment = svld1_u8_nopred_16(table + 16 * i);
          vmask = Eq
              ? svorr_b_z(pred, vmask, svmatch_u8(pred, vhaystack, vsegment))
              : svand_b_z(pred, vmask, svnmatch_u8(pred, vhaystack, vsegment));
        }
        // an important optimization that llvm-17 *could*, but doesn't, do for
        // sve
        if (meta_.rounds == 1) {
          auto const vsegment = svld1_u8_nopred_16(table);
          vmask = Eq
              ? svmatch_u8(pred, vhaystack, vsegment)
              : svnmatch_u8(pred, vhaystack, vsegment);
        }
        auto const count = svcntp_b8(pred, svbrkb_b_z(pred, vmask));
        if (count < 16) {
          return size + count;
        }
        size += 16;
      }
      if (input.size() < size) {
        compiler_may_unsafely_assume_unreachable();
      }
#endif
    }
    return scalar(input, size);
  }
};

} // namespace detail

template <typename CharT>
using basic_stdfind_scalar_finder_first_of =
    detail::stdfind_scalar_finder_first_of<CharT>;

template <typename CharT>
using basic_default_scalar_finder_first_of =
    detail::default_scalar_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_default_scalar_finder_first_not_of =
    detail::default_scalar_finder_first_op_of<CharT, false>;

template <typename CharT>
using basic_ltindex_scalar_finder_first_of =
    detail::ltindex_scalar_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_ltindex_scalar_finder_first_not_of =
    detail::ltindex_scalar_finder_first_op_of<CharT, false>;

template <typename CharT>
using basic_ltsparse_scalar_finder_first_of =
    detail::ltsparse_scalar_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_ltsparse_scalar_finder_first_not_of =
    detail::ltsparse_scalar_finder_first_op_of<CharT, false>;

template <typename CharT>
using basic_default_vector_finder_first_of =
    detail::default_vector_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_default_vector_finder_first_not_of =
    detail::default_vector_finder_first_op_of<CharT, false>;

template <typename CharT>
using basic_shuffle_vector_finder_first_of =
    detail::shuffle_vector_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_shuffle_vector_finder_first_not_of =
    detail::shuffle_vector_finder_first_op_of<CharT, false>;

template <typename CharT>
using basic_azmatch_vector_finder_first_of =
    detail::azmatch_vector_finder_first_op_of<CharT, true>;
template <typename CharT>
using basic_azmatch_vector_finder_first_not_of =
    detail::azmatch_vector_finder_first_op_of<CharT, false>;

using stdfind_scalar_finder_first_of =
    basic_stdfind_scalar_finder_first_of<char>;

using default_scalar_finder_first_of =
    basic_default_scalar_finder_first_of<char>;
using default_scalar_finder_first_not_of =
    basic_default_scalar_finder_first_not_of<char>;

using ltindex_scalar_finder_first_of =
    basic_ltindex_scalar_finder_first_of<char>;
using ltindex_scalar_finder_first_not_of =
    basic_ltindex_scalar_finder_first_not_of<char>;

using ltsparse_scalar_finder_first_of =
    basic_ltsparse_scalar_finder_first_of<char>;
using ltsparse_scalar_finder_first_not_of =
    basic_ltsparse_scalar_finder_first_not_of<char>;

using default_vector_finder_first_of =
    basic_default_vector_finder_first_of<char>;
using default_vector_finder_first_not_of =
    basic_default_vector_finder_first_not_of<char>;

using shuffle_vector_finder_first_of =
    basic_shuffle_vector_finder_first_of<char>;
using shuffle_vector_finder_first_not_of =
    basic_shuffle_vector_finder_first_not_of<char>;

using azmatch_vector_finder_first_of =
    basic_azmatch_vector_finder_first_of<char>;
using azmatch_vector_finder_first_not_of =
    basic_azmatch_vector_finder_first_not_of<char>;

/// composite_finder
///
/// A find-first-of finder which composes a vector finder with a scalar finder.
///
/// A vector finder requires a scalar finder for not-implemented-architecture
/// and for near-end-of-input. This combinator producers a finder which uses the
/// vector finder where possible and otherwise falls back to the scalar finder.
template <typename Vector, typename Scalar>
class composite_finder_first_of : private Vector, Scalar {
 private:
  using view = span<char const>;

 public:
  constexpr explicit composite_finder_first_of(view const alphabet) noexcept
      : Vector{alphabet}, Scalar{alphabet} {}

  size_t operator()(view const input, size_t const pos = 0) const noexcept {
    auto const& vector = static_cast<Vector const&>(*this);
    auto const& scalar = static_cast<Scalar const&>(*this);
    return vector(scalar, input, pos);
  }
};

} // namespace folly::simd