From 84d413ebc13e5c595f8c1b23d2b9eae7685cc349 Mon Sep 17 00:00:00 2001
From: Jingchen Ye <97littleleaf11@gmail.com>
Date: Wed, 11 Aug 2021 02:35:00 +0800
Subject: [PATCH 1/4] Support bytes slicing

---
 mypyc/irbuild/expression.py        |  3 +-
 mypyc/lib-rt/bytes_ops.c           | 26 +++++++++++++
 mypyc/primitives/bytes_ops.py      | 11 +++++-
 mypyc/test-data/fixtures/ir.py     |  3 ++
 mypyc/test-data/irbuild-bytes.test | 12 ++++++
 mypyc/test-data/run-bytes.test     | 60 ++++++++++++++++++++++++++++++
 6 files changed, 112 insertions(+), 3 deletions(-)

diff --git a/mypyc/irbuild/expression.py b/mypyc/irbuild/expression.py
index 9794cd058d480..f71fa9e010cfc 100644
--- a/mypyc/irbuild/expression.py
+++ b/mypyc/irbuild/expression.py
@@ -27,6 +27,7 @@
 from mypyc.irbuild.format_str_tokenizer import (
     tokenizer_printf_style, join_formatted_strings, convert_expr
 )
+from mypyc.primitives.bytes_ops import bytes_slice_op
 from mypyc.primitives.registry import CFunctionDescription, builtin_names, binary_ops
 from mypyc.primitives.generic_ops import iter_op
 from mypyc.primitives.misc_ops import new_slice_op, ellipsis_op, type_op, get_module_dict_op
@@ -441,7 +442,7 @@ def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Optio
             # Replace missing end index with the largest short integer
             # (a sequence can't be longer).
             end = builder.load_int(MAX_SHORT_INT)
-        candidates = [list_slice_op, tuple_slice_op, str_slice_op]
+        candidates = [list_slice_op, tuple_slice_op, str_slice_op, bytes_slice_op]
         return builder.builder.matching_call_c(candidates, [base, begin, end], index.line)
 
     return None
diff --git a/mypyc/lib-rt/bytes_ops.c b/mypyc/lib-rt/bytes_ops.c
index 64db425836ed9..f1d05772c3550 100644
--- a/mypyc/lib-rt/bytes_ops.c
+++ b/mypyc/lib-rt/bytes_ops.c
@@ -23,6 +23,32 @@ PyObject *CPyBytes_Concat(PyObject *a, PyObject *b) {
     }
 }
 
+#define CLAMP(a, b, c) (a < b ? b : (a >= c ? c : a))
+
+PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
+    if ((PyBytes_Check(obj) || PyByteArray_Check(obj))
+            && CPyTagged_CheckShort(start) && CPyTagged_CheckShort(end)) {
+        Py_ssize_t startn = CPyTagged_ShortAsSsize_t(start);
+        Py_ssize_t endn = CPyTagged_ShortAsSsize_t(end);
+        Py_ssize_t len = ((PyVarObject *)obj)->ob_size;
+        if (startn < 0) {
+            startn += len;
+        }
+        if (endn < 0) {
+            endn += len;
+        }
+        startn = CLAMP(startn, 0, len);
+        endn = CLAMP(endn, 0, len);
+        Py_ssize_t slice_len = endn - startn;
+        if (PyBytes_Check(obj)) {
+            return PyBytes_FromStringAndSize(PyBytes_AS_STRING(obj) + startn, slice_len);
+        } else {
+            return PyByteArray_FromStringAndSize(PyByteArray_AS_STRING(obj) + startn, slice_len);
+        }
+    }
+    return CPyObject_GetSlice(obj, start, end);
+}
+
 // Like _PyBytes_Join but fallback to dynamic call if 'sep' is not bytes
 // (mostly commonly, for bytearrays)
 PyObject *CPyBytes_Join(PyObject *sep, PyObject *iter) {
diff --git a/mypyc/primitives/bytes_ops.py b/mypyc/primitives/bytes_ops.py
index a95a7da3fa190..f8811de5987ab 100644
--- a/mypyc/primitives/bytes_ops.py
+++ b/mypyc/primitives/bytes_ops.py
@@ -3,10 +3,10 @@
 from mypyc.ir.ops import ERR_MAGIC
 from mypyc.ir.rtypes import (
     object_rprimitive, bytes_rprimitive, list_rprimitive, dict_rprimitive,
-    str_rprimitive, RUnion
+    str_rprimitive, RUnion, int_rprimitive
 )
 from mypyc.primitives.registry import (
-    load_address_op, function_op, method_op, binary_op
+    load_address_op, function_op, method_op, binary_op, custom_op
 )
 
 # Get the 'bytes' type object.
@@ -41,6 +41,13 @@
     error_kind=ERR_MAGIC,
     steals=[True, False])
 
+# bytes[begin:end]
+bytes_slice_op = custom_op(
+    arg_types=[bytes_rprimitive, int_rprimitive, int_rprimitive],
+    return_type=bytes_rprimitive,
+    c_function_name='CPyBytes_GetSlice',
+    error_kind=ERR_MAGIC)
+
 # bytes.join(obj)
 method_op(
     name='join',
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index f64db578c868b..5e2f8293e6591 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/mypyc/test-data/fixtures/ir.py
@@ -104,7 +104,10 @@ def __init__(self, x: object) -> None: ...
     def __add__(self, x: bytes) -> bytes: ...
     def __eq__(self, x: object) -> bool: ...
     def __ne__(self, x: object) -> bool: ...
+    @overload
     def __getitem__(self, i: int) -> int: ...
+    @overload
+    def __getitem__(self, i: slice) -> bytes: ...
     def join(self, x: Iterable[object]) -> bytes: ...
     def decode(self, x: str, y: str=...) -> str: ...
 
diff --git a/mypyc/test-data/irbuild-bytes.test b/mypyc/test-data/irbuild-bytes.test
index d86dc3227b4a3..e3fb7726febea 100644
--- a/mypyc/test-data/irbuild-bytes.test
+++ b/mypyc/test-data/irbuild-bytes.test
@@ -62,6 +62,18 @@ L0:
     c = r6
     return 1
 
+[case testBytesSlicing]
+def f(a: bytes, start: int, end: int) -> bytes:
+    return a[start:end]
+[out]
+def f(a, start, end):
+    a :: bytes
+    start, end :: int
+    r0 :: bytes
+L0:
+    r0 = CPyBytes_GetSlice(a, start, end)
+    return r0
+
 [case testBytesConcat]
 def f1(a: bytes, b: bytes) -> bytes:
     return a + b
diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test
index 01fd6c45c3f9d..9be27c047c84c 100644
--- a/mypyc/test-data/run-bytes.test
+++ b/mypyc/test-data/run-bytes.test
@@ -98,6 +98,36 @@ def test_len() -> None:
     assert len(b) == 3
     assert len(bytes()) == 0
 
+[case testBytesSlicing]
+def test_bytes_slicing() -> None:
+    b = b'abcdefg'
+    zero = int()
+    ten = 10 + zero
+    two = 2 + zero
+    five = 5 + zero
+    seven = 7 + zero
+    assert b[:ten] == b'abcdefg'
+    assert b[0:seven] == b'abcdefg'
+    assert b[two:five] == b'cde'
+    assert b[two:two] == b''
+    assert b[-two:-two] == b''
+    assert b[-ten:(-ten+1)] == b''
+    assert b[:-two] == b'abcde'
+    assert b[:two] == b'ab'
+    assert b[:] == b'abcdefg'
+    assert b[-two:] == b'fg'
+    assert b[zero:] == b'abcdefg'
+    assert b[:zero] == b''
+    assert b[-ten:] == b'abcdefg'
+    assert b[-ten:ten] == b'abcdefg'
+    big_int: int = 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000
+    assert b[1:big_int] == b'bcdefg'
+    assert b[big_int:] == b''
+    assert b[-big_int:-1] == b'abcdef'
+    assert type(b[-ten:]) == bytes
+    assert type(b[-big_int:-1]) == bytes
+    assert type(b[:]) == bytes
+
 [case testBytearrayBasics]
 from typing import Any
 
@@ -126,6 +156,36 @@ def test_bytearray_passed_into_bytes() -> None:
     brr1: Any = bytearray()
     assert f(brr1)
 
+[case testBytearraySlicing]
+def test_bytearray_slicing() -> None:
+    b: bytes = bytearray(b'abcdefg')
+    zero = int()
+    ten = 10 + zero
+    two = 2 + zero
+    five = 5 + zero
+    seven = 7 + zero
+    assert b[:ten] == b'abcdefg'
+    assert b[0:seven] == b'abcdefg'
+    assert b[two:five] == b'cde'
+    assert b[two:two] == b''
+    assert b[-two:-two] == b''
+    assert b[-ten:(-ten+1)] == b''
+    assert b[:-two] == b'abcde'
+    assert b[:two] == b'ab'
+    assert b[:] == b'abcdefg'
+    assert b[-two:] == b'fg'
+    assert b[zero:] == b'abcdefg'
+    assert b[:zero] == b''
+    assert b[-ten:] == b'abcdefg'
+    assert b[-ten:ten] == b'abcdefg'
+    big_int: int = 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000
+    assert b[1:big_int] == b'bcdefg'
+    assert b[big_int:] == b''
+    assert b[-big_int:-1] == b'abcdef'
+    assert type(b[-ten:]) == bytearray
+    assert type(b[-big_int:-1]) == bytearray
+    assert type(b[:]) == bytearray
+
 [case testBytesJoin]
 from typing import Any
 from testutil import assertRaises

From f5bbc5d7e0bab1947809b1d5c222116beaa0c479 Mon Sep 17 00:00:00 2001
From: Jingchen Ye <97littleleaf11@gmail.com>
Date: Wed, 11 Aug 2021 03:48:04 +0800
Subject: [PATCH 2/4] Add missing headers

---
 mypyc/lib-rt/CPy.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index 18a1ce76f26a5..591a1263dc261 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/mypyc/lib-rt/CPy.h
@@ -400,6 +400,7 @@ Py_ssize_t CPyStr_Size_size_t(PyObject *str);
 // Bytes operations
 
 
+PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 PyObject *CPyBytes_Concat(PyObject *a, PyObject *b);
 PyObject *CPyBytes_Join(PyObject *sep, PyObject *iter);
 

From 7e24ba3b9727409dc00a2a42935057568b66e890 Mon Sep 17 00:00:00 2001
From: Jingchen Ye <97littleleaf11@gmail.com>
Date: Wed, 11 Aug 2021 21:52:22 +0800
Subject: [PATCH 3/4] Add tests

---
 mypyc/lib-rt/bytes_ops.c       |  8 +++++---
 mypyc/test-data/run-bytes.test | 25 +++++++++++++++----------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/mypyc/lib-rt/bytes_ops.c b/mypyc/lib-rt/bytes_ops.c
index 0a1416f76d795..5d855805945bf 100644
--- a/mypyc/lib-rt/bytes_ops.c
+++ b/mypyc/lib-rt/bytes_ops.c
@@ -42,7 +42,9 @@ PyObject *CPyBytes_Concat(PyObject *a, PyObject *b) {
     }
 }
 
-#define CLAMP(a, b, c) (a < b ? b : (a >= c ? c : a))
+static inline Py_ssize_t Clamp(Py_ssize_t a, Py_ssize_t b, Py_ssize_t c) {
+    return a < b ? b : (a >= c ? c : a);
+}
 
 PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
     if ((PyBytes_Check(obj) || PyByteArray_Check(obj))
@@ -56,8 +58,8 @@ PyObject *CPyBytes_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
         if (endn < 0) {
             endn += len;
         }
-        startn = CLAMP(startn, 0, len);
-        endn = CLAMP(endn, 0, len);
+        startn = Clamp(startn, 0, len);
+        endn = Clamp(endn, 0, len);
         Py_ssize_t slice_len = endn - startn;
         if (PyBytes_Check(obj)) {
             return PyBytes_FromStringAndSize(PyBytes_AS_STRING(obj) + startn, slice_len);
diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test
index e86763e21f4cd..23ce56c534d77 100644
--- a/mypyc/test-data/run-bytes.test
+++ b/mypyc/test-data/run-bytes.test
@@ -119,6 +119,7 @@ def test_bytes_slicing() -> None:
     seven = 7 + zero
     assert b[:ten] == b'abcdefg'
     assert b[0:seven] == b'abcdefg'
+    assert b[0:(len(b)+1)] == b'abcdefg'
     assert b[two:five] == b'cde'
     assert b[two:two] == b''
     assert b[-two:-two] == b''
@@ -131,12 +132,14 @@ def test_bytes_slicing() -> None:
     assert b[:zero] == b''
     assert b[-ten:] == b'abcdefg'
     assert b[-ten:ten] == b'abcdefg'
-    big_int: int = 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000
-    assert b[1:big_int] == b'bcdefg'
-    assert b[big_int:] == b''
-    assert b[-big_int:-1] == b'abcdef'
+    big_ints = [1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000, 2**24, 2**63]
+    for big_int in big_ints:
+        assert b[1:big_int] == b'bcdefg'
+        assert b[big_int:] == b''
+        assert b[-big_int:-1] == b'abcdef'
+        assert b[-big_int:big_int] == b'abcdefg'
+        assert type(b[-big_int:-1]) == bytes
     assert type(b[-ten:]) == bytes
-    assert type(b[-big_int:-1]) == bytes
     assert type(b[:]) == bytes
 
 [case testBytearrayBasics]
@@ -190,12 +193,14 @@ def test_bytearray_slicing() -> None:
     assert b[:zero] == b''
     assert b[-ten:] == b'abcdefg'
     assert b[-ten:ten] == b'abcdefg'
-    big_int: int = 1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000
-    assert b[1:big_int] == b'bcdefg'
-    assert b[big_int:] == b''
-    assert b[-big_int:-1] == b'abcdef'
+    big_ints = [1000 * 1000 * 1000 * 1000 * 1000 * 1000 * 1000, 2**24, 2**63]
+    for big_int in big_ints:
+        assert b[1:big_int] == b'bcdefg'
+        assert b[big_int:] == b''
+        assert b[-big_int:-1] == b'abcdef'
+        assert b[-big_int:big_int] == b'abcdefg'
+        assert type(b[-big_int:-1]) == bytearray
     assert type(b[-ten:]) == bytearray
-    assert type(b[-big_int:-1]) == bytearray
     assert type(b[:]) == bytearray
 
 [case testBytearrayIndexing]

From 77f84120b4f492332a8415b10853a4091079a3de Mon Sep 17 00:00:00 2001
From: Jingchen Ye <97littleleaf11@gmail.com>
Date: Wed, 11 Aug 2021 22:27:09 +0800
Subject: [PATCH 4/4] Fix

---
 mypyc/test-data/run-bytes.test | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mypyc/test-data/run-bytes.test b/mypyc/test-data/run-bytes.test
index 23ce56c534d77..be8b62a8ec376 100644
--- a/mypyc/test-data/run-bytes.test
+++ b/mypyc/test-data/run-bytes.test
@@ -144,7 +144,6 @@ def test_bytes_slicing() -> None:
 
 [case testBytearrayBasics]
 from typing import Any
-from testutil import assertRaises
 
 def test_basics() -> None:
     brr1: bytes = bytearray(3)
@@ -204,6 +203,8 @@ def test_bytearray_slicing() -> None:
     assert type(b[:]) == bytearray
 
 [case testBytearrayIndexing]
+from testutil import assertRaises
+
 def test_bytearray_indexing() -> None:
     b: bytes = bytearray(b'\xae\x80\xfe\x15')
     assert b[0] == 174