Merge pull request #3371 from juliantaylor/mmap-slice-improve

Mmap slice improve

7 files changed, 86 insertions, 64 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index ae07f05da..39dd2205e 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -3648,6 +3648,32 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('min',
     """))
 
 
+add_newdoc('numpy.core.multiarray', 'may_share_memory',
+    """
+    Determine if two arrays can share memory
+
+    The memory-bounds of a and b are computed.  If they overlap then
+    this function returns True.  Otherwise, it returns False.
+
+    A return of True does not necessarily mean that the two arrays
+    share any element.  It just means that they *might*.
+
+    Parameters
+    ----------
+    a, b : ndarray
+
+    Returns
+    -------
+    out : bool
+
+    Examples
+    --------
+    >>> np.may_share_memory(np.array([1,2]), np.array([5,8,9]))
+    False
+
+    """)
+
+
 add_newdoc('numpy.core.multiarray', 'ndarray', ('newbyteorder',
     """
     arr.newbyteorder(new_order='S')
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py
index a187d7c5b..bb2f3d28d 100644
--- a/numpy/core/numeric.py
+++ b/numpy/core/numeric.py
@@ -40,7 +40,7 @@ __all__ = ['newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc',
            'Inf', 'inf', 'infty', 'Infinity',
            'nan', 'NaN', 'False_', 'True_', 'bitwise_not',
            'CLIP', 'RAISE', 'WRAP', 'MAXDIMS', 'BUFSIZE', 'ALLOW_THREADS',
-           'ComplexWarning']
+           'ComplexWarning', 'may_share_memory']
 
 if sys.version_info[0] < 3:
     __all__.extend(['getbuffer', 'newbuffer'])
@@ -252,6 +252,7 @@ fromstring = multiarray.fromstring
 fromiter = multiarray.fromiter
 fromfile = multiarray.fromfile
 frombuffer = multiarray.frombuffer
+may_share_memory = multiarray.may_share_memory
 if sys.version_info[0] < 3:
     newbuffer = multiarray.newbuffer
     getbuffer = multiarray.getbuffer
diff --git a/numpy/core/src/multiarray/array_assign.c b/numpy/core/src/multiarray/array_assign.c
index 4f350388b..6467b6cfd 100644
--- a/numpy/core/src/multiarray/array_assign.c
+++ b/numpy/core/src/multiarray/array_assign.c
@@ -21,6 +21,7 @@
 #include "shape.h"
 
 #include "array_assign.h"
+#include "common.h"
 
 /* See array_assign.h for parameter documentation */
 NPY_NO_EXPORT int
@@ -102,36 +103,14 @@ raw_array_is_aligned(int ndim, char *data, npy_intp *strides, int alignment)
 /* Gets a half-open range [start, end) which contains the array data */
 NPY_NO_EXPORT void
 get_array_memory_extents(PyArrayObject *arr,
-                    npy_uintp *out_start, npy_uintp *out_end)
+                         npy_uintp *out_start, npy_uintp *out_end)
 {
-    npy_uintp start, end;
-    npy_intp idim, ndim = PyArray_NDIM(arr);
-    npy_intp *dimensions = PyArray_DIMS(arr),
-            *strides = PyArray_STRIDES(arr);
-
-    /* Calculate with a closed range [start, end] */
-    start = end = (npy_uintp)PyArray_DATA(arr);
-    for (idim = 0; idim < ndim; ++idim) {
-        npy_intp stride = strides[idim], dim = dimensions[idim];
-        /* If the array size is zero, return an empty range */
-        if (dim == 0) {
-            *out_start = *out_end = (npy_uintp)PyArray_DATA(arr);
-            return;
-        }
-        /* Expand either upwards or downwards depending on stride */
-        else {
-            if (stride > 0) {
-                end += stride*(dim-1);
-            }
-            else if (stride < 0) {
-                start += stride*(dim-1);
-            }
-        }
-    }
-
-    /* Return a half-open range */
-    *out_start = start;
-    *out_end = end + PyArray_DESCR(arr)->elsize;
+    npy_intp low, upper;
+    offset_bounds_from_strides(PyArray_ITEMSIZE(arr), PyArray_NDIM(arr),
+                               PyArray_DIMS(arr), PyArray_STRIDES(arr),
+                               &low, &upper);
+    *out_start = (npy_uintp)PyArray_DATA(arr) + (npy_uintp)low;
+    *out_end = (npy_uintp)PyArray_DATA(arr) + (npy_uintp)upper;
 }
 
 /* Returns 1 if the arrays have overlapping data, 0 otherwise */
diff --git a/numpy/core/src/multiarray/common.c b/numpy/core/src/multiarray/common.c
index 4f25dc913..01bf9cd74 100644
--- a/numpy/core/src/multiarray/common.c
+++ b/numpy/core/src/multiarray/common.c
@@ -688,6 +688,7 @@ _IsWriteable(PyArrayObject *ap)
 }
 
 
+/* Gets a half-open range [start, end) of offsets from the data pointer */
 NPY_NO_EXPORT void
 offset_bounds_from_strides(const int itemsize, const int nd,
                            const npy_intp *dims, const npy_intp *strides,
@@ -699,11 +700,12 @@ offset_bounds_from_strides(const int itemsize, const int nd,
 
     for (i = 0; i < nd; i++) {
         if (dims[i] == 0) {
-            /* Empty array special case */
+            /* If the array size is zero, return an empty range */
             *lower_offset = 0;
             *upper_offset = 0;
             return;
         }
+        /* Expand either upwards or downwards depending on stride */
         max_axis_offset = strides[i] * (dims[i] - 1);
         if (max_axis_offset > 0) {
             upper += max_axis_offset;
@@ -712,6 +714,7 @@ offset_bounds_from_strides(const int itemsize, const int nd,
             lower += max_axis_offset;
         }
     }
+    /* Return a half-open range */
     upper += itemsize;
     *lower_offset = lower;
     *upper_offset = upper;
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index a1c9e7c30..7332a26d0 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -3544,6 +3544,32 @@ PyDataMem_RENEW(void *ptr, size_t size)
     return (char *)result;
 }
 
+static PyObject *
+array_may_share_memory(PyObject *NPY_UNUSED(ignored), PyObject *args)
+{
+    PyArrayObject * self = NULL;
+    PyArrayObject * other = NULL;
+    int overlap;
+
+    if (!PyArg_ParseTuple(args, "O&O&", PyArray_Converter, &self,
+                          PyArray_Converter, &other)) {
+        return NULL;
+    }
+
+    overlap = arrays_overlap(self, other);
+    Py_XDECREF(self);
+    Py_XDECREF(other);
+
+    if (overlap) {
+        Py_RETURN_TRUE;
+    }
+    else {
+        Py_RETURN_FALSE;
+    }
+}
+
+
+
 static struct PyMethodDef array_module_methods[] = {
     {"_get_ndarray_c_version",
         (PyCFunction)array__get_ndarray_c_version,
@@ -3644,6 +3670,9 @@ static struct PyMethodDef array_module_methods[] = {
     {"result_type",
         (PyCFunction)array_result_type,
         METH_VARARGS, NULL},
+    {"may_share_memory",
+        (PyCFunction)array_may_share_memory,
+        METH_VARARGS, NULL},
     /* Datetime-related functions */
     {"datetime_data",
         (PyCFunction)array_datetime_data,
diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index 3c270088f..a92ddde83 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -315,6 +315,21 @@ class TestTile(TestCase):
                 assert_equal(large, klarge)
 
 
+class TestMayShareMemory(TestCase):
+    def test_basic(self):
+        d = ones((50, 60))
+        d2 = ones((30, 60, 6))
+        self.assertTrue(may_share_memory(d, d))
+        self.assertTrue(may_share_memory(d, d[::-1]))
+        self.assertTrue(may_share_memory(d, d[::2]))
+        self.assertTrue(may_share_memory(d, d[1:, ::-1]))
+
+        self.assertFalse(may_share_memory(d[::-1], d2))
+        self.assertFalse(may_share_memory(d[::2], d2))
+        self.assertFalse(may_share_memory(d[1:, ::-1], d2))
+        self.assertTrue(may_share_memory(d2[1:, ::-1], d2))
+
+
 # Utility
 def compare_results(res,desired):
     for i in range(len(desired)):
diff --git a/numpy/lib/utils.py b/numpy/lib/utils.py
index f94abeeab..f19a47176 100644
--- a/numpy/lib/utils.py
+++ b/numpy/lib/utils.py
@@ -11,7 +11,7 @@ from numpy.core import product, ndarray, ufunc
 __all__ = ['issubclass_', 'issubsctype', 'issubdtype',
         'deprecate', 'deprecate_with_doc', 'get_numarray_include',
         'get_include', 'info', 'source', 'who', 'lookfor', 'byte_bounds',
-        'may_share_memory', 'safe_eval']
+        'safe_eval']
 
 def get_include():
     """
@@ -250,12 +250,11 @@ def byte_bounds(a):
     a_data = ai['data'][0]
     astrides = ai['strides']
     ashape = ai['shape']
-    nd_a = len(ashape)
     bytes_a = int(ai['typestr'][2:])
 
     a_low = a_high = a_data
     if astrides is None: # contiguous case
-        a_high += product(ashape, dtype=int)*bytes_a
+        a_high += a.size * bytes_a
     else:
         for shape, stride in zip(ashape, astrides):
             if stride < 0:
@@ -266,36 +265,6 @@ def byte_bounds(a):
     return a_low, a_high
 
 
-def may_share_memory(a, b):
-    """
-    Determine if two arrays can share memory
-
-    The memory-bounds of a and b are computed.  If they overlap then
-    this function returns True.  Otherwise, it returns False.
-
-    A return of True does not necessarily mean that the two arrays
-    share any element.  It just means that they *might*.
-
-    Parameters
-    ----------
-    a, b : ndarray
-
-    Returns
-    -------
-    out : bool
-
-    Examples
-    --------
-    >>> np.may_share_memory(np.array([1,2]), np.array([5,8,9]))
-    False
-
-    """
-    a_low, a_high = byte_bounds(a)
-    b_low, b_high = byte_bounds(b)
-    if b_low >= a_high or a_low >= b_high:
-        return False
-    return True
-
 #-----------------------------------------------------------------------------
 # Function for output and information on the variables used.
 #-----------------------------------------------------------------------------