pythongh-94906: Support multiple steps in math.nextafter

Author: matthiasgoergens

Doc/library/math.rst

@@ -224,9 +224,9 @@ Number-theoretic and representation functions
    of *x* and are floats.
 
 
-.. function:: nextafter(x, y)
+.. function:: nextafter(x, y, /, *, steps=1)
 
-   Return the next floating-point value after *x* towards *y*.
+   Return the floating-point value *steps* steps after *x* towards *y*.
 
    If *x* is equal to *y*, return *y*.
 
@@ -239,6 +239,9 @@ Number-theoretic and representation functions
 
    See also :func:`math.ulp`.
 
+   .. versionchanged:: 3.12
+      Added the *steps* argument.
+
    .. versionadded:: 3.9
 
 .. function:: perm(n, k=None)

Include/internal/pycore_global_objects_fini_generated.h

@@ -657,6 +657,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(stdin)
         STRUCT_FOR_ID(stdout)
         STRUCT_FOR_ID(step)
+        STRUCT_FOR_ID(steps)
         STRUCT_FOR_ID(store_name)
         STRUCT_FOR_ID(strategy)
         STRUCT_FOR_ID(strftime)

Include/internal/pycore_global_strings.h

@@ -2296,11 +2296,20 @@ def test_nextafter(self):
                          float.fromhex('0x1.fffffffffffffp-1'))
         self.assertEqual(math.nextafter(1.0, INF),
                          float.fromhex('0x1.0000000000001p+0'))
+        self.assertEqual(math.nextafter(1.0, -INF, steps=1),
+                         float.fromhex('0x1.fffffffffffffp-1'))
+        self.assertEqual(math.nextafter(1.0, INF, steps=1),
+                         float.fromhex('0x1.0000000000001p+0'))
+        self.assertEqual(math.nextafter(1.0, -INF, steps=3),
+                         float.fromhex('0x1.ffffffffffffdp-1'))
+        self.assertEqual(math.nextafter(1.0, INF, steps=3),
+                         float.fromhex('0x1.0000000000003p+0'))
 
         # x == y: y is returned
-        self.assertEqual(math.nextafter(2.0, 2.0), 2.0)
-        self.assertEqualSign(math.nextafter(-0.0, +0.0), +0.0)
-        self.assertEqualSign(math.nextafter(+0.0, -0.0), -0.0)
+        for steps in range(1, 5):
+            self.assertEqual(math.nextafter(2.0, 2.0, steps=steps), 2.0)
+            self.assertEqualSign(math.nextafter(-0.0, +0.0, steps=steps), +0.0)
+            self.assertEqualSign(math.nextafter(+0.0, -0.0, steps=steps), -0.0)
 
         # around 0.0
         smallest_subnormal = sys.float_info.min * sys.float_info.epsilon
@@ -2325,6 +2334,11 @@ def test_nextafter(self):
         self.assertIsNaN(math.nextafter(1.0, NAN))
         self.assertIsNaN(math.nextafter(NAN, NAN))
 
+        self.assertEqual(1.0, math.nextafter(1.0, INF, steps=0))
+        with self.assertRaises(ValueError):
+            math.nextafter(1.0, INF, steps=-1)
+
+
     @requires_IEEE_754
     def test_ulp(self):
         self.assertEqual(math.ulp(1.0), sys.float_info.epsilon)

Include/internal/pycore_runtime_init_generated.h

@@ -0,0 +1 @@
+Support multiple steps in :func:`math.nextafter`. Patch by Shantanu Jain and Matthias Gorgens.

Include/internal/pycore_unicodeobject_generated.h

@@ -3893,13 +3893,15 @@ math.nextafter
     x: double
     y: double
     /
+    *
+    steps: int = 1
 
-Return the next floating-point value after x towards y.
+Return the floating-point value the given number of steps after x towards y.
 [clinic start generated code]*/
 
 static PyObject *
-math_nextafter_impl(PyObject *module, double x, double y)
-/*[clinic end generated code: output=750c8266c1c540ce input=02b2d50cd1d9f9b6]*/
+math_nextafter_impl(PyObject *module, double x, double y, int steps)
+/*[clinic end generated code: output=14190eb869199e5a input=a794e7a79768ee25]*/
 {
 #if defined(_AIX)
     if (x == y) {
@@ -3914,7 +3916,66 @@ math_nextafter_impl(PyObject *module, double x, double y)
         return PyFloat_FromDouble(y);
     }
 #endif
-    return PyFloat_FromDouble(nextafter(x, y));
+    // fast path:
+    if (steps == 1) {
+        return PyFloat_FromDouble(nextafter(x, y));
+    }
+    if (steps < 0) {
+        PyErr_SetString(PyExc_ValueError, "steps must be >= 0");
+        return NULL;
+    }
+    if (steps == 0)
+        return PyFloat_FromDouble(x);
+    if (Py_IS_NAN(x) || Py_IS_NAN(y))
+        return PyFloat_FromDouble(x+y);
+
+    uint64_t usteps = steps;
+
+    union pun {double f; uint64_t i;};
+    union pun ux = {x}, uy = {y};
+    if(ux.i == uy.i) {
+        return PyFloat_FromDouble(x);
+    }
+
+    const uint64_t sign_bit = 1ULL<<63;
+
+    uint64_t ax = ux.i & ~sign_bit;
+	uint64_t ay = uy.i & ~sign_bit;
+
+    // opposite signs
+    if (((ux.i ^ uy.i) & sign_bit)) {
+        if (ax + ay <= usteps) {
+            return PyFloat_FromDouble(uy.f);
+        // We can't use <= here, because of +0.0 vs 0.0
+        } else if (ax < usteps) {
+            union pun result = {.i = (uy.i & sign_bit) | (usteps - ax)};
+            return PyFloat_FromDouble(result.f);
+        } else {
+            ux.i -= usteps;
+            return PyFloat_FromDouble(ux.f);
+        }
+    // same sign now.
+    // going down:
+    } else if (ax > ay) {
+        // going too far:
+        // the addition is not UB,
+        // because we have an extra bit at the top of ax and usteps.
+        if (ax >= ay + usteps) {
+            ux.i -= usteps;
+            return PyFloat_FromDouble(ux.f);
+        } else {
+            return PyFloat_FromDouble(uy.f);
+        }
+    } else { // go up in magnitude. same sign.
+        // the addition is not UB,
+        // because we have an extra bit at the top of ax and usteps.
+        if (ax + usteps <= ay) {
+            ux.i += usteps;
+            return PyFloat_FromDouble(ux.f);
+        } else {
+            return PyFloat_FromDouble(uy.f);
+        }
+    }
 }