Skip to content

Commit 78731c0

Browse files
gh-102221: Optimize math.lcm() for multiple arguments
1 parent 654b8d9 commit 78731c0

1 file changed

Lines changed: 31 additions & 26 deletions

File tree

Modules/mathmodule.c

Lines changed: 31 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -802,41 +802,46 @@ math_lcm_impl(PyObject *module, PyObject * const *args,
802802
Py_ssize_t args_length)
803803
/*[clinic end generated code: output=c8a59a5c2e55c816 input=3e4f4b7cdf948a98]*/
804804
{
805-
PyObject *res, *x;
806-
Py_ssize_t i;
807-
808805
if (args_length == 0) {
809806
return PyLong_FromLong(1);
810807
}
811-
res = PyNumber_Index(args[0]);
812-
if (res == NULL) {
813-
return NULL;
808+
PyObject *res;
809+
PyObject *stack[8 * sizeof(Py_ssize_t)];
810+
int top = 0;
811+
Py_ssize_t i = 0;
812+
while (1) {
813+
size_t j = i;
814+
res = PyNumber_Index(args[i++]);
815+
if (res == NULL) {
816+
goto error;
817+
}
818+
if (i >= args_length) {
819+
j = ((size_t)1 << top) - 1;
820+
}
821+
for (; j & 1; j >>= 1) {
822+
top--;
823+
Py_SETREF(res, long_lcm(res, stack[top]));
824+
if (res == NULL) {
825+
goto error;
826+
}
827+
Py_DECREF(stack[top]);
828+
}
829+
if (i >= args_length) {
830+
break;
831+
}
832+
stack[top++] = res;
814833
}
815834
if (args_length == 1) {
816835
Py_SETREF(res, PyNumber_Absolute(res));
817-
return res;
818836
}
837+
return res;
819838

820-
PyObject *zero = _PyLong_GetZero(); // borrowed ref
821-
for (i = 1; i < args_length; i++) {
822-
x = PyNumber_Index(args[i]);
823-
if (x == NULL) {
824-
Py_DECREF(res);
825-
return NULL;
826-
}
827-
if (res == zero) {
828-
/* Fast path: just check arguments.
829-
It is okay to use identity comparison here. */
830-
Py_DECREF(x);
831-
continue;
832-
}
833-
Py_SETREF(res, long_lcm(res, x));
834-
Py_DECREF(x);
835-
if (res == NULL) {
836-
return NULL;
837-
}
839+
error:
840+
while (top > 0) {
841+
top--;
842+
Py_DECREF(stack[top]);
838843
}
839-
return res;
844+
return NULL;
840845
}
841846

842847

0 commit comments

Comments
 (0)