/* stringlib: split implementation */

#ifndef STRINGLIB_FASTSEARCH_H
#error must include "stringlib/fastsearch.h" before including this module
#endif

#include "pycore_list.h"          // _PyList_AppendTakeRef()

/* Overallocate the initial list to reduce the number of reallocs for small
   split sizes.  Eg, "A A A A A A A A A A".split() (10 elements) has three
   resizes, to sizes 4, 8, then 16.  Most observed string splits are for human
   text (roughly 11 words per line) and field delimited data (usually 1-10
   fields).  For large strings the split algorithms are bandwidth limited
   so increasing the preallocation likely will not improve things.*/

#define MAX_PREALLOC 12

/* 5 splits gives 6 elements */
#define PREALLOC_SIZE(maxsplit) \
    (maxsplit >= MAX_PREALLOC ? MAX_PREALLOC : maxsplit+1)

#define SPLIT_APPEND(data, left, right)         \
    sub = STRINGLIB_NEW((data) + (left),        \
                        (right) - (left));      \
    if (sub == NULL)                            \
        goto onError;                           \
    if (_PyList_AppendTakeRef((PyListObject *)list, sub)) \
        goto onError;

#define SPLIT_ADD(data, left, right) {          \
    sub = STRINGLIB_NEW((data) + (left),        \
                        (right) - (left));      \
    if (sub == NULL)                            \
        goto onError;                           \
    if (count < MAX_PREALLOC) {                 \
        PyList_SET_ITEM(list, count, sub);      \
    } else {                                    \
        if (_PyList_AppendTakeRef((PyListObject *)list, sub)) \
            goto onError;                       \
    }                                           \
    count++; }


/* Always force the list to the expected size. */
#define FIX_PREALLOC_SIZE(list) Py_SET_SIZE(list, count)

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(split_whitespace)(PyObject* str_obj,
                           const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                           Py_ssize_t maxcount)
{
    Py_ssize_t i, j, count=0;
    PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
    PyObject *sub;

    if (list == NULL)
        return NULL;

    i = j = 0;
    while (maxcount-- > 0) {
        while (i < str_len && STRINGLIB_ISSPACE(str[i]))
            i++;
        if (i == str_len) break;
        j = i; i++;
        while (i < str_len && !STRINGLIB_ISSPACE(str[i]))
            i++;
#if !STRINGLIB_MUTABLE
        if (j == 0 && i == str_len && STRINGLIB_CHECK_EXACT(str_obj)) {
            /* No whitespace in str_obj, so just use it as list[0] */
            Py_INCREF(str_obj);
            PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
            count++;
            break;
        }
#endif
        SPLIT_ADD(str, j, i);
    }

    if (i < str_len) {
        /* Only occurs when maxcount was reached */
        /* Skip any remaining whitespace and copy to end of string */
        while (i < str_len && STRINGLIB_ISSPACE(str[i]))
            i++;
        if (i != str_len)
            SPLIT_ADD(str, i, str_len);
    }
    FIX_PREALLOC_SIZE(list);
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(split_char)(PyObject* str_obj,
                     const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                     const STRINGLIB_CHAR ch,
                     Py_ssize_t maxcount)
{
    Py_ssize_t i, j, count=0;
    PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
    PyObject *sub;

    if (list == NULL)
        return NULL;

    i = j = 0;
    while ((j < str_len) && (maxcount-- > 0)) {
        for(; j < str_len; j++) {
            /* I found that using memchr makes no difference */
            if (str[j] == ch) {
                SPLIT_ADD(str, i, j);
                i = j = j + 1;
                break;
            }
        }
    }
#if !STRINGLIB_MUTABLE
    if (count == 0 && STRINGLIB_CHECK_EXACT(str_obj)) {
        /* ch not in str_obj, so just use str_obj as list[0] */
        Py_INCREF(str_obj);
        PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
        count++;
    } else
#endif
    if (i <= str_len) {
        SPLIT_ADD(str, i, str_len);
    }
    FIX_PREALLOC_SIZE(list);
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(split)(PyObject* str_obj,
                const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                const STRINGLIB_CHAR* sep, Py_ssize_t sep_len,
                Py_ssize_t maxcount)
{
    Py_ssize_t i, j, pos, count=0;
    PyObject *list, *sub;

    if (sep_len == 0) {
        PyErr_SetString(PyExc_ValueError, "empty separator");
        return NULL;
    }
    else if (sep_len == 1)
        return STRINGLIB(split_char)(str_obj, str, str_len, sep[0], maxcount);

    list = PyList_New(PREALLOC_SIZE(maxcount));
    if (list == NULL)
        return NULL;

    i = j = 0;
    while (maxcount-- > 0) {
        pos = FASTSEARCH(str+i, str_len-i, sep, sep_len, -1, FAST_SEARCH);
        if (pos < 0)
            break;
        j = i + pos;
        SPLIT_ADD(str, i, j);
        i = j + sep_len;
    }
#if !STRINGLIB_MUTABLE
    if (count == 0 && STRINGLIB_CHECK_EXACT(str_obj)) {
        /* No match in str_obj, so just use it as list[0] */
        Py_INCREF(str_obj);
        PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
        count++;
    } else
#endif
    {
        SPLIT_ADD(str, i, str_len);
    }
    FIX_PREALLOC_SIZE(list);
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(rsplit_whitespace)(PyObject* str_obj,
                            const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                            Py_ssize_t maxcount)
{
    Py_ssize_t i, j, count=0;
    PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
    PyObject *sub;

    if (list == NULL)
        return NULL;

    i = j = str_len - 1;
    while (maxcount-- > 0) {
        while (i >= 0 && STRINGLIB_ISSPACE(str[i]))
            i--;
        if (i < 0) break;
        j = i; i--;
        while (i >= 0 && !STRINGLIB_ISSPACE(str[i]))
            i--;
#if !STRINGLIB_MUTABLE
        if (j == str_len - 1 && i < 0 && STRINGLIB_CHECK_EXACT(str_obj)) {
            /* No whitespace in str_obj, so just use it as list[0] */
            Py_INCREF(str_obj);
            PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
            count++;
            break;
        }
#endif
        SPLIT_ADD(str, i + 1, j + 1);
    }

    if (i >= 0) {
        /* Only occurs when maxcount was reached */
        /* Skip any remaining whitespace and copy to beginning of string */
        while (i >= 0 && STRINGLIB_ISSPACE(str[i]))
            i--;
        if (i >= 0)
            SPLIT_ADD(str, 0, i + 1);
    }
    FIX_PREALLOC_SIZE(list);
    if (PyList_Reverse(list) < 0)
        goto onError;
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(rsplit_char)(PyObject* str_obj,
                      const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                      const STRINGLIB_CHAR ch,
                      Py_ssize_t maxcount)
{
    Py_ssize_t i, j, count=0;
    PyObject *list = PyList_New(PREALLOC_SIZE(maxcount));
    PyObject *sub;

    if (list == NULL)
        return NULL;

    i = j = str_len - 1;
    while ((i >= 0) && (maxcount-- > 0)) {
        for(; i >= 0; i--) {
            if (str[i] == ch) {
                SPLIT_ADD(str, i + 1, j + 1);
                j = i = i - 1;
                break;
            }
        }
    }
#if !STRINGLIB_MUTABLE
    if (count == 0 && STRINGLIB_CHECK_EXACT(str_obj)) {
        /* ch not in str_obj, so just use str_obj as list[0] */
        Py_INCREF(str_obj);
        PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
        count++;
    } else
#endif
    if (j >= -1) {
        SPLIT_ADD(str, 0, j + 1);
    }
    FIX_PREALLOC_SIZE(list);
    if (PyList_Reverse(list) < 0)
        goto onError;
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(rsplit)(PyObject* str_obj,
                 const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                 const STRINGLIB_CHAR* sep, Py_ssize_t sep_len,
                 Py_ssize_t maxcount)
{
    Py_ssize_t j, pos, count=0;
    PyObject *list, *sub;

    if (sep_len == 0) {
        PyErr_SetString(PyExc_ValueError, "empty separator");
        return NULL;
    }
    else if (sep_len == 1)
        return STRINGLIB(rsplit_char)(str_obj, str, str_len, sep[0], maxcount);

    list = PyList_New(PREALLOC_SIZE(maxcount));
    if (list == NULL)
        return NULL;

    j = str_len;
    while (maxcount-- > 0) {
        pos = FASTSEARCH(str, j, sep, sep_len, -1, FAST_RSEARCH);
        if (pos < 0)
            break;
        SPLIT_ADD(str, pos + sep_len, j);
        j = pos;
    }
#if !STRINGLIB_MUTABLE
    if (count == 0 && STRINGLIB_CHECK_EXACT(str_obj)) {
        /* No match in str_obj, so just use it as list[0] */
        Py_INCREF(str_obj);
        PyList_SET_ITEM(list, 0, (PyObject *)str_obj);
        count++;
    } else
#endif
    {
        SPLIT_ADD(str, 0, j);
    }
    FIX_PREALLOC_SIZE(list);
    if (PyList_Reverse(list) < 0)
        goto onError;
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}

Py_LOCAL_INLINE(PyObject *)
STRINGLIB(splitlines)(PyObject* str_obj,
                     const STRINGLIB_CHAR* str, Py_ssize_t str_len,
                     int keepends)
{
    /* This does not use the preallocated list because splitlines is
       usually run with hundreds of newlines.  The overhead of
       switching between PyList_SET_ITEM and append causes about a
       2-3% slowdown for that common case.  A smarter implementation
       could move the if check out, so the SET_ITEMs are done first
       and the appends only done when the prealloc buffer is full.
       That's too much work for little gain.*/

    Py_ssize_t i;
    Py_ssize_t j;
    PyObject *list = PyList_New(0);
    PyObject *sub;

    if (list == NULL)
        return NULL;

    for (i = j = 0; i < str_len; ) {
        Py_ssize_t eol;

        /* Find a line and append it */
        while (i < str_len && !STRINGLIB_ISLINEBREAK(str[i]))
            i++;

        /* Skip the line break reading CRLF as one line break */
        eol = i;
        if (i < str_len) {
            if (str[i] == '\r' && i + 1 < str_len && str[i+1] == '\n')
                i += 2;
            else
                i++;
            if (keepends)
                eol = i;
        }
#if !STRINGLIB_MUTABLE
        if (j == 0 && eol == str_len && STRINGLIB_CHECK_EXACT(str_obj)) {
            /* No linebreak in str_obj, so just use it as list[0] */
            if (PyList_Append(list, str_obj))
                goto onError;
            break;
        }
#endif
        SPLIT_APPEND(str, j, eol);
        j = i;
    }
    return list;

  onError:
    Py_DECREF(list);
    return NULL;
}