Implement FOR_ITER_SET

corona10 · Jun 21, 2022 · Jun 22, 2022 · Jun 22, 2022 · Jun 22, 2022 · Jun 22, 2022
commit 5f147c61b5d4f38b963dc461f032f5f0b0d92b75
diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
diff --git a/Include/internal/pycore_set.h b/Include/internal/pycore_set.h
@@ -0,0 +1,24 @@
+#ifndef Py_INTERNAL_SET_H
+#define Py_INTERNAL_SET_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+typedef struct {
+    PyObject_HEAD
+    PySetObject *si_set; /* Set to NULL when iterator is exhausted */
+    Py_ssize_t si_used;
+    Py_ssize_t si_pos;
+    Py_ssize_t len;
+} _PySetIterObject;
+
+int _PySetIter_GetNext(_PySetIterObject *si, PyObject**);
+
+#ifdef __cplusplus
+}
+#endif
+#endif   /* !Py_INTERNAL_SET_H */
diff --git a/Include/opcode.h b/Include/opcode.h
diff --git a/Lib/opcode.py b/Lib/opcode.py
@@ -282,6 +282,7 @@ def jabs_op(name, op):
         "FOR_ITER_ADAPTIVE",
         "FOR_ITER_LIST",
         "FOR_ITER_RANGE",
+        "FOR_ITER_SET",
     ],
     "JUMP_BACKWARD": [
         "JUMP_BACKWARD_QUICK",

@@ -1631,6 +1631,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_range.h \
 		$(srcdir)/Include/internal/pycore_runtime.h \
 		$(srcdir)/Include/internal/pycore_runtime_init.h \
+		$(srcdir)/Include/internal/pycore_set.h \
 		$(srcdir)/Include/internal/pycore_signal.h \
 		$(srcdir)/Include/internal/pycore_sliceobject.h \
 		$(srcdir)/Include/internal/pycore_strhex.h \

@@ -32,12 +32,12 @@
 */
 
 #include "Python.h"
+#include "pycore_set.h"           // struct _PySetIterObject
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include <stddef.h>               // offsetof()
 
 /* Object used as dummy key to fill deleted entries */
 static PyObject _dummy_struct;
-
 #define dummy (&_dummy_struct)
 
 
@@ -730,16 +730,8 @@ frozenset_hash(PyObject *self)
 
 /***** Set iterator type ***********************************************/
 
-typedef struct {
-    PyObject_HEAD
-    PySetObject *si_set; /* Set to NULL when iterator is exhausted */
-    Py_ssize_t si_used;
-    Py_ssize_t si_pos;
-    Py_ssize_t len;
-} setiterobject;
-
 static void
-setiter_dealloc(setiterobject *si)
+setiter_dealloc(_PySetIterObject *si)
 {
     /* bpo-31095: UnTrack is needed before calling any callbacks */
     _PyObject_GC_UNTRACK(si);
@@ -748,30 +740,30 @@ setiter_dealloc(setiterobject *si)
 }
 
 static int
-setiter_traverse(setiterobject *si, visitproc visit, void *arg)
+setiter_traverse(_PySetIterObject *si, visitproc visit, void *arg)
 {
     Py_VISIT(si->si_set);
     return 0;
 }
 
 static PyObject *
-setiter_len(setiterobject *si, PyObject *Py_UNUSED(ignored))
+setiter_len(_PySetIterObject *si, PyObject *Py_UNUSED(ignored))
 {
     Py_ssize_t len = 0;
     if (si->si_set != NULL && si->si_used == si->si_set->used)
         len = si->len;
     return PyLong_FromSsize_t(len);
 }
 
-PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
+static PyObject *setiter_iternext(_PySetIterObject *si);
 
-static PyObject *setiter_iternext(setiterobject *si);
+PyDoc_STRVAR(length_hint_doc, "Private method returning an estimate of len(list(it)).");
 
 static PyObject *
-setiter_reduce(setiterobject *si, PyObject *Py_UNUSED(ignored))
+setiter_reduce(_PySetIterObject *si, PyObject *Py_UNUSED(ignored))
 {
     /* copy the iterator state */
-    setiterobject tmp = *si;
+    _PySetIterObject tmp = *si;
     Py_XINCREF(tmp.si_set);
 
     /* iterate the temporary into a list */
@@ -791,48 +783,60 @@ static PyMethodDef setiter_methods[] = {
     {NULL,              NULL}           /* sentinel */
 };
 
-static PyObject *setiter_iternext(setiterobject *si)
+int _PySetIter_GetNext(_PySetIterObject *si, PyObject **stack)
 {
-    PyObject *key;
+    assert(Py_IS_TYPE((PyObject *)si, &PySetIter_Type));
     Py_ssize_t i, mask;
     setentry *entry;
     PySetObject *so = si->si_set;
 
-    if (so == NULL)
-        return NULL;
+    if (so == NULL) {
+        return 0;
+    }
     assert (PyAnySet_Check(so));
 
     if (si->si_used != so->used) {
         PyErr_SetString(PyExc_RuntimeError,
                         "Set changed size during iteration");
         si->si_used = -1; /* Make this state sticky */
-        return NULL;
+        return -1;
     }
 
     i = si->si_pos;
     assert(i>=0);
     entry = so->table;
     mask = so->mask;
-    while (i <= mask && (entry[i].key == NULL || entry[i].key == dummy))
+    while (i <= mask && (entry[i].key == NULL || entry[i].key == dummy)) {
         i++;
+    }
     si->si_pos = i+1;
-    if (i > mask)
+    if (i > mask) {
         goto fail;
+    }
     si->len--;
-    key = entry[i].key;
-    Py_INCREF(key);
-    return key;
-
+    stack[0] = entry[i].key;
+    Py_INCREF(stack[0]);
+    return 1;
 fail:
     si->si_set = NULL;
     Py_DECREF(so);
-    return NULL;
+    return 0;
+}
+
+static PyObject *setiter_iternext(_PySetIterObject *si) {
+    PyObject *key = NULL;
+    PyObject *stack[1];
+    int err = _PySetIter_GetNext(si, stack);
+    if (err <= 0) {
+        return NULL;
+    }
+    return stack[0];
 }
-static PyObject *setiter_iternext(_PySetIterObject *si) {
-    PyObject *stack[1];
-    int err = _PySetIter_GetNext(si, stack);
-    if (err <= 0) {
-        return NULL;
-    }
-    return stack[0];
-}
+static PyObject *setiter_iternext(_PySetIterObject *si) {
+    PyObject *key;
+    if (_PySetIter_GetNext(si, &key) <= 0) {
+        return NULL;
+    }
+    return key;
+}
-static PyObject *setiter_iternext(_PySetIterObject *si) {
-    PyObject *stack[1];
-    int err = _PySetIter_GetNext(si, stack);
-    if (err <= 0) {
-        return NULL;
-    }
-    return stack[0];
-}
+static PyObject *setiter_iternext(_PySetIterObject *si) {
+    PyObject *key;
+    if (_PySetIter_GetNext(si, &key) <= 0) {
+        return NULL;
+    }
+    return key;
+}
 
 PyTypeObject PySetIter_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     "set_iterator",                             /* tp_name */
-    sizeof(setiterobject),                      /* tp_basicsize */
+    sizeof(_PySetIterObject),                   /* tp_basicsize */
     0,                                          /* tp_itemsize */
     /* methods */
     (destructor)setiter_dealloc,                /* tp_dealloc */
@@ -865,7 +869,7 @@ PyTypeObject PySetIter_Type = {
 static PyObject *
 set_iter(PySetObject *so)
 {
-    setiterobject *si = PyObject_GC_New(setiterobject, &PySetIter_Type);
+    _PySetIterObject *si = PyObject_GC_New(_PySetIterObject, &PySetIter_Type);
     if (si == NULL)
         return NULL;
     Py_INCREF(so);