Skip to content

Commit 2dd7568

Browse files
dg-pbdg-pb
committed
V4
1 parent d840ad7 commit 2dd7568

File tree

1 file changed

+96
-66
lines changed

1 file changed

+96
-66
lines changed

‎Modules/_functoolsmodule.c‎

Lines changed: 96 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,27 @@ partial_descr_get(PyObject *self, PyObject *obj, PyObject *type)
334334
return PyMethod_New(self, obj);
335335
}
336336

337+
#define ALLOCATE_STACK(type, size, small_stack, stack) \
338+
do { \
339+
if (size <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) { \
340+
stack = small_stack; \
341+
} \
342+
else { \
343+
stack = PyMem_Malloc(size * sizeof(type *)); \
344+
if (stack == NULL) { \
345+
PyErr_NoMemory(); \
346+
return NULL; \
347+
} \
348+
} \
349+
} while (0)
350+
351+
#define DEALLOCATE_STACK(small_stack, stack) \
352+
do { \
353+
if (stack != small_stack) { \
354+
PyMem_Free(stack); \
355+
} \
356+
} while (0)
357+
337358
static PyObject *
338359
partial_vectorcall(partialobject *pto, PyObject *const *args,
339360
size_t nargsf, PyObject *kwnames)
@@ -355,13 +376,10 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
355376
return NULL;
356377
}
357378

358-
Py_ssize_t tot_nargs = pto_nargs + nargs - pto_phcount;
359-
Py_ssize_t tot_nkwds;
360-
Py_ssize_t tot_nargskw;
361379
PyObject **pto_args = _PyTuple_ITEMS(pto->args);
362380
PyObject *ret;
363381

364-
/* Divergence on whether pto has kwds */
382+
/* Special cases */
365383
if (pto_nkwds == 0) {
366384
/* Fast path if we're called without arguments */
367385
if (nargskw == 0) {
@@ -370,7 +388,8 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
370388
}
371389

372390
/* Use PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
373-
* positional argument */
391+
* positional argument.
392+
* NOTE: Could add (Placeholder, arg) case in similar manner */
374393
if (pto_nargs == 1 && (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET)) {
375394
PyObject **newargs = (PyObject **)args - 1;
376395
PyObject *tmp = newargs[0];
@@ -381,22 +400,21 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
381400
return ret;
382401
}
383402

403+
}
404+
405+
/* Size variables */
406+
Py_ssize_t tot_nargs = pto_nargs + nargs - pto_phcount;
407+
Py_ssize_t tot_nkwds;
408+
Py_ssize_t tot_nargskw;
409+
410+
/* Divergence on whether pto has kwds */
411+
if (pto_nkwds == 0) {
384412
tot_nkwds = nkwds;
385413
tot_nargskw = tot_nargs + tot_nkwds;
386414

387415
/* Allocate Stack */
388-
PyObject *small_stack[_PY_FASTCALL_SMALL_STACK];
389-
PyObject **stack;
390-
if (tot_nargskw <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
391-
stack = small_stack;
392-
}
393-
else {
394-
stack = PyMem_Malloc(tot_nargskw * sizeof(PyObject *));
395-
if (stack == NULL) {
396-
PyErr_NoMemory();
397-
return NULL;
398-
}
399-
}
416+
PyObject **stack, *small_stack[_PY_FASTCALL_SMALL_STACK];
417+
ALLOCATE_STACK(PyObject, tot_nargskw, small_stack, stack);
400418

401419
/* Copy Positionals to new stack */
402420
if (pto_phcount) {
@@ -420,56 +438,52 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
420438
memcpy(stack, pto_args, pto_nargs * sizeof(PyObject*));
421439
memcpy(stack + pto_nargs, args, nargskw * sizeof(PyObject*));
422440
}
423-
ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack, tot_nargs, kwnames);
424441

425-
/* Free stack & Return */
426-
if (stack != small_stack) {
427-
PyMem_Free(stack);
428-
}
442+
/* Call / Maintenance / Return */
443+
ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack, tot_nargs, kwnames);
444+
DEALLOCATE_STACK(small_stack, stack);
429445
return ret;
430446
}
431447
else {
432-
PyObject *tot_kw = NULL;
433-
PyObject *tot_kwnames = NULL;
434448
PyObject *key, *val;
435-
int has_nkwds = 0;
436449

437-
if (!nkwds) {
438-
tot_kw = pto->kw;
439-
tot_nkwds = pto_nkwds;
440-
tot_nargskw = tot_nargs + tot_nkwds;
450+
PyObject **pto_kwkeys, *small_pto_kwkeys[_PY_FASTCALL_SMALL_STACK];
451+
PyObject **pto_kwvals, *small_pto_kwvals[_PY_FASTCALL_SMALL_STACK];
452+
ALLOCATE_STACK(PyObject, pto_nkwds, small_pto_kwkeys, pto_kwkeys);
453+
ALLOCATE_STACK(PyObject, pto_nkwds, small_pto_kwvals, pto_kwvals);
454+
Py_ssize_t pos = 0, i = 0;
455+
while (PyDict_Next(pto->kw, &pos, &key, &val)) {
456+
pto_kwkeys[i] = key;
457+
pto_kwvals[i] = val;
458+
i++;
441459
}
442-
else {
443-
has_nkwds = 1;
444-
tot_kw = PyDict_Copy(pto->kw);
460+
461+
/* Calculate total kw size and positions of items to override */
462+
Py_ssize_t *positions, small_positions[_PY_FASTCALL_SMALL_STACK];
463+
tot_nkwds = pto_nkwds + nkwds;
464+
if (nkwds) {
465+
ALLOCATE_STACK(Py_ssize_t, nkwds, small_positions, positions);
445466
for (Py_ssize_t i=0; i < nkwds; i++) {
467+
positions[i] = -1;
446468
key = PyTuple_GET_ITEM(kwnames, i);
447-
val = args[nargs + i];
448-
if (PyDict_SetItem(tot_kw, key, val)) {
449-
Py_DECREF(tot_kw);
450-
return NULL;
469+
if (PyDict_Contains(pto->kw, key)) {
470+
tot_nkwds--;
471+
for (Py_ssize_t j=0; j < pto_nkwds; j++) {
472+
if (PyObject_RichCompareBool(key, pto_kwkeys[j], Py_EQ)) {
473+
positions[i] = j;
474+
break;
475+
}
476+
}
451477
}
452478
}
453-
454-
tot_nkwds = PyDict_GET_SIZE(tot_kw);
455-
tot_nargskw = tot_nargs + tot_nkwds;
456479
}
480+
tot_nargskw = tot_nargs + tot_nkwds;
457481

458482
/* Allocate Stack */
459-
PyObject *small_stack[_PY_FASTCALL_SMALL_STACK];
460-
PyObject **stack;
461-
if (tot_nargskw <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
462-
stack = small_stack;
463-
}
464-
else {
465-
stack = PyMem_Malloc(tot_nargskw * sizeof(PyObject *));
466-
if (stack == NULL) {
467-
PyErr_NoMemory();
468-
return NULL;
469-
}
470-
}
483+
PyObject **stack, *small_stack[_PY_FASTCALL_SMALL_STACK];
484+
ALLOCATE_STACK(PyObject, tot_nargskw, small_stack, stack);
471485

472-
/* Copy Positionals to new stack */
486+
/* Copy Positionals to stack */
473487
if (pto_phcount) {
474488
Py_ssize_t j = 0; // New args index
475489
for (Py_ssize_t i = 0; i < pto_nargs; i++) {
@@ -492,26 +506,43 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
492506
memcpy(stack + pto_nargs, args, nargs * sizeof(PyObject*));
493507
}
494508

495-
/* Copy Keywords to new stack */
496-
tot_kwnames = PyTuple_New(tot_nkwds);
497-
Py_ssize_t pos = 0, i = 0;
498-
while (PyDict_Next(tot_kw, &pos, &key, &val)) {
509+
/* Copy Pto Keywords to stack */
510+
memcpy(stack + tot_nargs, pto_kwvals, pto_nkwds * sizeof(PyObject*));
511+
DEALLOCATE_STACK(small_pto_kwvals, pto_kwvals);
512+
513+
/* Copy New Keywords to stack */
514+
PyObject *tot_kwnames = PyTuple_New(tot_nkwds);
515+
for (Py_ssize_t i=0; i < pto_nkwds; i++) {
516+
key = pto_kwkeys[i];
499517
PyTuple_SET_ITEM(tot_kwnames, i, key);
500518
Py_INCREF(key);
501-
stack[tot_nargs + i] = val;
502-
i += 1;
503519
}
504-
if (has_nkwds) {
505-
Py_DECREF(tot_kw);
520+
DEALLOCATE_STACK(small_pto_kwkeys, pto_kwkeys);
521+
522+
if (nkwds) {
523+
Py_ssize_t k = 0;
524+
Py_ssize_t j;
525+
for (Py_ssize_t i=0; i < nkwds; i++) {
526+
key = PyTuple_GET_ITEM(kwnames, i);
527+
val = args[nargs + i];
528+
j = positions[i];
529+
if (j != -1) {
530+
stack[tot_nargs + j] = val;
531+
}
532+
else {
533+
PyTuple_SET_ITEM(tot_kwnames, pto_nkwds + k, key);
534+
Py_INCREF(key);
535+
stack[tot_nargs + pto_nkwds + k] = val;
536+
k++;
537+
}
538+
}
539+
DEALLOCATE_STACK(small_positions, positions);
506540
}
507541

542+
/* Call / Maintenance / Return */
508543
ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack, tot_nargs, tot_kwnames);
544+
DEALLOCATE_STACK(small_stack, stack);
509545
Py_DECREF(tot_kwnames);
510-
511-
/* Free stack & Return */
512-
if (stack != small_stack) {
513-
PyMem_Free(stack);
514-
}
515546
return ret;
516547
}
517548
}
@@ -532,7 +563,6 @@ partial_setvectorcall(partialobject *pto)
532563
}
533564
}
534565

535-
536566
// Not converted to argument clinic, because of `*args, **kwargs` arguments.
537567
static PyObject *
538568
partial_call(partialobject *pto, PyObject *args, PyObject *kwargs)

0 commit comments

Comments
 (0)