// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you 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.

// This file is copied from

// https://github.com/ClickHouse/ClickHouse/blob/master/src/Common/StringSearcher.h

// and modified by Doris

#pragma once

#include <stdint.h>

#include <string.h>

#include <algorithm>

#include <limits>

#include <vector>

#include "core/string_ref.h"

#include "exec/common/string_utils/string_utils.h"

#include "util/sse_util.hpp"

namespace doris {

#include "common/compile_check_begin.h"

// namespace ErrorCodes

// {

//     extern const int BAD_ARGUMENTS;

// }

/** Variants for searching a substring in a string.

  */

class StringSearcherBase {

public:

    bool force_fallback = false;

#if defined(__SSE2__) || defined(__aarch64__)

protected:

    static constexpr auto n = sizeof(__m128i);

    const long page_size = sysconf(_SC_PAGESIZE); //::getPageSize();

    bool page_safe(const void* const ptr) const {

        return ((page_size - 1) & reinterpret_cast<std::uintptr_t>(ptr)) <= page_size - n;

    }

#endif

};

/// Performs case-sensitive and case-insensitive search of UTF-8 strings

template <bool CaseSensitive, bool ASCII>

class StringSearcher;

/// Case-sensitive searcher (both ASCII and UTF-8)

template <bool ASCII>

class StringSearcher<true, ASCII> : public StringSearcherBase {

private:

    /// string to be searched for

    const uint8_t* const needle;

    const uint8_t* const needle_end;

    /// first character in `needle`

    uint8_t first {};

#if defined(__SSE4_1__) || defined(__aarch64__)

    uint8_t second {};

    /// vector filled `first` or `second` for determining leftmost position of the first and second symbols

    __m128i first_pattern;

    __m128i second_pattern;

    /// vector of first 16 characters of `needle`

    __m128i cache = _mm_setzero_si128();

    int cachemask {};

#endif

public:

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    StringSearcher(const CharT* needle_, const size_t needle_size)

            : needle {reinterpret_cast<const uint8_t*>(needle_)},

              needle_end {needle + needle_size} {

        if (0 == needle_size) return;

        first = *needle;

#if defined(__SSE4_1__) || defined(__aarch64__)

        first_pattern = _mm_set1_epi8(first);

        if (needle + 1 < needle_end) {

            second = *(needle + 1);

            second_pattern = _mm_set1_epi8(second);

        }

        const auto* needle_pos = needle;

        //for (const auto i : collections::range(0, n))

        for (size_t i = 0; i < n; i++) {

            cache = _mm_srli_si128(cache, 1);

            if (needle_pos != needle_end) {

                cache = _mm_insert_epi8(cache, *needle_pos, n - 1);

                cachemask |= 1 << i;

                ++needle_pos;

            }

        }

#endif

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, size_t haystack_size) const {

        // cast to unsigned int8 to be consitent with needle type

        // ensure unsigned type compare

        return reinterpret_cast<const CharT*>(

                _search(reinterpret_cast<const uint8_t*>(haystack), haystack_size));

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const CharT* haystack_end) const {

        // cast to unsigned int8 to be consitent with needle type

        // ensure unsigned type compare

        return reinterpret_cast<const CharT*>(

                _search(reinterpret_cast<const uint8_t*>(haystack),

                        reinterpret_cast<const uint8_t*>(haystack_end)));

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    ALWAYS_INLINE bool compare(const CharT* haystack, const CharT* haystack_end, CharT* pos) const {

        // cast to unsigned int8 to be consitent with needle type

        // ensure unsigned type compare

        return _compare(reinterpret_cast<const uint8_t*>(haystack),

                        reinterpret_cast<const uint8_t*>(haystack_end),

                        reinterpret_cast<const uint8_t*>(pos));

    }

private:

    ALWAYS_INLINE bool _compare(uint8_t* /*haystack*/, uint8_t* /*haystack_end*/,

                                uint8_t* pos) const {

#if defined(__SSE4_1__) || defined(__aarch64__)

        if (needle_end - needle > n && page_safe(pos)) {

            const auto v_haystack = _mm_loadu_si128(reinterpret_cast<const __m128i*>(pos));

            const auto v_against_cache = _mm_cmpeq_epi8(v_haystack, cache);

            const auto mask = _mm_movemask_epi8(v_against_cache);

            if (0xffff == cachemask) {

                if (mask == cachemask) {

                    pos += n;

                    const auto* needle_pos = needle + n;

                    while (needle_pos < needle_end && *pos == *needle_pos) ++pos, ++needle_pos;

                    if (needle_pos == needle_end) return true;

                }

            } else if ((mask & cachemask) == cachemask)

                return true;

            return false;

        }

#endif

        if (*pos == first) {

            ++pos;

            const auto* needle_pos = needle + 1;

            while (needle_pos < needle_end && *pos == *needle_pos) ++pos, ++needle_pos;

            if (needle_pos == needle_end) return true;

        }

        return false;

    }

    const uint8_t* _search(const uint8_t* haystack, const uint8_t* haystack_end) const {

        if (needle == needle_end) return haystack;

        const auto needle_size = needle_end - needle;

#if defined(__SSE4_1__) || defined(__aarch64__)

        /// Here is the quick path when needle_size is 1.

        if (needle_size == 1) {

            while (haystack < haystack_end) {

                if (haystack + n <= haystack_end && page_safe(haystack)) {

                    const auto v_haystack =

                            _mm_loadu_si128(reinterpret_cast<const __m128i*>(haystack));

                    const auto v_against_pattern = _mm_cmpeq_epi8(v_haystack, first_pattern);

                    const auto mask = _mm_movemask_epi8(v_against_pattern);

                    if (mask == 0) {

                        haystack += n;

                        continue;

                    }

                    const auto offset = __builtin_ctz(mask);

                    haystack += offset;

                    return haystack;

                }

                if (haystack == haystack_end) {

                    return haystack_end;

                }

                if (*haystack == first) {

                    return haystack;

                }

                ++haystack;

            }

            return haystack_end;

        }

#endif

        while (haystack < haystack_end && haystack_end - haystack >= needle_size) {

#if defined(__SSE4_1__) || defined(__aarch64__)

            if ((haystack + 1 + n) <= haystack_end && page_safe(haystack)) {

                /// find first and second characters

                const auto v_haystack_block_first =

                        _mm_loadu_si128(reinterpret_cast<const __m128i*>(haystack));

                const auto v_haystack_block_second =

                        _mm_loadu_si128(reinterpret_cast<const __m128i*>(haystack + 1));

                const auto v_against_pattern_first =

                        _mm_cmpeq_epi8(v_haystack_block_first, first_pattern);

                const auto v_against_pattern_second =

                        _mm_cmpeq_epi8(v_haystack_block_second, second_pattern);

                const auto mask = _mm_movemask_epi8(

                        _mm_and_si128(v_against_pattern_first, v_against_pattern_second));

                /// first and second characters not present in 16 octets starting at `haystack`

                if (mask == 0) {

                    haystack += n;

                    continue;

                }

                const auto offset = __builtin_ctz(mask);

                haystack += offset;

                if (haystack + n <= haystack_end && page_safe(haystack)) {

                    /// check for first 16 octets

                    const auto v_haystack_offset =

                            _mm_loadu_si128(reinterpret_cast<const __m128i*>(haystack));

                    const auto v_against_cache = _mm_cmpeq_epi8(v_haystack_offset, cache);

                    const auto mask_offset = _mm_movemask_epi8(v_against_cache);

                    if (0xffff == cachemask) {

                        if (mask_offset == cachemask) {

                            const auto* haystack_pos = haystack + n;

                            const auto* needle_pos = needle + n;

                            while (haystack_pos < haystack_end && needle_pos < needle_end &&

                                   *haystack_pos == *needle_pos)

                                ++haystack_pos, ++needle_pos;

                            if (needle_pos == needle_end) return haystack;

                        }

                    } else if ((mask_offset & cachemask) == cachemask)

                        return haystack;

                    ++haystack;

                    continue;

                }

            }

#endif

            if (haystack == haystack_end) return haystack_end;

            if (*haystack == first) {

                const auto* haystack_pos = haystack + 1;

                const auto* needle_pos = needle + 1;

                while (haystack_pos < haystack_end && needle_pos < needle_end &&

                       *haystack_pos == *needle_pos)

                    ++haystack_pos, ++needle_pos;

                if (needle_pos == needle_end) return haystack;

            }

            ++haystack;

        }

        return haystack_end;

    }

    const uint8_t* _search(const uint8_t* haystack, const size_t haystack_size) const {

        return _search(haystack, haystack + haystack_size);

    }

};

// Searches for needle surrounded by token-separators.

// Separators are anything inside ASCII (0-128) and not alphanum.

// Any value outside of basic ASCII (>=128) is considered a non-separator symbol, hence UTF-8 strings

// should work just fine. But any Unicode whitespace is not considered a token separtor.

template <typename StringSearcher>

class TokenSearcher : public StringSearcherBase {

    StringSearcher searcher;

    size_t needle_size;

public:

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    TokenSearcher(const CharT* needle_, const size_t needle_size_)

            : searcher {needle_, needle_size_}, needle_size(needle_size_) {

        if (std::any_of(needle_, needle_ + needle_size_, isTokenSeparator)) {

            //throw Exception{"Needle must not contain whitespace or separator characters", ErrorCodes::BAD_ARGUMENTS};

        }

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    ALWAYS_INLINE bool compare(const CharT* haystack, const CharT* haystack_end,

                               const CharT* pos) const {

        // use searcher only if pos is in the beginning of token and pos + searcher.needle_size is end of token.

        if (isToken(haystack, haystack_end, pos))

            return searcher.compare(haystack, haystack_end, pos);

        return false;

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const CharT* const haystack_end) const {

        // use searcher.search(), then verify that returned value is a token

        // if it is not, skip it and re-run

        const auto* pos = haystack;

        while (pos < haystack_end) {

            pos = searcher.search(pos, haystack_end);

            if (pos == haystack_end || isToken(haystack, haystack_end, pos)) return pos;

            // assuming that heendle does not contain any token separators.

            pos += needle_size;

        }

        return haystack_end;

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const size_t haystack_size) const {

        return search(haystack, haystack + haystack_size);

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    ALWAYS_INLINE bool isToken(const CharT* haystack, const CharT* const haystack_end,

                               const CharT* p) const {

        return (p == haystack || isTokenSeparator(*(p - 1))) &&

               (p + needle_size >= haystack_end || isTokenSeparator(*(p + needle_size)));

    }

    ALWAYS_INLINE static bool isTokenSeparator(const uint8_t c) {

        return !(is_alpha_numeric_ascii(c) || !is_ascii(c));

    }

};

using ASCIICaseSensitiveStringSearcher = StringSearcher<true, true>;

// using ASCIICaseInsensitiveStringSearcher = StringSearcher<false, true>;

using UTF8CaseSensitiveStringSearcher = StringSearcher<true, false>;

// using UTF8CaseInsensitiveStringSearcher = StringSearcher<false, false>;

using ASCIICaseSensitiveTokenSearcher = TokenSearcher<ASCIICaseSensitiveStringSearcher>;

// using ASCIICaseInsensitiveTokenSearcher = TokenSearcher<ASCIICaseInsensitiveStringSearcher>;

/** Uses functions from libc.

  * It makes sense to use only with short haystacks when cheap initialization is required.

  * There is no option for case-insensitive search for UTF-8 strings.

  * It is required that strings are zero-terminated.

  */

struct LibCASCIICaseSensitiveStringSearcher : public StringSearcherBase {

    const char* const needle;

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    LibCASCIICaseSensitiveStringSearcher(const CharT* const needle_, const size_t /* needle_size */)

            : needle(reinterpret_cast<const char*>(needle_)) {}

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const CharT* const haystack_end) const {

        const auto* res = strstr(reinterpret_cast<const char*>(haystack),

                                 reinterpret_cast<const char*>(needle));

        if (!res) return haystack_end;

        return reinterpret_cast<const CharT*>(res);

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const size_t haystack_size) const {

        return search(haystack, haystack + haystack_size);

    }

};

struct LibCASCIICaseInsensitiveStringSearcher : public StringSearcherBase {

    const char* const needle;

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    LibCASCIICaseInsensitiveStringSearcher(const CharT* const needle_,

                                           const size_t /* needle_size */)

            : needle(reinterpret_cast<const char*>(needle_)) {}

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const CharT* const haystack_end) const {

        const auto* res = strcasestr(reinterpret_cast<const char*>(haystack),

                                     reinterpret_cast<const char*>(needle));

        if (!res) return haystack_end;

        return reinterpret_cast<const CharT*>(res);

    }

    template <typename CharT>

    // requires (sizeof(CharT) == 1)

    const CharT* search(const CharT* haystack, const size_t haystack_size) const {

        return search(haystack, haystack + haystack_size);

    }

};

#include "common/compile_check_end.h"

} // namespace doris