omlish 0.0.0.dev493__py3-none-any.whl → 0.0.0.dev506__py3-none-any.whl

omlish/typedvalues/_collection.cc

@@ -0,0 +1,500 @@
+// @omlish-cext
+// @omlish-llm-author "claude-sonnet-4-5-20250929"
+#define PY_SSIZE_T_CLEAN
+#include "Python.h"
+#include "structmember.h"
+
+#include <memory>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+//
+
+#define _MODULE_NAME "_collection"
+#define _PACKAGE_NAME "omlish.typedvalues"
+#define _MODULE_FULL_NAME _PACKAGE_NAME "." _MODULE_NAME
+
+typedef struct collection_state {
+    PyObject *typed_value_type;
+    PyObject *unique_typed_value_type;
+    PyObject *duplicate_error_type;
+} collection_state;
+
+static inline collection_state * get_collection_state(PyObject *module)
+{
+    void *state = PyModule_GetState(module);
+    assert(state != NULL);
+    return (collection_state *)state;
+}
+
+//
+
+// Helper struct to track unique typed values during processing
+struct UniqueInfo {
+    PyObject *unique_tv_cls;             // The unique class (borrowed from map key)
+    PyObject *tv;                        // The typed value (borrowed from args)
+    size_t idx;                          // Index in unique_lst when added
+};
+
+using TmpItem = std::variant<PyObject*, std::unique_ptr<UniqueInfo>>;
+
+PyDoc_STRVAR(init_typed_values_collection_doc,
+"init_typed_values_collection(*tvs, override=False, check_type=None)\n"
+"\n"
+"Initialize a typed values collection.\n"
+"\n"
+"Returns a tuple of (tuple, dict, dict2) where:\n"
+"- tuple: ordered sequence of typed values\n"
+"- dict: mapping of type -> typed value or tuple of typed values\n"
+"- dict2: extended dict including instance types for unique values");
+
+static PyObject * init_typed_values_collection(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    collection_state *state = get_collection_state(module);
+
+    // Parse arguments
+    bool override = false;
+    PyObject *check_type = nullptr;
+
+    // Handle keyword arguments
+    Py_ssize_t nkwargs = kwnames ? PyTuple_GET_SIZE(kwnames) : 0;
+
+    for (Py_ssize_t i = 0; i < nkwargs; i++) {
+        PyObject *key = PyTuple_GET_ITEM(kwnames, i);
+        PyObject *value = args[nargs + i];
+        const char *key_str = PyUnicode_AsUTF8(key);
+
+        if (key_str == nullptr) {
+            return nullptr;
+        }
+
+        if (strcmp(key_str, "override") == 0) {
+            int override_result = PyObject_IsTrue(value);
+            if (override_result == -1) {
+                return nullptr;
+            }
+            override = override_result;
+        } else if (strcmp(key_str, "check_type") == 0) {
+            check_type = value;
+        } else {
+            PyErr_Format(PyExc_TypeError, "unexpected keyword argument: %s", key_str);
+            return nullptr;
+        }
+    }
+
+    // If no positional arguments, return empty tuple and dicts
+    if (nargs == 0) {
+        PyObject *empty_tuple = PyTuple_New(0);
+        if (empty_tuple == nullptr) {
+            return nullptr;
+        }
+
+        PyObject *empty_dict = PyDict_New();
+        if (empty_dict == nullptr) {
+            Py_DECREF(empty_tuple);
+            return nullptr;
+        }
+
+        PyObject *empty_dict2 = PyDict_New();
+        if (empty_dict2 == nullptr) {
+            Py_DECREF(empty_tuple);
+            Py_DECREF(empty_dict);
+            return nullptr;
+        }
+
+        PyObject *result = PyTuple_Pack(3, empty_tuple, empty_dict, empty_dict2);
+        Py_DECREF(empty_tuple);
+        Py_DECREF(empty_dict);
+        Py_DECREF(empty_dict2);
+        return result;
+    }
+
+    // Temporary storage
+    std::vector<TmpItem> tmp_lst;
+    tmp_lst.reserve(nargs);
+
+    // Map from unique type to list of typed values
+    std::unordered_map<PyObject*, std::vector<PyObject*>> unique_dct;
+
+    struct UniqueKeysCleaner {
+        std::unordered_map<PyObject*, std::vector<PyObject*>> &map;
+        ~UniqueKeysCleaner() { for (auto &pair : map) Py_DECREF(pair.first); }
+    } keys_cleaner{unique_dct};
+
+    // Process each typed value
+    for (Py_ssize_t i = 0; i < nargs; i++) {
+        PyObject *tv = args[i];
+
+        // Check type if requested
+        if (check_type != nullptr && check_type != Py_None) {
+            int result = PyObject_IsInstance(tv, check_type);
+            if (result == -1) {
+                return nullptr;
+            }
+            if (!result) {
+                PyErr_SetObject(PyExc_TypeError, tv);
+                return nullptr;
+            }
+        }
+
+        // Check if it's a UniqueTypedValue
+        int is_unique = PyObject_IsInstance(tv, state->unique_typed_value_type);
+        if (is_unique == -1) {
+            return nullptr;
+        }
+
+        if (is_unique) {
+            // Get the _unique_typed_value_cls attribute
+            PyObject *unique_tv_cls = PyObject_GetAttrString(tv, "_unique_typed_value_cls");
+            if (unique_tv_cls == nullptr) {
+                return nullptr;
+            }
+
+            // Check for duplicates if not override
+            if (!override) {
+                auto it = unique_dct.find(unique_tv_cls);
+                if (it != unique_dct.end() && !it->second.empty()) {
+                    // Duplicate found - raise error
+                    PyObject *old_tv = it->second[0];
+
+                    // Create DuplicateUniqueTypedValueError
+                    PyObject *error = PyObject_CallFunction(
+                        state->duplicate_error_type,
+                        "OOO",
+                        unique_tv_cls,
+                        tv,
+                        old_tv
+                    );
+
+                    Py_DECREF(unique_tv_cls);
+
+                    if (error == nullptr) {
+                        return nullptr;
+                    }
+
+                    PyErr_SetObject(state->duplicate_error_type, error);
+                    Py_DECREF(error);
+                    return nullptr;
+                }
+            }
+
+            // Add to unique_dct
+            // Use insert to avoid reference leak: if key already exists, we need to decref the new reference
+            auto insertion = unique_dct.insert({unique_tv_cls, std::vector<PyObject*>()});
+            if (!insertion.second) {
+                // Key already existed, decref the newly acquired reference
+                Py_DECREF(unique_tv_cls);
+            }
+            std::vector<PyObject*> &unique_lst = insertion.first->second;
+            unique_lst.push_back(tv);
+
+            // Add to tmp_lst using make_unique
+            // unique_ptr automatically handles deletion, variant handles moves
+            tmp_lst.emplace_back(std::make_unique<UniqueInfo>(UniqueInfo{
+                unique_tv_cls,  // Borrowed from map
+                tv,
+                unique_lst.size()
+            }));
+
+        } else {
+            // Check if it's a TypedValue
+            int is_typed_value = PyObject_IsInstance(tv, state->typed_value_type);
+            if (is_typed_value == -1) {
+                return nullptr;
+            }
+
+            if (!is_typed_value) {
+                PyErr_SetObject(PyExc_TypeError, tv);
+                return nullptr;
+            }
+
+            // Non-unique typed value
+            tmp_lst.emplace_back(tv);
+        }
+    }
+
+    // Build output structures
+    PyObject *lst = PyList_New(0);
+    if (lst == nullptr) {
+        return nullptr;
+    }
+
+    PyObject *tmp_dct = PyDict_New();
+    if (tmp_dct == nullptr) {
+        Py_DECREF(lst);
+        return nullptr;
+    }
+
+    // Process tmp_lst to build output list and tmp_dct
+    for (auto &item : tmp_lst) {
+        // Check if item holds unique_ptr<UniqueInfo>
+        if (auto *uniq_ptr_wrapper = std::get_if<std::unique_ptr<UniqueInfo>>(&item)) {
+            UniqueInfo *info = uniq_ptr_wrapper->get();
+
+            // Look up the vector now (after all insertions are done, no more rehashing)
+            auto it = unique_dct.find(info->unique_tv_cls);
+            assert(it != unique_dct.end());
+
+            // Last-in-wins: only include if this is the last one
+            if (info->idx == it->second.size()) {
+                // Add to output list
+                if (PyList_Append(lst, info->tv) == -1) {
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+
+                // Add to tmp_dct (scalar value for unique types)
+                if (PyDict_SetItem(tmp_dct, info->unique_tv_cls, info->tv) == -1) {
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+            }
+        } else {
+            // Must be PyObject*
+            PyObject *tv = std::get<PyObject*>(item);
+
+            // Add to output list
+            if (PyList_Append(lst, tv) == -1) {
+                Py_DECREF(lst);
+                Py_DECREF(tmp_dct);
+                return nullptr;
+            }
+
+            // Add to tmp_dct (accumulating list for non-unique types)
+            PyObject *tv_type = (PyObject *)Py_TYPE(tv);
+            PyObject *existing = PyDict_GetItem(tmp_dct, tv_type);  // Borrowed reference
+
+            if (existing == nullptr) {
+                // Check if an error occurred during lookup
+                if (PyErr_Occurred()) {
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+
+                // Create new list
+                PyObject *new_list = PyList_New(0);
+                if (new_list == nullptr) {
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+
+                if (PyList_Append(new_list, tv) == -1) {
+                    Py_DECREF(new_list);
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+
+                if (PyDict_SetItem(tmp_dct, tv_type, new_list) == -1) {
+                    Py_DECREF(new_list);
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+
+                Py_DECREF(new_list);
+            } else {
+                // Append to existing list
+                if (PyList_Append(existing, tv) == -1) {
+                    Py_DECREF(lst);
+                    Py_DECREF(tmp_dct);
+                    return nullptr;
+                }
+            }
+        }
+    }
+
+    // Convert list to tuple
+    PyObject *result_tuple = PyList_AsTuple(lst);
+    Py_DECREF(lst);
+    if (result_tuple == nullptr) {
+        Py_DECREF(tmp_dct);
+        return nullptr;
+    }
+
+    // Build dct: convert lists to tuples
+    PyObject *dct = PyDict_New();
+    if (dct == nullptr) {
+        Py_DECREF(result_tuple);
+        Py_DECREF(tmp_dct);
+        return nullptr;
+    }
+
+    PyObject *key, *value;
+    Py_ssize_t pos = 0;
+
+    while (PyDict_Next(tmp_dct, &pos, &key, &value)) {
+        if (PyList_Check(value)) {
+            PyObject *tuple_value = PyList_AsTuple(value);
+            if (tuple_value == nullptr) {
+                Py_DECREF(result_tuple);
+                Py_DECREF(tmp_dct);
+                Py_DECREF(dct);
+                return nullptr;
+            }
+
+            if (PyDict_SetItem(dct, key, tuple_value) == -1) {
+                Py_DECREF(tuple_value);
+                Py_DECREF(result_tuple);
+                Py_DECREF(tmp_dct);
+                Py_DECREF(dct);
+                return nullptr;
+            }
+
+            Py_DECREF(tuple_value);
+        } else {
+            // Scalar value (unique type)
+            if (PyDict_SetItem(dct, key, value) == -1) {
+                Py_DECREF(result_tuple);
+                Py_DECREF(tmp_dct);
+                Py_DECREF(dct);
+                return nullptr;
+            }
+        }
+    }
+
+    Py_DECREF(tmp_dct);
+
+    // Build dct2: dct + unique values keyed by instance type
+    PyObject *dct2 = PyDict_Copy(dct);
+    if (dct2 == nullptr) {
+        Py_DECREF(result_tuple);
+        Py_DECREF(dct);
+        return nullptr;
+    }
+
+    // Add unique values by their instance type
+    pos = 0;
+    while (PyDict_Next(dct, &pos, &key, &value)) {
+        // Check if value is a UniqueTypedValue (not a tuple)
+        if (!PyTuple_Check(value)) {
+            int is_unique = PyObject_IsInstance(value, state->unique_typed_value_type);
+            if (is_unique == -1) {
+                Py_DECREF(result_tuple);
+                Py_DECREF(dct);
+                Py_DECREF(dct2);
+                return nullptr;
+            }
+
+            if (is_unique) {
+                PyObject *value_type = (PyObject *)Py_TYPE(value);
+                if (PyDict_SetItem(dct2, value_type, value) == -1) {
+                    Py_DECREF(result_tuple);
+                    Py_DECREF(dct);
+                    Py_DECREF(dct2);
+                    return nullptr;
+                }
+            }
+        }
+    }
+
+    // Build final result
+    PyObject *result = PyTuple_Pack(3, result_tuple, dct, dct2);
+    Py_DECREF(result_tuple);
+    Py_DECREF(dct);
+    Py_DECREF(dct2);
+
+    return result;
+}
+
+//
+
+PyDoc_STRVAR(collection_doc, "Native C++ implementations for omlish.typedvalues.collection");
+
+static int collection_exec(PyObject *module)
+{
+    collection_state *state = get_collection_state(module);
+
+    // Import TypedValue
+    PyObject *values_module = PyImport_ImportModule("omlish.typedvalues.values");
+    if (values_module == nullptr) {
+        return -1;
+    }
+
+    state->typed_value_type = PyObject_GetAttrString(values_module, "TypedValue");
+    if (state->typed_value_type == nullptr) {
+        Py_DECREF(values_module);
+        return -1;
+    }
+
+    state->unique_typed_value_type = PyObject_GetAttrString(values_module, "UniqueTypedValue");
+    Py_DECREF(values_module);
+    if (state->unique_typed_value_type == nullptr) {
+        return -1;
+    }
+
+    // Import DuplicateUniqueTypedValueError
+    PyObject *collection_module = PyImport_ImportModule("omlish.typedvalues.collection");
+    if (collection_module == nullptr) {
+        return -1;
+    }
+
+    state->duplicate_error_type = PyObject_GetAttrString(collection_module, "DuplicateUniqueTypedValueError");
+    Py_DECREF(collection_module);
+    if (state->duplicate_error_type == nullptr) {
+        return -1;
+    }
+
+    return 0;
+}
+
+static int collection_traverse(PyObject *module, visitproc visit, void *arg)
+{
+    collection_state *state = get_collection_state(module);
+    Py_VISIT(state->typed_value_type);
+    Py_VISIT(state->unique_typed_value_type);
+    Py_VISIT(state->duplicate_error_type);
+    return 0;
+}
+
+static int collection_clear(PyObject *module)
+{
+    collection_state *state = get_collection_state(module);
+    Py_CLEAR(state->typed_value_type);
+    Py_CLEAR(state->unique_typed_value_type);
+    Py_CLEAR(state->duplicate_error_type);
+    return 0;
+}
+
+static void collection_free(void *module)
+{
+    collection_clear((PyObject *)module);
+}
+
+static PyMethodDef collection_methods[] = {
+    {"init_typed_values_collection", (PyCFunction)(void(*)(void))init_typed_values_collection, METH_FASTCALL | METH_KEYWORDS, init_typed_values_collection_doc},
+    {NULL, NULL, 0, NULL}
+};
+
+static struct PyModuleDef_Slot collection_slots[] = {
+    {Py_mod_exec, (void *) collection_exec},
+    {Py_mod_gil, Py_MOD_GIL_NOT_USED},
+    {Py_mod_multiple_interpreters, Py_MOD_MULTIPLE_INTERPRETERS_SUPPORTED},
+    {0, NULL}
+};
+
+static struct PyModuleDef collection_module = {
+    .m_base = PyModuleDef_HEAD_INIT,
+    .m_name = _MODULE_NAME,
+    .m_doc = collection_doc,
+    .m_size = sizeof(collection_state),
+    .m_methods = collection_methods,
+    .m_slots = collection_slots,
+    .m_traverse = collection_traverse,
+    .m_clear = collection_clear,
+    .m_free = collection_free,
+};
+
+extern "C" {
+
+PyMODINIT_FUNC PyInit__collection(void)
+{
+    return PyModuleDef_Init(&collection_module);
+}
+
+}