Add 'len', 'iter' support to IntArrayList

Add 'refcnt' method to Unsafe

pyfastutil/ints.pyi

@@ -1,4 +1,4 @@
-from typing import overload, Iterable, SupportsIndex, final
+from typing import overload, Iterable, SupportsIndex, final, Iterator, _T_co
 
 
 @final
@@ -133,3 +133,46 @@ class IntArrayList(list[int]):
             [1, 2]
         """
         pass
+
+
+@final
+class IntArrayListIter(Iterator[int]):
+    """
+    Iterator for IntArrayList.
+
+    This class provides an iterator over an `IntArrayList`, allowing you to iterate over the elements
+    of the list one by one.
+
+    Note:
+        This class cannot be directly instantiated by users. It is designed to be used internally by
+        `IntArrayList` and can only be obtained by calling the `__iter__` method on an `IntArrayList` object.
+
+    Example:
+        >>> int_list = IntArrayList([1, 2, 3])
+        >>> iter_obj = iter(int_list)  # This returns an IntArrayListIter instance
+        >>> next(iter_obj)
+        1
+        >>> next(iter_obj)
+        2
+        >>> next(iter_obj)
+        3
+        >>> next(iter_obj)  # Raises StopIteration
+
+    Raises:
+        TypeError: If attempted to be instantiated directly.
+    """
+
+    def __next__(self) -> int:
+        """
+        Return the next element in the iteration.
+
+        This method retrieves the next element from the `IntArrayList` that this iterator is associated with.
+        If all elements have been iterated over, it raises a `StopIteration` exception.
+
+        Returns:
+            int: The next integer in the `IntArrayList`.
+
+        Raises:
+            StopIteration: If there are no more elements to iterate over.
+        """
+        pass

pyfastutil/src/PyFastUtil.cpp

@@ -5,6 +5,7 @@
 #include "PyFastUtil.h"
 #include "utils/simd/BitonicSort.h"
 #include "ints/IntArrayList.h"
+#include "ints/IntArrayListIter.h"
 #include "unsafe/Unsafe.h"
 
 static struct PyModuleDef pyfastutilModule = {
@@ -26,6 +27,8 @@ PyMODINIT_FUNC PyInit___pyfastutil() {
         return nullptr;
 
     PyModule_AddObject(parent, "IntArrayList", PyInit_IntArrayList());
+    PyModule_AddObject(parent, "IntArrayListIter", PyInit_IntArrayListIter());
+
     PyModule_AddObject(parent, "Unsafe", PyInit_Unsafe());
 
     return parent;

pyfastutil/src/ints/IntArrayList.cpp

@@ -6,18 +6,14 @@
 #include <vector>
 #include <algorithm>
 #include <stdexcept>
+#include <iostream>
 #include "utils/PythonUtils.h"
-#include "utils/ParseUtils.h"
 #include "utils/TimSort.h"
 #include "utils/simd/BitonicSort.h"
 #include "utils/memory/AlignedAllocator.h"
+#include "ints/IntArrayListIter.h"
 
 extern "C" {
-typedef struct IntArrayList {
-    PyObject_HEAD;
-    // we use 64 bytes memory aligned to support faster SIMD, suggestion by ChatGPT.
-    std::vector<int, AlignedAllocator<int, 64>> vector;
-} IntArrayList;
 
 static PyTypeObject IntArrayListType = {
         PyVarObject_HEAD_INIT(&PyType_Type, 0)
@@ -454,35 +450,54 @@ static PyObject *IntArrayList_sort(PyObject *pySelf, PyObject *args, PyObject *k
     Py_RETURN_NONE;
 }
 
+Py_ssize_t IntArrayList_len(PyObject *pySelf) {
+    auto *self = reinterpret_cast<IntArrayList *>(pySelf);
+
+    return static_cast<Py_ssize_t>(self->vector.size());
+}
+
+
+PyObject *IntArrayList_iter(PyObject *pySelf) {
+    auto *self = reinterpret_cast<IntArrayList *>(pySelf);
+
+    return reinterpret_cast<PyObject *>(IntArrayListIter_create(self));
+}
+
 
 static PyMethodDef IntArrayList_methods[] = {
         {"from_range", (PyCFunction) IntArrayList_from_range, METH_VARARGS | METH_STATIC},
-        {"resize",     (PyCFunction) IntArrayList_resize,     METH_O},
-        {"copy",       (PyCFunction) IntArrayList_copy,       METH_NOARGS},
-        {"append",     (PyCFunction) IntArrayList_append,     METH_O},
-        {"extend",     (PyCFunction) IntArrayList_extend,     METH_O},
-        {"pop",        (PyCFunction) IntArrayList_pop,        METH_VARARGS},
-        {"index",      (PyCFunction) IntArrayList_index,      METH_VARARGS},
-        {"count",      (PyCFunction) IntArrayList_count,      METH_O},
-        {"insert",     (PyCFunction) IntArrayList_insert,     METH_VARARGS},
-        {"remove",     (PyCFunction) IntArrayList_remove,     METH_O},
-        {"sort",       (PyCFunction) IntArrayList_sort,       METH_VARARGS | METH_KEYWORDS},
+        {"resize", (PyCFunction) IntArrayList_resize, METH_O},
+        {"copy", (PyCFunction) IntArrayList_copy, METH_NOARGS},
+        {"append", (PyCFunction) IntArrayList_append, METH_O},
+        {"extend", (PyCFunction) IntArrayList_extend, METH_O},
+        {"pop", (PyCFunction) IntArrayList_pop, METH_VARARGS},
+        {"index", (PyCFunction) IntArrayList_index, METH_VARARGS},
+        {"count", (PyCFunction) IntArrayList_count, METH_O},
+        {"insert", (PyCFunction) IntArrayList_insert, METH_VARARGS},
+        {"remove", (PyCFunction) IntArrayList_remove, METH_O},
+        {"sort", (PyCFunction) IntArrayList_sort, METH_VARARGS | METH_KEYWORDS},
         {nullptr}
 };
 
 static struct PyModuleDef IntArrayList_module = {
         PyModuleDef_HEAD_INIT,
         "__pyfastutil.IntArrayList",
-        "A IntArrayList_module that creates an IntArrayList",
+        "An IntArrayList_module that creates an IntArrayList",
         -1,
         nullptr, nullptr, nullptr, nullptr, nullptr
 };
 
+static PySequenceMethods IntArrayList_asSequence = {
+        .sq_length = (lenfunc) IntArrayList_len
+};
+
 void initializeIntArrayListType(PyTypeObject &type) {
     type.tp_name = "IntArrayList";
     type.tp_basicsize = sizeof(IntArrayList);
     type.tp_itemsize = 0;
     type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
+    type.tp_as_sequence = &IntArrayList_asSequence;
+    type.tp_iter = IntArrayList_iter;
     type.tp_methods = IntArrayList_methods;
     type.tp_init = (initproc) IntArrayList_init;
     type.tp_new = PyType_GenericNew;

pyfastutil/src/ints/IntArrayList.h

@@ -6,6 +6,16 @@
 #define PYFASTUTIL_INTARRAYLIST_H
 
 #include "utils/PythonPCH.h"
+#include "utils/memory/AlignedAllocator.h"
+#include <vector>
+
+extern "C" {
+typedef struct IntArrayList {
+    PyObject_HEAD;
+    // we use 64 bytes memory aligned to support faster SIMD, suggestion by ChatGPT.
+    std::vector<int, AlignedAllocator<int, 64>> vector;
+} IntArrayList;
+}
 
 PyMODINIT_FUNC PyInit_IntArrayList();
 

pyfastutil/src/ints/IntArrayListIter.cpp

@@ -0,0 +1,91 @@
+//
+// Created by xia__mc on 2024/11/13.
+//
+
+#include "IntArrayListIter.h"
+#include "utils/PythonUtils.h"
+
+extern "C" {
+
+static PyTypeObject IntArrayListIterType = {
+        PyVarObject_HEAD_INIT(&PyType_Type, 0)
+};
+
+IntArrayListIter *IntArrayListIter_create(IntArrayList *list) {
+    auto *instance = Py_CreateObjNoInit<IntArrayListIter>(IntArrayListIterType);
+
+    Py_INCREF(list);
+    instance->container = list;
+    instance->index = 0;
+
+    return instance;
+}
+
+static void IntArrayListIter_dealloc(IntArrayListIter *self) {
+    SAFE_DECREF(self->container);
+    Py_TYPE(self)->tp_free((PyObject *) self);
+}
+
+static PyObject *IntArrayListIter_next(PyObject *pySelf) {
+    auto *self = reinterpret_cast<IntArrayListIter *>(pySelf);
+
+    if (self->index >= self->container->vector.size()) {
+        PyErr_SetNone(PyExc_StopIteration);
+        return nullptr;
+    }
+
+    auto element = self->container->vector[self->index];
+    self->index++;
+    return PyLong_FromLong(static_cast<long>(element));
+}
+
+static PyObject *IntArrayListIter_iter(PyObject *pySelf) {
+    return pySelf;
+}
+
+static PyMethodDef IntArrayListIter_methods[] = {
+        {nullptr}
+};
+
+static struct PyModuleDef IntArrayListIter_module = {
+        PyModuleDef_HEAD_INIT,
+        "__pyfastutil.IntArrayListIter",
+        "An IntArrayListIter_module that creates an IntArrayList",
+        -1,
+        nullptr, nullptr, nullptr, nullptr, nullptr
+};
+
+void initializeIntArrayListIterType(PyTypeObject &type) {
+    type.tp_name = "IntArrayListIter";
+    type.tp_basicsize = sizeof(IntArrayListIter);
+    type.tp_itemsize = 0;
+    type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
+    type.tp_iter = IntArrayListIter_iter;
+    type.tp_iternext = IntArrayListIter_next;
+    type.tp_methods = IntArrayListIter_methods;
+    type.tp_dealloc = (destructor) IntArrayListIter_dealloc;
+    type.tp_alloc = PyType_GenericAlloc;
+    type.tp_free = PyObject_Del;
+}
+
+#pragma clang diagnostic push
+#pragma ide diagnostic ignored "OCUnusedGlobalDeclarationInspection"
+PyMODINIT_FUNC PyInit_IntArrayListIter() {
+    initializeIntArrayListIterType(IntArrayListIterType);
+
+    PyObject *object = PyModule_Create(&IntArrayListIter_module);
+    if (object == nullptr)
+        return nullptr;
+
+    Py_INCREF(&IntArrayListIterType);
+    if (PyModule_AddObject(object, "IntArrayList", (PyObject *) &IntArrayListIterType) < 0) {
+        Py_DECREF(&IntArrayListIterType);
+        Py_DECREF(object);
+        return nullptr;
+    }
+
+    return object;
+}
+#pragma clang diagnostic pop
+
+}

pyfastutil/src/ints/IntArrayListIter.h

@@ -0,0 +1,24 @@
+//
+// Created by xia__mc on 2024/11/13.
+//
+
+#ifndef PYFASTUTIL_INTARRAYLISTITER_H
+#define PYFASTUTIL_INTARRAYLISTITER_H
+
+#include "utils/PythonPCH.h"
+#include "IntArrayList.h"
+
+extern "C" {
+typedef struct IntArrayListIter {
+    PyObject_HEAD;
+    IntArrayList *container;
+    size_t index;
+} IntArrayListIter;
+
+IntArrayListIter *IntArrayListIter_create(IntArrayList *list);
+
+}
+
+PyMODINIT_FUNC PyInit_IntArrayListIter();
+
+#endif //PYFASTUTIL_INTARRAYLISTITER_H

pyfastutil/src/unsafe/Unsafe.cpp

@@ -6,9 +6,6 @@
 #include "mutex"
 
 extern "C" {
-typedef struct Unsafe {
-    PyObject_HEAD;
-} Unsafe;
 
 static PyTypeObject UnsafeType = {
         PyVarObject_HEAD_INIT(&PyType_Type, 0)
@@ -159,6 +156,10 @@ static PyObject *Unsafe_decref([[maybe_unused]] PyObject *self, PyObject *pyObje
     Py_RETURN_NONE;
 }
 
+static PyObject *Unsafe_refcnt([[maybe_unused]] PyObject *self, PyObject *pyObject) {
+    return PyLong_FromSsize_t(Py_REFCNT(pyObject));
+}
+
 static PyObject *Unsafe_fputs([[maybe_unused]] PyObject *pySelf, PyObject *pyStr) {
     if (!PyUnicode_Check(pyStr)) {
         PyErr_SetString(PyExc_TypeError, "Argument must be a string.");
@@ -217,6 +218,7 @@ static PyMethodDef Unsafe_methods[] = {
         {"memcpy",      (PyCFunction) Unsafe_memcpy,     METH_VARARGS, nullptr},
         {"incref",      (PyCFunction) Unsafe_incref,     METH_O,       nullptr},
         {"decref",      (PyCFunction) Unsafe_decref,     METH_O,       nullptr},
+        {"refcnt",      (PyCFunction) Unsafe_refcnt,     METH_O,       nullptr},
         {"fputs",       (PyCFunction) Unsafe_fputs,      METH_O,       nullptr},
         {"fflush",      (PyCFunction) Unsafe_fflush,     METH_NOARGS,  nullptr},
         {"fgets",       (PyCFunction) Unsafe_fgets,      METH_O,       nullptr},

pyfastutil/src/unsafe/Unsafe.h

@@ -7,6 +7,12 @@
 
 #include "utils/PythonPCH.h"
 
+extern "C" {
+typedef struct Unsafe {
+    PyObject_HEAD;
+} Unsafe;
+}
+
 PyMODINIT_FUNC PyInit_Unsafe();
 
 #endif //PYFASTUTIL_UNSAFE_H

pyfastutil/src/utils/PythonPCH.h

@@ -11,4 +11,6 @@
 
 #pragma warning(pop)
 
+#include "Compat.h"
+
 #endif //PYFASTUTIL_PYTHONPCH_H