Repository: numba/numba

Issue:
ASCII flag for Numba Unicode
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

For more general "how do I do X?" type questions, please speak to us in real
time on https://gitter.im/numba/numba or post to the Numba mailing list
https://groups.google.com/a/continuum.io/forum/#!forum/numba-users.

-->

## Feature request

<!--

Please include details of the feature you would like to see, why you would
like to see it/the use case

-->
CPython's Unicode implementation includes an `ascii` flag ([here](https://github.com/python/cpython/blob/31e8d69bfe7cf5d4ffe0967cb225d2a8a229cc97/Include/cpython/unicodeobject.h#L211)), which indicates that all character values are less than 128 (note that 1 byte kind can be up to 256). This flag enables significant optimizations; for example, an ASCII string is a valid UTF-8 string and no conversion is required.


Discussion:

Resolving patch:
diff --git a/numba/pythonapi.py b/numba/pythonapi.py
--- a/numba/pythonapi.py
+++ b/numba/pythonapi.py
@@ -1141,20 +1141,24 @@ def string_as_string_size_and_kind(self, strobj):
         if PYVERSION >= (3, 3):
             p_length = cgutils.alloca_once(self.builder, self.py_ssize_t)
             p_kind = cgutils.alloca_once(self.builder, Type.int())
+            p_ascii = cgutils.alloca_once(self.builder, Type.int())
             p_hash = cgutils.alloca_once(self.builder, self.py_hash_t)
             fnty = Type.function(self.cstring, [self.pyobj,
                                                 self.py_ssize_t.as_pointer(),
                                                 Type.int().as_pointer(),
+                                                Type.int().as_pointer(),
                                                 self.py_hash_t.as_pointer()])
             fname = "numba_extract_unicode"
             fn = self._get_function(fnty, name=fname)
 
-            buffer = self.builder.call(fn, [strobj, p_length, p_kind, p_hash])
+            buffer = self.builder.call(
+                fn, [strobj, p_length, p_kind, p_ascii, p_hash])
             ok = self.builder.icmp_unsigned('!=',
                                             ir.Constant(buffer.type, None),
                                             buffer)
             return (ok, buffer, self.builder.load(p_length),
-                    self.builder.load(p_kind), self.builder.load(p_hash))
+                    self.builder.load(p_kind), self.builder.load(p_ascii),
+                    self.builder.load(p_hash))
         else:
             assert False, 'not supported on Python < 3.3'
 
diff --git a/numba/unicode.py b/numba/unicode.py
--- a/numba/unicode.py
+++ b/numba/unicode.py
@@ -42,6 +42,7 @@ def __init__(self, dmm, fe_type):
             ('data', types.voidptr),
             ('length', types.intp),
             ('kind', types.int32),
+            ('is_ascii', types.uint32),
             ('hash', _Py_hash_t),
             ('meminfo', types.MemInfoPointer(types.voidptr)),
             # A pointer to the owner python str/unicode object
@@ -53,6 +54,7 @@ def __init__(self, dmm, fe_type):
 make_attribute_wrapper(types.UnicodeType, 'data', '_data')
 make_attribute_wrapper(types.UnicodeType, 'length', '_length')
 make_attribute_wrapper(types.UnicodeType, 'kind', '_kind')
+make_attribute_wrapper(types.UnicodeType, 'is_ascii', '_is_ascii')
 make_attribute_wrapper(types.UnicodeType, 'hash', '_hash')
 
 
@@ -71,25 +73,27 @@ def compile_time_get_string_data(obj):
     the string data into the LLVM module.
     """
     from ctypes import (
-        CFUNCTYPE, c_void_p, c_int, c_ssize_t, c_ubyte, py_object,
+        CFUNCTYPE, c_void_p, c_int, c_uint, c_ssize_t, c_ubyte, py_object,
         POINTER, byref,
     )
 
     extract_unicode_fn = c_helpers['extract_unicode']
     proto = CFUNCTYPE(c_void_p, py_object, POINTER(c_ssize_t), POINTER(c_int),
-                      POINTER(c_ssize_t))
+                      POINTER(c_uint), POINTER(c_ssize_t))
     fn = proto(extract_unicode_fn)
     length = c_ssize_t()
     kind = c_int()
+    is_ascii = c_uint()
     hashv = c_ssize_t()
-    data = fn(obj, byref(length), byref(kind), byref(hashv))
+    data = fn(obj, byref(length), byref(kind), byref(is_ascii), byref(hashv))
     if data is None:
         raise ValueError("cannot extract unicode data from the given string")
     length = length.value
     kind = kind.value
+    is_ascii = is_ascii.value
     nbytes = (length + 1) * _kind_to_byte_width(kind)
     out = (c_ubyte * nbytes).from_address(data)
-    return bytes(out), length, kind, hashv.value
+    return bytes(out), length, kind, is_ascii, hashv.value
 
 
 def make_string_from_constant(context, builder, typ, literal_string):
@@ -97,7 +101,7 @@ def make_string_from_constant(context, builder, typ, literal_string):
     Get string data by `compile_time_get_string_data()` and return a
     unicode_type LLVM value
     """
-    databytes, length, kind, hashv = \
+    databytes, length, kind, is_ascii, hashv = \
         compile_time_get_string_data(literal_string)
     mod = builder.module
     gv = context.insert_const_bytes(mod, databytes)
@@ -105,6 +109,7 @@ def make_string_from_constant(context, builder, typ, literal_string):
     uni_str.data = gv
     uni_str.length = uni_str.length.type(length)
     uni_str.kind = uni_str.kind.type(kind)
+    uni_str.is_ascii = uni_str.is_ascii.type(is_ascii)
     uni_str.hash = uni_str.hash.type(hashv)
     return uni_str._getvalue()
 
@@ -131,11 +136,13 @@ def unbox_unicode_str(typ, obj, c):
     """
     Convert a unicode str object to a native unicode structure.
     """
-    ok, data, length, kind, hashv = c.pyapi.string_as_string_size_and_kind(obj)
+    ok, data, length, kind, is_ascii, hashv = \
+        c.pyapi.string_as_string_size_and_kind(obj)
     uni_str = cgutils.create_struct_proxy(typ)(c.context, c.builder)
     uni_str.data = data
     uni_str.length = length
     uni_str.kind = kind
+    uni_str.is_ascii = is_ascii
     uni_str.hash = hashv
     uni_str.meminfo = c.pyapi.nrt_meminfo_new_from_pyobject(
         data,  # the borrowed data pointer
@@ -195,13 +202,13 @@ def deref_uint32(typingctx, data, offset):
 
 
 @intrinsic
-def _malloc_string(typingctx, kind, char_bytes, length):
+def _malloc_string(typingctx, kind, char_bytes, length, is_ascii):
     """make empty string with data buffer of size alloc_bytes.
 
     Must set length and kind values for string after it is returned
     """
     def details(context, builder, signature, args):
-        [kind_val, char_bytes_val, length_val] = args
+        [kind_val, char_bytes_val, length_val, is_ascii_val] = args
 
         # fill the struct
         uni_str_ctor = cgutils.create_struct_proxy(types.unicode_type)
@@ -212,6 +219,7 @@ def details(context, builder, signature, args):
                                              Constant(length_val.type, 1)))
         uni_str.meminfo = context.nrt.meminfo_alloc(builder, nbytes_val)
         uni_str.kind = kind_val
+        uni_str.is_ascii = is_ascii_val
         uni_str.length = length_val
         # empty string has hash value -1 to indicate "need to compute hash"
         uni_str.hash = context.get_constant(_Py_hash_t, -1)
@@ -220,14 +228,14 @@ def details(context, builder, signature, args):
         uni_str.parent = cgutils.get_null_value(uni_str.parent.type)
         return uni_str._getvalue()
 
-    sig = types.unicode_type(types.int32, types.intp, types.intp)
+    sig = types.unicode_type(types.int32, types.intp, types.intp, types.uint32)
     return sig, details
 
 
 @njit
-def _empty_string(kind, length):
+def _empty_string(kind, length, is_ascii=0):
     char_width = _kind_to_byte_width(kind)
-    s = _malloc_string(kind, char_width, length)
+    s = _malloc_string(kind, char_width, length, is_ascii)
     _set_code_point(s, length, np.uint32(0))    # Write NULL character
     return s
 
@@ -316,6 +324,13 @@ def _pick_kind(kind1, kind2):
         raise AssertionError("Unexpected unicode representation in _pick_kind")
 
 
+@njit
+def _pick_ascii(is_ascii1, is_ascii2):
+    if is_ascii1 == 1 and is_ascii2 == 1:
+        return types.uint32(1)
+    return types.uint32(0)
+
+
 @njit
 def _kind_to_byte_width(kind):
     if kind == PY_UNICODE_1BYTE_KIND:
@@ -588,11 +603,13 @@ def join_list(sep, parts):
     sep_len = len(sep)
     length = (parts_len - 1) * sep_len
     kind = sep._kind
+    is_ascii = sep._is_ascii
     for p in parts:
         length += len(p)
         kind = _pick_kind(kind, p._kind)
+        is_ascii = _pick_ascii(is_ascii, p._is_ascii)
 
-    result = _empty_string(kind, length)
+    result = _empty_string(kind, length, is_ascii)
 
     # populate string
     part = parts[0]
@@ -726,6 +743,7 @@ def codegen(context, builder, sig, args):
             types.unicode_type)(context, builder)
         view_str.meminfo = in_str.meminfo
         view_str.kind = in_str.kind
+        view_str.is_ascii = in_str.is_ascii
         view_str.length = length
         # hash value -1 to indicate "need to compute hash"
         view_str.hash = context.get_constant(_Py_hash_t, -1)
@@ -754,7 +772,7 @@ def unicode_getitem(s, idx):
         if isinstance(idx, types.Integer):
             def getitem_char(s, idx):
                 idx = normalize_str_idx(idx, len(s))
-                ret = _empty_string(s._kind, 1)
+                ret = _empty_string(s._kind, 1, s._is_ascii)
                 _set_code_point(ret, 0, _get_code_point(s, idx))
                 return ret
             return getitem_char
@@ -766,7 +784,7 @@ def getitem_slice(s, idx):
                 if slice_idx.step == 1:
                     return _get_str_slice_view(s, slice_idx.start, span)
                 else:
-                    ret = _empty_string(s._kind, span)
+                    ret = _empty_string(s._kind, span, s._is_ascii)
                     cur = slice_idx.start
                     for i in range(span):
                         _set_code_point(ret, i, _get_code_point(s, cur))
@@ -782,7 +800,8 @@ def unicode_concat(a, b):
         def concat_impl(a, b):
             new_length = a._length + b._length
             new_kind = _pick_kind(a._kind, b._kind)
-            result = _empty_string(new_kind, new_length)
+            new_ascii = _pick_ascii(a._is_ascii, b._is_ascii)
+            result = _empty_string(new_kind, new_length, new_ascii)
             for i in range(len(a)):
                 _set_code_point(result, i, _get_code_point(a, i))
             for j in range(len(b)):
@@ -800,7 +819,7 @@ def _repeat_impl(str_arg, mult_arg):
     else:
         new_length = str_arg._length * mult_arg
         new_kind = str_arg._kind
-        result = _empty_string(new_kind, new_length)
+        result = _empty_string(new_kind, new_length, str_arg._is_ascii)
         # make initial copy into result
         len_a = len(str_arg)
         _strncpy(result, 0, str_arg, 0, len_a)


Test patch:
diff --git a/numba/tests/test_unicode.py b/numba/tests/test_unicode.py
--- a/numba/tests/test_unicode.py
+++ b/numba/tests/test_unicode.py
@@ -16,6 +16,9 @@
 _py34_or_later = sys.version_info[:2] >= (3, 4)
 
 
+isascii = lambda s: all(ord(c) < 128 for c in s)
+
+
 def literal_usecase():
     return '大处着眼，小处着手。'
 
@@ -723,6 +726,53 @@ def f():
 
         self.assertEqual(f.py_func(), f())
 
+    def test_ascii_flag_unbox(self):
+        @njit
+        def f(s):
+            return s._is_ascii
+
+        for s in UNICODE_EXAMPLES:
+            self.assertEqual(f(s), isascii(s))
+
+    def test_ascii_flag_join(self):
+        @njit
+        def f():
+            s1 = 'abc'
+            s2 = '123'
+            s3 = '🐍⚡'
+            s4 = '大处着眼，小处着手。'
+            return (",".join([s1, s2])._is_ascii,
+                    "🐍⚡".join([s1, s2])._is_ascii,
+                    ",".join([s1, s3])._is_ascii,
+                    ",".join([s3, s4])._is_ascii)
+
+        self.assertEqual(f(), (1, 0, 0, 0))
+
+    def test_ascii_flag_getitem(self):
+        @njit
+        def f():
+            s1 = 'abc123'
+            s2 = '🐍⚡🐍⚡🐍⚡'
+            return (s1[0]._is_ascii, s1[2:]._is_ascii, s2[0]._is_ascii,
+                    s2[2:]._is_ascii)
+
+        self.assertEqual(f(), (1, 1, 0, 0))
+
+    def test_ascii_flag_add_mul(self):
+        @njit
+        def f():
+            s1 = 'abc'
+            s2 = '123'
+            s3 = '🐍⚡'
+            s4 = '大处着眼，小处着手。'
+            return ((s1 + s2)._is_ascii,
+                    (s1 + s3)._is_ascii,
+                    (s3 + s4)._is_ascii,
+                    (s1 * 2)._is_ascii,
+                    (s3 * 2)._is_ascii)
+
+        self.assertEqual(f(), (1, 0, 0, 1, 0))
+
 
 @unittest.skipUnless(_py34_or_later,
                      'unicode support requires Python 3.4 or later')