BUG: core: implement a long-int loop for ldexp, for cases where int != long (#1633)

long != int on many 64-bit platforms, so a second ufunc loop is needed
to handle ldexp long int inputs.

5 files changed, 80 insertions, 12 deletions

diff --git a/numpy/core/src/umath/loops.c.src b/numpy/core/src/umath/loops.c.src
index 850126482..c61d16ae4 100644
--- a/numpy/core/src/umath/loops.c.src
+++ b/numpy/core/src/umath/loops.c.src
@@ -1286,6 +1286,36 @@ NPY_NO_EXPORT void
         *((@type@ *)op1) = ldexp@c@(in1, in2);
     }
 }
+
+NPY_NO_EXPORT void
+@TYPE@_ldexp_long(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func))
+{
+    /*
+     * Additional loop to handle long integer inputs (cf. #866, #1633).
+     * long != int on many 64-bit platforms, so we need this second loop
+     * to handle the default integer type.
+     */
+    BINARY_LOOP {
+        const @type@ in1 = *(@type@ *)ip1;
+        const long in2 = *(long *)ip2;
+        if (((int)in2) == in2) {
+            /* Range OK */
+            *((@type@ *)op1) = ldexp@c@(in1, ((int)in2));
+        }
+        else {
+            /*
+             * Outside int range -- also ldexp will overflow in this case,
+             * given that exponent has less bits than int.
+             */
+            if (in2 > 0) {
+                *((@type@ *)op1) = ldexp@c@(in1, NPY_MAX_INT);
+            }
+            else {
+                *((@type@ *)op1) = ldexp@c@(in1, NPY_MIN_INT);
+            }
+        }
+    }
+}
 #endif
 
 #define @TYPE@_true_divide @TYPE@_divide
diff --git a/numpy/core/src/umath/loops.h b/numpy/core/src/umath/loops.h
index 59e9556c2..7cb4b22cc 100644
--- a/numpy/core/src/umath/loops.h
+++ b/numpy/core/src/umath/loops.h
@@ -1675,6 +1675,8 @@ FLOAT_frexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
 #ifdef HAVE_LDEXPF
 NPY_NO_EXPORT void
 FLOAT_ldexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
+NPY_NO_EXPORT void
+FLOAT_ldexp_long(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
 #endif
 
 #define FLOAT_true_divide FLOAT_divide
@@ -1827,6 +1829,8 @@ DOUBLE_frexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func))
 #ifdef HAVE_LDEXP
 NPY_NO_EXPORT void
 DOUBLE_ldexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
+NPY_NO_EXPORT void
+DOUBLE_ldexp_long(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
 #endif
 
 #define DOUBLE_true_divide DOUBLE_divide
@@ -1979,6 +1983,8 @@ LONGDOUBLE_frexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(fu
 #ifdef HAVE_LDEXPL
 NPY_NO_EXPORT void
 LONGDOUBLE_ldexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
+NPY_NO_EXPORT void
+LONGDOUBLE_ldexp_long(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
 #endif
 
 #define LONGDOUBLE_true_divide LONGDOUBLE_divide
diff --git a/numpy/core/src/umath/loops.h.src b/numpy/core/src/umath/loops.h.src
index ca3fb37ee..33dfe882e 100644
--- a/numpy/core/src/umath/loops.h.src
+++ b/numpy/core/src/umath/loops.h.src
@@ -268,6 +268,8 @@ NPY_NO_EXPORT void
 #ifdef HAVE_LDEXP@C@
 NPY_NO_EXPORT void
 @TYPE@_ldexp(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
+NPY_NO_EXPORT void
+@TYPE@_ldexp_long(char **args, intp *dimensions, intp *steps, void *NPY_UNUSED(func));
 #endif
 
 #define @TYPE@_true_divide @TYPE@_divide
diff --git a/numpy/core/src/umath/umathmodule.c.src b/numpy/core/src/umath/umathmodule.c.src
index 9694f521d..c3da9f3d3 100644
--- a/numpy/core/src/umath/umathmodule.c.src
+++ b/numpy/core/src/umath/umathmodule.c.src
@@ -164,6 +164,8 @@ static PyUFuncGenericFunction frexp_functions[] = {
 };
 
 static void * blank3_data[] = { (void *)NULL, (void *)NULL, (void *)NULL};
+static void * blank6_data[] = { (void *)NULL, (void *)NULL, (void *)NULL,
+                                (void *)NULL, (void *)NULL, (void *)NULL};
 static char frexp_signatures[] = {
 #ifdef HAVE_FREXPF
     PyArray_FLOAT, PyArray_FLOAT, PyArray_INT,
@@ -174,22 +176,35 @@ static char frexp_signatures[] = {
 #endif
 };
 
+#if NPY_SIZEOF_LONG == NPY_SIZEOF_INT
+#define LDEXP_LONG(typ) typ##_ldexp
+#else
+#define LDEXP_LONG(typ) typ##_ldexp_long
+#endif
+
 static PyUFuncGenericFunction ldexp_functions[] = {
 #ifdef HAVE_LDEXPF
     FLOAT_ldexp,
+    LDEXP_LONG(FLOAT),
 #endif
-    DOUBLE_ldexp
+    DOUBLE_ldexp,
+    LDEXP_LONG(DOUBLE)
 #ifdef HAVE_LDEXPL
-    ,LONGDOUBLE_ldexp
+    ,
+    LONGDOUBLE_ldexp,
+    LDEXP_LONG(LONGDOUBLE)
 #endif
 };
 
 static char ldexp_signatures[] = {
 #ifdef HAVE_LDEXPF
     PyArray_FLOAT, PyArray_INT, PyArray_FLOAT,
+    PyArray_FLOAT, PyArray_LONG, PyArray_FLOAT,
 #endif
+    PyArray_DOUBLE, PyArray_INT, PyArray_DOUBLE,
     PyArray_DOUBLE, PyArray_LONG, PyArray_DOUBLE
 #ifdef HAVE_LDEXPL
+    ,PyArray_LONGDOUBLE, PyArray_INT, PyArray_LONGDOUBLE
     ,PyArray_LONGDOUBLE, PyArray_LONG, PyArray_LONGDOUBLE
 #endif
 };
@@ -213,14 +228,14 @@ InitOtherOperators(PyObject *dictionary) {
     PyDict_SetItemString(dictionary, "frexp", f);
     Py_DECREF(f);
 
-    num = 1;
+    num = 2;
 #ifdef HAVE_LDEXPL
-    num += 1;
+    num += 2;
 #endif
 #ifdef HAVE_LDEXPF
-    num += 1;
+    num += 2;
 #endif
-    f = PyUFunc_FromFuncAndData(ldexp_functions, blank3_data, ldexp_signatures, num,
+    f = PyUFunc_FromFuncAndData(ldexp_functions, blank6_data, ldexp_signatures, num,
                                 2, 1, PyUFunc_None, "ldexp",
                                 "Compute y = x1 * 2**x2.",0);
     PyDict_SetItemString(dictionary, "ldexp", f);
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index 2dace8c16..1245030a4 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -379,14 +379,29 @@ class TestArctan2SpecialValues(TestCase):
 
 
 class TestLdexp(TestCase):
+    def _check_ldexp(self, tp):
+        assert_almost_equal(ncu.ldexp(np.array(2., np.float32),
+                                      np.array(3, tp)), 16.)
+        assert_almost_equal(ncu.ldexp(np.array(2., np.float64),
+                                      np.array(3, tp)), 16.)
+        assert_almost_equal(ncu.ldexp(np.array(2., np.longdouble),
+                                      np.array(3, tp)), 16.)
+
     def test_ldexp(self):
+        # The default Python int type should work
         assert_almost_equal(ncu.ldexp(2., 3),  16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.float32), np.array(3, np.int16)), 16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.float32), np.array(3, np.int32)), 16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.float64), np.array(3, np.int16)), 16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.float64), np.array(3, np.int32)), 16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.longdouble), np.array(3, np.int16)), 16.)
-        assert_almost_equal(ncu.ldexp(np.array(2., np.longdouble), np.array(3, np.int32)), 16.)
+        # The following int types should all be accepted
+        self._check_ldexp(np.int8)
+        self._check_ldexp(np.int16)
+        self._check_ldexp(np.int32)
+        self._check_ldexp('i')
+        self._check_ldexp('l')
+
+    def test_ldexp_overflow(self):
+        imax = np.iinfo(np.dtype('l')).max
+        imin = np.iinfo(np.dtype('l')).min
+        assert_equal(ncu.ldexp(2., imax), np.inf)
+        assert_equal(ncu.ldexp(2., imin), 0)
 
 
 class TestMaximum(TestCase):