gh-94906: Support multiple steps in math.nextafter (#103881)

This PR updates `math.nextafter` to add a new `steps` argument. The behaviour is as though `math.nextafter` had been called `steps` times in succession.

---------

Co-authored-by: Mark Dickinson <mdickinson@enthought.com>
This commit is contained in:
Matthias Görgens 2023-05-20 04:03:49 +08:00 committed by GitHub
parent c3f43bfb4b
commit 6e39fa1955
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 223 additions and 18 deletions

View file

@ -3864,13 +3864,20 @@ math.nextafter
x: double
y: double
/
*
steps: object = None
Return the next floating-point value after x towards y.
Return the floating-point value the given number of steps after x towards y.
If steps is not specified or is None, it defaults to 1.
Raises a TypeError, if x or y is not a double, or if steps is not an integer.
Raises ValueError if steps is negative.
[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, PyObject *steps)
/*[clinic end generated code: output=cc6511f02afc099e input=7f2a5842112af2b4]*/
{
#if defined(_AIX)
if (x == y) {
@ -3885,7 +3892,101 @@ math_nextafter_impl(PyObject *module, double x, double y)
return PyFloat_FromDouble(y);
}
#endif
return PyFloat_FromDouble(nextafter(x, y));
if (steps == Py_None) {
// fast path: we default to one step.
return PyFloat_FromDouble(nextafter(x, y));
}
steps = PyNumber_Index(steps);
if (steps == NULL) {
return NULL;
}
assert(PyLong_CheckExact(steps));
if (_PyLong_IsNegative((PyLongObject *)steps)) {
PyErr_SetString(PyExc_ValueError,
"steps must be a non-negative integer");
Py_DECREF(steps);
return NULL;
}
unsigned long long usteps_ull = PyLong_AsUnsignedLongLong(steps);
// Conveniently, uint64_t and double have the same number of bits
// on all the platforms we care about.
// So if an overflow occurs, we can just use UINT64_MAX.
Py_DECREF(steps);
if (usteps_ull >= UINT64_MAX) {
// This branch includes the case where an error occurred, since
// (unsigned long long)(-1) = ULLONG_MAX >= UINT64_MAX. Note that
// usteps_ull can be strictly larger than UINT64_MAX on a machine
// where unsigned long long has width > 64 bits.
if (PyErr_Occurred()) {
if (PyErr_ExceptionMatches(PyExc_OverflowError)) {
PyErr_Clear();
}
else {
return NULL;
}
}
usteps_ull = UINT64_MAX;
}
assert(usteps_ull <= UINT64_MAX);
uint64_t usteps = (uint64_t)usteps_ull;
if (usteps == 0) {
return PyFloat_FromDouble(x);
}
if (Py_IS_NAN(x)) {
return PyFloat_FromDouble(x);
}
if (Py_IS_NAN(y)) {
return PyFloat_FromDouble(y);
}
// We assume that double and uint64_t have the same endianness.
// This is not guaranteed by the C-standard, but it is true for
// all platforms we care about. (The most likely form of violation
// would be a "mixed-endian" double.)
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)) {
// NOTE: ax + ay can never overflow, because their most significant bit
// ain't set.
if (ax + ay <= usteps) {
return PyFloat_FromDouble(uy.f);
// This comparison has to use <, because <= would get +0.0 vs -0.0
// wrong.
} 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
} else if (ax > ay) {
if (ax - ay >= usteps) {
ux.i -= usteps;
return PyFloat_FromDouble(ux.f);
} else {
return PyFloat_FromDouble(uy.f);
}
} else {
if (ay - ax >= usteps) {
ux.i += usteps;
return PyFloat_FromDouble(ux.f);
} else {
return PyFloat_FromDouble(uy.f);
}
}
}