[mypyc] Add primitives for list, str and tuple slicing (#9283)

This speeds up some microbenchmarks from 40% (list) to 100% (str) 
when I run them on Ubuntu 20.04.

Non-default strides aren't optimized, since they are fairly rare.

`a[::-1]` for lists might be worth special casing in the future. Also, once
we have primitives for `bytes`, it should also be special cased.

Fixes mypyc/mypyc#725.

Author: JukkaL

mypyc/common.py

@@ -1,5 +1,6 @@
 import sys
 from typing import Dict, Any
+import sys
 
 from typing_extensions import Final
 
@@ -22,7 +23,11 @@
 
 # Max short int we accept as a literal is based on 32-bit platforms,
 # so that we can just always emit the same code.
-MAX_LITERAL_SHORT_INT = (1 << 30) - 1  # type: Final
+
+# Maximum value for a short tagged integer.
+#
+# Note: Assume that the compiled code uses the same bit width as mypyc.
+MAX_LITERAL_SHORT_INT = sys.maxsize >> 1  # type: Final
 
 TOP_LEVEL_NAME = '__top_level__'  # type: Final # Special function representing module top level
 

mypyc/doc/list_operations.rst

@@ -29,6 +29,10 @@ Get item by integer index:
 
 * ``lst[n]``
 
+Slicing:
+
+* ``lst[n:m]``, ``lst[n:]``, ``lst[:m]``, ``lst[:]``
+
 Repeat list ``n`` times:
 
 * ``lst * n``, ``n * lst``

mypyc/doc/str_operations.rst

@@ -24,6 +24,10 @@ Indexing:
 
 * ``s[n]`` (integer index)
 
+Slicing:
+
+* ``s[n:m]``, ``s[n:]``, ``s[:m]``
+
 Comparisons:
 
 * ``s1 == s2``, ``s1 != s2``

mypyc/doc/tuple_operations.rst

@@ -20,6 +20,7 @@ Operators
 ---------
 
 * ``tup[n]`` (integer index)
+* ``tup[n:m]``, ``tup[n:]``, ``tup[:m]`` (slicing)
 
 Statements
 ----------

mypyc/irbuild/builder.py

@@ -185,6 +185,9 @@ def py_get_attr(self, obj: Value, attr: str, line: int) -> Value:
     def load_static_unicode(self, value: str) -> Value:
         return self.builder.load_static_unicode(value)
 
+    def load_static_int(self, value: int) -> Value:
+        return self.builder.load_static_int(value)
+
     def primitive_op(self, desc: OpDescription, args: List[Value], line: int) -> Value:
         return self.builder.primitive_op(desc, args, line)
 

mypyc/irbuild/expression.py

@@ -15,18 +15,22 @@
 )
 from mypy.types import TupleType, get_proper_type
 
+from mypyc.common import MAX_LITERAL_SHORT_INT
 from mypyc.ir.ops import (
     Value, TupleGet, TupleSet, BasicBlock, Assign, LoadAddress
 )
-from mypyc.ir.rtypes import RTuple, object_rprimitive, is_none_rprimitive, is_int_rprimitive
+from mypyc.ir.rtypes import (
+    RTuple, object_rprimitive, is_none_rprimitive, int_rprimitive, is_int_rprimitive
+)
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
 from mypyc.primitives.registry import CFunctionDescription, builtin_names
 from mypyc.primitives.generic_ops import iter_op
 from mypyc.primitives.misc_ops import new_slice_op, ellipsis_op, type_op
-from mypyc.primitives.list_ops import list_append_op, list_extend_op
-from mypyc.primitives.tuple_ops import list_tuple_op
+from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
+from mypyc.primitives.tuple_ops import list_tuple_op, tuple_slice_op
 from mypyc.primitives.dict_ops import dict_new_op, dict_set_item_op
 from mypyc.primitives.set_ops import new_set_op, set_add_op, set_update_op
+from mypyc.primitives.str_ops import str_slice_op
 from mypyc.primitives.int_ops import int_comparison_op_mapping
 from mypyc.irbuild.specialize import specializers
 from mypyc.irbuild.builder import IRBuilder
@@ -323,15 +327,59 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
 
 def transform_index_expr(builder: IRBuilder, expr: IndexExpr) -> Value:
     base = builder.accept(expr.base)
+    index = expr.index
+
+    if isinstance(base.type, RTuple) and isinstance(index, IntExpr):
+        return builder.add(TupleGet(base, index.value, expr.line))
 
-    if isinstance(base.type, RTuple) and isinstance(expr.index, IntExpr):
-        return builder.add(TupleGet(base, expr.index.value, expr.line))
+    if isinstance(index, SliceExpr):
+        value = try_gen_slice_op(builder, base, index)
+        if value:
+            return value
 
     index_reg = builder.accept(expr.index)
     return builder.gen_method_call(
         base, '__getitem__', [index_reg], builder.node_type(expr), expr.line)
 
 
+def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Optional[Value]:
+    """Generate specialized slice op for some index expressions.
+
+    Return None if a specialized op isn't available.
+
+    This supports obj[x:y], obj[:x], and obj[x:] for a few types.
+    """
+    if index.stride:
+        # We can only handle the default stride of 1.
+        return None
+
+    if index.begin_index:
+        begin_type = builder.node_type(index.begin_index)
+    else:
+        begin_type = int_rprimitive
+    if index.end_index:
+        end_type = builder.node_type(index.end_index)
+    else:
+        end_type = int_rprimitive
+
+    # Both begin and end index must be int (or missing).
+    if is_int_rprimitive(begin_type) and is_int_rprimitive(end_type):
+        if index.begin_index:
+            begin = builder.accept(index.begin_index)
+        else:
+            begin = builder.load_static_int(0)
+        if index.end_index:
+            end = builder.accept(index.end_index)
+        else:
+            # Replace missing end index with the largest short integer
+            # (a sequence can't be longer).
+            end = builder.load_static_int(MAX_LITERAL_SHORT_INT)
+        candidates = [list_slice_op, tuple_slice_op, str_slice_op]
+        return builder.builder.matching_call_c(candidates, [base, begin, end], index.line)
+
+    return None
+
+
 def transform_conditional_expr(builder: IRBuilder, expr: ConditionalExpr) -> Value:
     if_body, else_body, next = BasicBlock(), BasicBlock(), BasicBlock()
 

mypyc/lib-rt/CPy.h

@@ -303,6 +303,7 @@ CPyTagged CPyObject_Hash(PyObject *o);
 PyObject *CPyObject_GetAttr3(PyObject *v, PyObject *name, PyObject *defl);
 PyObject *CPyIter_Next(PyObject *iter);
 PyObject *CPyNumber_Power(PyObject *base, PyObject *index);
+PyObject *CPyObject_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
 
 // List operations
@@ -318,6 +319,7 @@ CPyTagged CPyList_Count(PyObject *obj, PyObject *value);
 PyObject *CPyList_Extend(PyObject *o1, PyObject *o2);
 PyObject *CPySequence_Multiply(PyObject *seq, CPyTagged t_size);
 PyObject *CPySequence_RMultiply(CPyTagged t_size, PyObject *seq);
+PyObject *CPyList_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
 
 // Dict operations
@@ -367,6 +369,7 @@ static inline char CPyDict_CheckSize(PyObject *dict, CPyTagged size) {
 PyObject *CPyStr_GetItem(PyObject *str, CPyTagged index);
 PyObject *CPyStr_Split(PyObject *str, PyObject *sep, CPyTagged max_split);
 PyObject *CPyStr_Append(PyObject *o1, PyObject *o2);
+PyObject *CPyStr_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
 
 // Set operations
@@ -379,6 +382,7 @@ bool CPySet_Remove(PyObject *set, PyObject *key);
 
 
 PyObject *CPySequenceTuple_GetItem(PyObject *tuple, CPyTagged index);
+PyObject *CPySequenceTuple_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
 
 // Exception operations

mypyc/lib-rt/generic_ops.c

@@ -40,3 +40,20 @@ PyObject *CPyNumber_Power(PyObject *base, PyObject *index)
 {
     return PyNumber_Power(base, index, Py_None);
 }
+
+PyObject *CPyObject_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
+    PyObject *start_obj = CPyTagged_AsObject(start);
+    PyObject *end_obj = CPyTagged_AsObject(end);
+    if (unlikely(start_obj == NULL || end_obj == NULL)) {
+        return NULL;
+    }
+    PyObject *slice = PySlice_New(start_obj, end_obj, NULL);
+    Py_DECREF(start_obj);
+    Py_DECREF(end_obj);
+    if (unlikely(slice == NULL)) {
+        return NULL;
+    }
+    PyObject *result = PyObject_GetItem(obj, slice);
+    Py_DECREF(slice);
+    return result;
+}

mypyc/lib-rt/list_ops.c

@@ -123,3 +123,19 @@ PyObject *CPySequence_Multiply(PyObject *seq, CPyTagged t_size) {
 PyObject *CPySequence_RMultiply(CPyTagged t_size, PyObject *seq) {
     return CPySequence_Multiply(seq, t_size);
 }
+
+PyObject *CPyList_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
+    if (likely(PyList_CheckExact(obj)
+               && CPyTagged_CheckShort(start) && CPyTagged_CheckShort(end))) {
+        Py_ssize_t startn = CPyTagged_ShortAsSsize_t(start);
+        Py_ssize_t endn = CPyTagged_ShortAsSsize_t(end);
+        if (startn < 0) {
+            startn += PyList_GET_SIZE(obj);
+        }
+        if (endn < 0) {
+            endn += PyList_GET_SIZE(obj);
+        }
+        return PyList_GetSlice(obj, startn, endn);
+    }
+    return CPyObject_GetSlice(obj, start, end);
+}

mypyc/lib-rt/setup.py

@@ -17,7 +17,7 @@
       version='0.1',
       ext_modules=[Extension(
           'test_capi',
-          ['test_capi.cc', 'init.c', 'int_ops.c', 'list_ops.c', 'exc_ops.c'],
+          ['test_capi.cc', 'init.c', 'int_ops.c', 'list_ops.c', 'exc_ops.c', 'generic_ops.c'],
           depends=['CPy.h', 'mypyc_util.h', 'pythonsupport.h'],
           extra_compile_args=['-Wno-unused-function', '-Wno-sign-compare'] + compile_args,
           library_dirs=['../external/googletest/make'],