<commit_before><commit_msg>Implement getitem, len, eq, ne operations for bytes/unicode array items<commit_after>""" Implements support for bytes and str (unicode) arrays.
"""
import operator
import numpy as np
from llvmlite import ir

from numba import njit, types, cgutils
from numba.extending import overload, intrinsic

# bytes and str arrays items are of type CharSeq and UnicodeCharSeq,
# respectively.  See numpy/types/npytypes.py for CharSeq,
# UnicodeCharSeq definitions.  The corresponding data models are
# defined in numpy/datamodel/models.py

s1_dtype = np.dtype('S1')
assert s1_dtype.itemsize == 1

u1_dtype = np.dtype('U1')
unicode_byte_width = u1_dtype.itemsize
unicode_uint = {1: np.uint8, 2: np.uint16, 4: np.uint32}[unicode_byte_width]


# this is modified version of numba.unicode.make_deref_codegen
def make_deref_codegen(bitsize):
    def codegen(context, builder, signature, args):
        data, idx = args
        # XXX how to access data (of type CharSeq/UnicodeCharSeq)
        # without alloca?
        rawptr = cgutils.alloca_once_value(builder, value=data)
        ptr = builder.bitcast(rawptr, ir.IntType(bitsize).as_pointer())
        ch = builder.load(builder.gep(ptr, [idx]))
        return builder.zext(ch, ir.IntType(32))
    return codegen


@intrinsic
def deref_uint8(typingctx, data, offset):
    sig = types.uint32(data, types.intp)
    return sig, make_deref_codegen(8)


@intrinsic
def deref_uint16(typingctx, data, offset):
    sig = types.uint32(data, types.intp)
    return sig, make_deref_codegen(16)


@intrinsic
def deref_uint32(typingctx, data, offset):
    sig = types.uint32(data, types.intp)
    return sig, make_deref_codegen(32)


@njit(_nrt=False)
def charseq_get_code(a, i):
    """Access i-th item of CharSeq instance via code value
    """
    return deref_uint8(a, i)


@njit
def charseq_get_value(a, i):
    """Access i-th item of CharSeq instance via code value.

    null code is interpreted as IndexError
    """
    code = charseq_get_code(a, i)
    if code == 0:
        raise IndexError('index out of range')
    return code


@njit(_nrt=False)
def unicode_charseq_get_code(a, i):
    """Access i-th item of UnicodeCharSeq instance via code value
    """
    if unicode_byte_width == 4:
        return deref_uint32(a, i)
    elif unicode_byte_width == 2:
        return deref_uint16(a, i)
    elif unicode_byte_width == 1:
        return deref_uint8(a, i)
    else:
        raise NotImplementedError(
            'unicode_charseq_get_code: unicode_byte_width not in [1, 2, 4]')


@njit
def unicode_charseq_get_value(a, i):
    """Access i-th item of UnicodeCharSeq instance via unicode value

    null code is interpreted as IndexError
    """
    code = unicode_charseq_get_code(a, i)
    if code == 0:
        raise IndexError('index out of range')
    # Return numpy equivalent of `chr(code)`
    return np.array(code, unicode_uint).view(u1_dtype)[()]

#
#   Operators on bytes/unicode array items
#


@overload(operator.getitem)
def charseq_getitem(s, i):
    return_impl = False
    if isinstance(i, types.Integer):
        if isinstance(s, types.CharSeq):
            return_impl = True
            get_value = charseq_get_value
        if isinstance(s, types.UnicodeCharSeq):
            return_impl = True
            get_value = unicode_charseq_get_value
    if return_impl:
        max_i = s.count
        msg = 'index out of range [0, %s]' % (max_i - 1)

        def getitem_impl(s, i):
            if i < max_i and i >= 0:
                return get_value(s, i)
            raise IndexError(msg)
        return getitem_impl


@overload(len)
def charseq_len(s):
    return_impl = False
    if isinstance(s, types.CharSeq):
        return_impl = True
        get_code = charseq_get_code
    if isinstance(s, types.UnicodeCharSeq):
        return_impl = True
        get_code = unicode_charseq_get_code
    if return_impl:
        n = s.count
        if n == 0:
            def len_impl(s):
                return 0
        else:
            def len_impl(s):
                # return the index of the last non-null value (numpy
                # behavior)
                i = n
                code = 0
                while code == 0:
                    i = i - 1
                    code = get_code(s, i)
                return i + 1
        return len_impl


@overload(operator.eq)
def charseq_eq(a, b):
    return_impl = False
    if isinstance(a, types.CharSeq) and isinstance(b, types.CharSeq):
        return_impl = True
        left_code = charseq_get_code
        right_code = charseq_get_code
    if isinstance(a, types.UnicodeCharSeq) and isinstance(b, types.UnicodeCharSeq):
        return_impl = True
        left_code = unicode_charseq_get_code
        right_code = unicode_charseq_get_code
    if return_impl:
        def eq_impl(a, b):
            n = len(a)
            if n != len(b):
                return False
            for i in range(n):
                if left_code(a, i) != right_code(b, i):
                    return False
            return True
        return eq_impl


@overload(operator.ne)
def charseq_ne(a, b):
    if ((isinstance(a, types.CharSeq) and isinstance(b, types.CharSeq))
        or (isinstance(a, types.UnicodeCharSeq)
            and isinstance(b, types.UnicodeCharSeq))):
        def ne_impl(a, b):
            return not (a == b)
        return ne_impl
{'fraction_non_alphanumeric': 0.09244791666666667, 'fraction_numerical': 0.010602678571428572, 'mean_word_length': 2.5235910878112713, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5348240', 'n_tokens_mistral': 1798, 'n_tokens_neox': 1723, 'n_words': 481}