Repository: numba/numba

Issue:
Unicode characters as dict keys don't work reliably due to character width ambiguity
The following code breaks with a KeyError:

```
@njit
def test():
    s = 'a\u1234'
    d = {s[0] : 1}
    print(d['a'])

test()
```

The presence of a wider than 1 byte character in the string "s" triggers the problem. Quote from the documentation:

> String slices also use the same character width as the original string, even if the slice could be represented with a narrower character width. (These details are invisible to the user, of course.)

I suspect that the different character width results in a different hash, which breaks dictionary lookups using such strings as keys. I found this workaround for my application:

```
@njit
def fixuni(c):
    u = c.upper()
    return u if u == c else c.lower()

@njit
def test():
    s = 'a\u1234'
    d = {fixuni(s[0]) : 1}
    print(d['a'])

test()
```



Discussion:
@fxkuehl thanks for submitting this, your efforts to help to improve Numba are appreciated! RE formatting in GitHub issues, for future reference triple backticks are needed around a code block (with optional language syntax specified), for example: ` ```python <your code> ``` `. GitHub also has a basic guide to its markdown syntax available [here](https://help.github.com/articles/basic-writing-and-formatting-syntax). I have updated the issue for you.
@fxkuehl thanks for submitting this, I can confirm I am able to reproduce this.
I wonder how CPython handles this.  We copied both the string storage logic and the hashing function from Python 3...
Whilst this shows one bug, there's probably two stemming from the same thing. The OP shows a `getitem` with `int`, but a slice would do similar.

------------
Bug 1: string getitem(int) returns a kind the size of the original string.

This is the CPython getitem for an int index:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L11985-L12007
note it reads the value based on the kind of the string being index, but returns via `unicode_char`:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L2160-L2184
which is based on the minimal kind representation of the character presented.

Numba just reads in the original string size:
https://github.com/numba/numba/blob/b9efc8323c4b5915be55b7a375c4b0b8fdd84fef/numba/cpython/unicode.py#L1639-L1643

------------
Bug 2: string getitem(slice) returns a kind the size of the original string.

For slices, this is the abstract "get a slice from a sequence", note it calls the mapping method `mp_subscript`:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/abstract.c#L1727-L1748

which for unicode objects is defined here:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L14512-L14516

and this is the actual implementation:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L14440-L14510

the non-ASCII patch checks the slice for the max char value:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L14480-L14491

allocates a new buffer based on this:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L14494

and then queries the exact kind and length of the new buffer for doing the write:
https://github.com/python/cpython/blob/c3dd7e45cc5d36bbe2295c2840faabb5c75d83e4/Objects/unicodeobject.c#L14497-L14503

basically, the slice ends up being a string of the minimal kind required to hold the maximum char value.

For a stride 1 slice Numba returns a view, so it'll have the byte width of the original even if the slice doesn't need it. For non-unit stride the allocation will be for a string of the kind of the string being sliced.
https://github.com/numba/numba/blob/b9efc8323c4b5915be55b7a375c4b0b8fdd84fef/numba/cpython/unicode.py#L1646-L1658

Resolving patch:
diff --git a/numba/cpython/unicode.py b/numba/cpython/unicode.py
--- a/numba/cpython/unicode.py
+++ b/numba/cpython/unicode.py
@@ -1644,8 +1644,11 @@ 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, s._is_ascii)
-                _set_code_point(ret, 0, _get_code_point(s, idx))
+                cp = _get_code_point(s, idx)
+                kind = _codepoint_to_kind(cp)
+                is_ascii = _codepoint_is_ascii(cp)
+                ret = _empty_string(kind, 1, is_ascii)
+                _set_code_point(ret, 0, cp)
                 return ret
             return getitem_char
         elif isinstance(idx, types.SliceType):
@@ -1653,15 +1656,38 @@ def getitem_slice(s, idx):
                 slice_idx = _normalize_slice(idx, len(s))
                 span = _slice_span(slice_idx)
 
-                if slice_idx.step == 1:
+                cp = _get_code_point(s, slice_idx.start)
+                kind = _codepoint_to_kind(cp)
+                is_ascii = _codepoint_is_ascii(cp)
+
+                # Check slice to see if it's homogeneous in kind
+                for i in range(slice_idx.start + slice_idx.step,
+                               slice_idx.stop, slice_idx.step):
+                    cp = _get_code_point(s, i)
+                    is_ascii |= _codepoint_is_ascii(cp)
+                    new_kind = _codepoint_to_kind(cp)
+                    if kind != new_kind:
+                        kind = _pick_kind(kind, new_kind)
+                    # TODO: it might be possible to break here if the kind
+                    # is PY_UNICODE_4BYTE_KIND but there are potentially
+                    # strings coming from other internal functions that are
+                    # this wide and also actually ASCII (i.e. kind is larger
+                    # than actually required for storing the code point), so
+                    # it's necessary to continue.
+
+                if slice_idx.step == 1 and kind == s._kind:
+                    # Can return a view, the slice has the same kind as the
+                    # string itself and it's a stride slice 1.
                     return _get_str_slice_view(s, slice_idx.start, span)
                 else:
-                    ret = _empty_string(s._kind, span, s._is_ascii)
+                    # It's heterogeneous in kind OR stride != 1
+                    ret = _empty_string(kind, span, is_ascii)
                     cur = slice_idx.start
                     for i in range(span):
                         _set_code_point(ret, i, _get_code_point(s, cur))
                         cur += slice_idx.step
                     return ret
+
             return getitem_slice
 
 


Test patch:
diff --git a/numba/tests/test_dictobject.py b/numba/tests/test_dictobject.py
--- a/numba/tests/test_dictobject.py
+++ b/numba/tests/test_dictobject.py
@@ -912,6 +912,24 @@ def bar():
 
         self.assertEqual(foo(), [1, 2])
 
+    def test_024_unicode_getitem_keys(self):
+        # See issue #6135
+        @njit
+        def foo():
+            s = 'a\u1234'
+            d = {s[0] : 1}
+            return d['a']
+
+        self.assertEqual(foo(), foo.py_func())
+
+        @njit
+        def foo():
+            s = 'abc\u1234'
+            d = {s[:1] : 1}
+            return d['a']
+
+        self.assertEqual(foo(), foo.py_func())
+
 
 class TestDictTypeCasting(TestCase):
     def check_good(self, fromty, toty):
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
@@ -42,6 +42,10 @@ def getitem_usecase(x, i):
     return x[i]
 
 
+def getitem_check_kind_usecase(x, i):
+    return hash(x[i])
+
+
 def zfill_usecase(x, y):
     return x.zfill(y)
 
@@ -1173,7 +1177,20 @@ def test_getitem(self):
         cfunc = njit(pyfunc)
 
         for s in UNICODE_EXAMPLES:
-            for i in range(-len(s)):
+            for i in range(-len(s), len(s)):
+                self.assertEqual(pyfunc(s, i),
+                                 cfunc(s, i),
+                                 "'%s'[%d]?" % (s, i))
+
+    def test_getitem_scalar_kind(self):
+        # See issue #6135, make sure that getitem returns a char of the minimal
+        # kind required to represent the "got" item, this is done via the use
+        # of `hash` in the test function as it is sensitive to kind.
+        pyfunc = getitem_check_kind_usecase
+        cfunc = njit(pyfunc)
+        samples = ['a\u1234', '¡着']
+        for s in samples:
+            for i in range(-len(s), len(s)):
                 self.assertEqual(pyfunc(s, i),
                                  cfunc(s, i),
                                  "'%s'[%d]?" % (s, i))
@@ -1217,6 +1234,20 @@ def test_slice2_error(self):
                                      cfunc(s, sl),
                                      "'%s'[%d:%d]?" % (s, i, j))
 
+    def test_getitem_slice2_kind(self):
+        # See issue #6135. Also see note in test_getitem_scalar_kind regarding
+        # testing.
+        pyfunc = getitem_check_kind_usecase
+        cfunc = njit(pyfunc)
+        samples = ['abc\u1234\u1234', '¡¡¡着着着']
+        for s in samples:
+            for i in [-2, -1, 0, 1, 2, len(s), len(s) + 1]:
+                for j in [-2, -1, 0, 1, 2, len(s), len(s) + 1]:
+                    sl = slice(i, j)
+                    self.assertEqual(pyfunc(s, sl),
+                                     cfunc(s, sl),
+                                     "'%s'[%d:%d]?" % (s, i, j))
+
     def test_slice3(self):
         pyfunc = getitem_usecase
         cfunc = njit(pyfunc)
@@ -1230,6 +1261,27 @@ def test_slice3(self):
                                          cfunc(s, sl),
                                          "'%s'[%d:%d:%d]?" % (s, i, j, k))
 
+    def test_getitem_slice3_kind(self):
+        # See issue #6135. Also see note in test_getitem_scalar_kind regarding
+        # testing.
+        pyfunc = getitem_check_kind_usecase
+        cfunc = njit(pyfunc)
+        samples = ['abc\u1234\u1234',
+                   'a\u1234b\u1234c'
+                   '¡¡¡着着着',
+                   '¡着¡着¡着',
+                   '着a着b着c',
+                   '¡着a¡着b¡着c',
+                   '¡着a着¡c',]
+        for s in samples:
+            for i in range(-len(s), len(s)):
+                for j in range(-len(s), len(s)):
+                    for k in [-2, -1, 1, 2]:
+                        sl = slice(i, j, k)
+                        self.assertEqual(pyfunc(s, sl),
+                                         cfunc(s, sl),
+                                         "'%s'[%d:%d:%d]?" % (s, i, j, k))
+
     def test_slice3_error(self):
         pyfunc = getitem_usecase
         cfunc = njit(pyfunc)