Support index-slicing indirect dimensions if slice data pointer can be moved

Cython/Compiler/ExprNodes.py

@@ -2554,13 +2554,6 @@ class IndexNode(ExprNode):
                     indices[i] = index
                     new_indices.append(index)
 
-                    if access in ('ptr', 'generic') and i != 0 and have_slices:
-                        self.type = error_type
-                        return error(index.pos,
-                                     "Indexing of non-leading indirect or generic "
-                                     "dimensions not supported yet, "
-                                     "try slicing with i:i+1")
-
                 else:
                     self.type = error_type
                     return error(index.pos, "Invalid index for memoryview specified")

Cython/Compiler/MemoryView.py

@@ -305,14 +305,22 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
             if not isinstance(index, ExprNodes.SliceNode):
                 # normal index
                 idx = index.result()
-                d = locals()
+
                 access, packing = self.type.axes[dim]
-                if access != 'direct':
-                    return error(index.pos,
-                                 "Dimension cannot be indexed away, "
-                                 "must be direct")
+                if access == 'direct':
+                    indirect = False
+                else:
+                    indirect = True
+                    generic = (access == 'full')
+                    if new_ndim != 0:
+                        return error(index.pos,
+                                     "All preceding dimensions must be "
+                                     "indexed and not sliced")
+
+                d = locals()
                 code.put(load_slice_util("SliceIndex", d))
             else:
+
                 # slice, unspecified dimension, or part of ellipsis
                 d = locals()
                 for s in "start stop step".split():

Cython/Utility/MemoryView.pyx

@@ -730,12 +730,12 @@ cdef int slice_memviewslice({{memviewslice_name}} *src,
 
     if suboffset >= 0:
         if not is_slice:
-            with gil:
-                raise IndexError(
-                    "Cannot make indirect dimension %d disappear "
-                    "through indexing, consider slicing with %d:%d" %
-                                                (dim, start, start + 1))
-
+            if new_ndim == 0:
+                dst.data = (<char **> dst.data)[0] + suboffset
+            else:
+                with gil:
+                    raise IndexError("All dimensions preceding dimension %d "
+                                     "must be indexed and not sliced" % dim)
         else:
             suboffset_dim[0] = new_ndim
 

Cython/Utility/MemoryView_C.c

@@ -729,29 +729,59 @@ if (unlikely(__pyx_memoryview_slice_memviewslice(
 
 {
     Py_ssize_t __pyx_tmp_idx = {{idx}};
-
+    Py_ssize_t __pyx_tmp_shape = {{src}}.shape[{{dim}}];
+    Py_ssize_t __pyx_tmp_stride = {{src}}.strides[{{dim}}];
     if (__pyx_tmp_idx < 0)
-        __pyx_tmp_idx += {{src}}.shape[{{dim}}];
+        __pyx_tmp_idx += __pyx_tmp_shape;
+
+    if (__pyx_tmp_idx < 0 || __pyx_tmp_idx >= __pyx_tmp_shape) {
+        {{if not have_gil}}
+            #ifdef WITH_THREAD
+            PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure();
+            #endif
+        {{endif}}
 
-    if (__pyx_tmp_idx < 0 || __pyx_tmp_idx >= {{src}}.shape[{{dim}}]) {
-    {{if not have_gil}}
-        #ifdef WITH_THREAD
-        PyGILState_STATE __pyx_gilstate_save = PyGILState_Ensure();
-        #endif
-    {{endif}}
         PyErr_SetString(PyExc_IndexError, "Index out of bounds (axis {{dim}})");
-    {{if not have_gil}}
-        #ifdef WITH_THREAD
-        PyGILState_Release(__pyx_gilstate_save);
-        #endif
-    {{endif}}
+
+        {{if not have_gil}}
+            #ifdef WITH_THREAD
+            PyGILState_Release(__pyx_gilstate_save);
+            #endif
+        {{endif}}
+
         {{error_goto}}
     }
 
     {{if all_dimensions_direct}}
-        {{dst}}.data += __pyx_tmp_idx * {{src}}.strides[{{dim}}];
+        {{dst}}.data += __pyx_tmp_idx * __pyx_tmp_stride;
     {{else}}
-        if ({{suboffset_dim}} < 0) {{dst}}.data += __pyx_tmp_idx * {{src}}.strides[{{dim}}];
-        else {{dst}}.suboffsets[{{suboffset_dim}}] = __pyx_tmp_idx * {{src}}.strides[{{dim}}];
+        if ({{suboffset_dim}} < 0) {
+            {{dst}}.data += __pyx_tmp_idx * __pyx_tmp_stride;
+
+            /* This dimension is the first dimension, or is preceded by    */
+            /* direct or indirect dimensions that are indexed away.        */
+            /* Hence suboffset_dim must be less than zero, and we can have */
+            /* our data pointer refer to another block by dereferencing.   */
+            /*   slice.data -> B -> C     becomes     slice.data -> C      */
+
+            {{if indirect}}
+              {
+                Py_ssize_t __pyx_tmp_suboffset = {{src}}.suboffsets[{{dim}}];
+
+                {{if generic}}
+                    if (__pyx_tmp_suboffset >= 0)
+                {{endif}}
+
+                    {{dst}}.data = *((char **) {{dst}}.data) + __pyx_tmp_suboffset;
+              }
+            {{endif}}
+
+        } else {
+            {{dst}}.suboffsets[{{suboffset_dim}}] += __pyx_tmp_idx * __pyx_tmp_stride;
+
+            /* Note: dimension can not be indirect, the compiler will have */
+            /*       issued an error */
+        }
+
     {{endif}}
 }

tests/run/memoryview.pyx

@@ -652,6 +652,11 @@ def test_indirect_slicing(arg):
     (5, 3, 2)
     0 0 -1
     58
+    56
+    58
+    58
+    58
+    58
     released A
 
     >>> test_indirect_slicing(IntMockBuffer("A", shape_9_14_21_list, shape=(9, 14, 21)))
@@ -659,6 +664,11 @@ def test_indirect_slicing(arg):
     (5, 14, 3)
     0 16 -1
     2412
+    2410
+    2412
+    2412
+    2412
+    2412
     released A
     """
     cdef int[::view.indirect, ::view.indirect, :] _a = arg
@@ -669,6 +679,11 @@ def test_indirect_slicing(arg):
     print_int_offsets(*b.suboffsets)
 
     print b[4, 2, 1]
+    print b[..., 0][4, 2]
+    print b[..., 1][4, 2]
+    print b[..., 1][4][2]
+    print b[4][2][1]
+    print b[4, 2][1]
 
 def test_direct_slicing(arg):
     """

tests/run/memslice.pyx

@@ -1240,6 +1240,13 @@ def test_indirect_slicing(arg):
     0 0 -1
     58
     56
+    58
+    index away indirect
+    58
+    58
+    index away generic
+    58
+    58
     released A
 
     >>> test_indirect_slicing(IntMockBuffer("A", shape_9_14_21_list, shape=(9, 14, 21)))
@@ -1248,12 +1255,27 @@ def test_indirect_slicing(arg):
     0 16 -1
     2412
     2410
+    2412
+    index away indirect
+    2412
+    2412
+    index away generic
+    2412
+    2412
     released A
     """
     cdef int[::view.indirect, ::view.indirect, :] a = arg
     cdef int[::view.indirect, ::view.indirect, :] b = a[-5:, ..., -5:100:2]
+    cdef int[::view.generic , :: view.generic, :] generic_b = a[-5:, ..., -5:100:2]
     cdef int[::view.indirect, ::view.indirect] c = b[..., 0]
 
+    # try indexing away leading indirect dimensions
+    cdef int[::view.indirect, :] d = b[4]
+    cdef int[:] e = b[4, 2]
+
+    cdef int[::view.generic, :] generic_d = generic_b[4]
+    cdef int[:] generic_e = generic_b[4, 2]
+
     print b.shape[0], b.shape[1], b.shape[2]
     print b.suboffsets[0] // sizeof(int *),
     print b.suboffsets[1] // sizeof(int),
@@ -1261,6 +1283,90 @@ def test_indirect_slicing(arg):
 
     print b[4, 2, 1]
     print c[4, 2]
+    # test adding offset from last dimension to suboffset
+    print b[..., 1][4, 2]
+
+    print "index away indirect"
+    print d[2, 1]
+    print e[1]
+
+    print "index away generic"
+    print generic_d[2, 1]
+    print generic_e[1]
+
+cdef class TestIndexSlicingDirectIndirectDims(object):
+    "Test a int[:, ::view.indirect, :] slice"
+
+    cdef Py_ssize_t[3] shape, strides, suboffsets
+
+    cdef int c_array[5]
+    cdef int *myarray[5][5]
+    cdef bytes format
+
+    def __init__(self):
+        cdef int i
+        self.c_array[3] = 20
+        self.myarray[1][2] = self.c_array
+
+        for i in range(3):
+            self.shape[i] = 5
+
+        self.strides[0] = sizeof(int *) * 5
+        self.strides[1] = sizeof(int *)
+        self.strides[2] = sizeof(int)
+
+        self.suboffsets[0] = -1
+        self.suboffsets[1] = 0
+        self.suboffsets[2] = -1
+
+        self.format = b"i"
+
+    def __getbuffer__(self, Py_buffer *info, int flags):
+        info.buf = <void *> self.myarray
+        info.len = 5 * 5 * 5
+        info.ndim = 3
+        info.shape = self.shape
+        info.strides = self.strides
+        info.suboffsets = self.suboffsets
+        info.itemsize = sizeof(int)
+        info.readonly = 0
+        info.obj = self
+        info.format = self.format
+
+@testcase
+def test_index_slicing_away_direct_indirect():
+    """
+    >>> test_index_slicing_away_direct_indirect()
+    20
+    20
+    20
+    20
+    <BLANKLINE>
+    20
+    20
+    20
+    20
+    All dimensions preceding dimension 1 must be indexed and not sliced
+    """
+    cdef int[:, ::view.indirect, :] a = TestIndexSlicingDirectIndirectDims()
+    a_obj = a
+
+    print a[1][2][3]
+    print a[1, 2, 3]
+    print a[1, 2][3]
+    print a[..., 3][1, 2]
+
+    print
+
+    print a_obj[1][2][3]
+    print a_obj[1, 2, 3]
+    print a_obj[1, 2][3]
+    print a_obj[..., 3][1, 2]
+
+    try:
+        print a_obj[1:, 2][3]
+    except IndexError, e:
+        print e.args[0]
 
 @testcase
 def test_direct_slicing(arg):