gh-111178: fix UBSan failures in `Python/hamt.c` (GH-128247)

* fix UBSan failures for `PyHamtObject`
* fix UBSan failures for `PyHamtNode_Array`
* fix UBSan failures for `PyHamtNode_Collision`
* fix UBSan failures for `PyHamtNode_Bitmap`

diff --git a/Python/hamt.c b/Python/hamt.c
index cfd211f45414467a7e683f71ff984c9c7c733c06..ed43a0449d7a016dba48b8907cd76d243d635eeb 100644 (file)
--- a/Python/hamt.c
+++ b/Python/hamt.c
@@ -319,6 +319,8 @@ typedef struct {
     Py_ssize_t a_count;
 } PyHamtNode_Array;
 
+#define _PyHamtNode_Array_CAST(op)      ((PyHamtNode_Array *)(op))
+
 
 typedef struct {
     PyObject_VAR_HEAD
@@ -326,6 +328,8 @@ typedef struct {
     PyObject *c_array[1];
 } PyHamtNode_Collision;
 
+#define _PyHamtNode_Collision_CAST(op)  ((PyHamtNode_Collision *)(op))
+
 
 static PyHamtObject *
 hamt_alloc(void);
@@ -479,6 +483,8 @@ error:
 #endif  /* Py_DEBUG */
 /////////////////////////////////// Bitmap Node
 
+#define _PyHamtNode_Bitmap_CAST(op)     ((PyHamtNode_Bitmap *)(op))
+
 
 static PyHamtNode *
 hamt_node_bitmap_new(Py_ssize_t size)
@@ -1083,30 +1089,27 @@ hamt_node_bitmap_find(PyHamtNode_Bitmap *self,
 }
 
 static int
-hamt_node_bitmap_traverse(PyHamtNode_Bitmap *self, visitproc visit, void *arg)
+hamt_node_bitmap_traverse(PyObject *op, visitproc visit, void *arg)
 {
     /* Bitmap's tp_traverse */
-
-    Py_ssize_t i;
-
-    for (i = Py_SIZE(self); --i >= 0; ) {
+    PyHamtNode_Bitmap *self = _PyHamtNode_Bitmap_CAST(op);
+    for (Py_ssize_t i = Py_SIZE(self); --i >= 0;) {
         Py_VISIT(self->b_array[i]);
     }
-
     return 0;
 }
 
 static void
-hamt_node_bitmap_dealloc(PyHamtNode_Bitmap *self)
+hamt_node_bitmap_dealloc(PyObject *self)
 {
     /* Bitmap's tp_dealloc */
 
-    Py_ssize_t len = Py_SIZE(self);
-    Py_ssize_t i;
+    PyHamtNode_Bitmap *node = _PyHamtNode_Bitmap_CAST(self);
+    Py_ssize_t i, len = Py_SIZE(self);
 
-    if (Py_SIZE(self) == 0) {
+    if (len == 0) {
         /* The empty node is statically allocated. */
-        assert(self == &_Py_SINGLETON(hamt_bitmap_node_empty));
+        assert(node == &_Py_SINGLETON(hamt_bitmap_node_empty));
 #ifdef Py_DEBUG
         _Py_FatalRefcountError("deallocating the empty hamt node bitmap singleton");
 #else
@@ -1120,11 +1123,11 @@ hamt_node_bitmap_dealloc(PyHamtNode_Bitmap *self)
     if (len > 0) {
         i = len;
         while (--i >= 0) {
-            Py_XDECREF(self->b_array[i]);
+            Py_XDECREF(node->b_array[i]);
         }
     }
 
-    Py_TYPE(self)->tp_free((PyObject *)self);
+    Py_TYPE(self)->tp_free(self);
     Py_TRASHCAN_END
 }
 
@@ -1489,38 +1492,30 @@ hamt_node_collision_find(PyHamtNode_Collision *self,
 
 
 static int
-hamt_node_collision_traverse(PyHamtNode_Collision *self,
-                             visitproc visit, void *arg)
+hamt_node_collision_traverse(PyObject *op, visitproc visit, void *arg)
 {
     /* Collision's tp_traverse */
-
-    Py_ssize_t i;
-
-    for (i = Py_SIZE(self); --i >= 0; ) {
+    PyHamtNode_Collision *self = _PyHamtNode_Collision_CAST(op);
+    for (Py_ssize_t i = Py_SIZE(self); --i >= 0; ) {
         Py_VISIT(self->c_array[i]);
     }
-
     return 0;
 }
 
 static void
-hamt_node_collision_dealloc(PyHamtNode_Collision *self)
+hamt_node_collision_dealloc(PyObject *self)
 {
     /* Collision's tp_dealloc */
-
     Py_ssize_t len = Py_SIZE(self);
-
     PyObject_GC_UnTrack(self);
     Py_TRASHCAN_BEGIN(self, hamt_node_collision_dealloc)
-
     if (len > 0) {
-
+        PyHamtNode_Collision *node = _PyHamtNode_Collision_CAST(self);
         while (--len >= 0) {
-            Py_XDECREF(self->c_array[len]);
+            Py_XDECREF(node->c_array[len]);
         }
     }
-
-    Py_TYPE(self)->tp_free((PyObject *)self);
+    Py_TYPE(self)->tp_free(self);
     Py_TRASHCAN_END
 }
 
@@ -1868,35 +1863,27 @@ hamt_node_array_find(PyHamtNode_Array *self,
 }
 
 static int
-hamt_node_array_traverse(PyHamtNode_Array *self,
-                         visitproc visit, void *arg)
+hamt_node_array_traverse(PyObject *op, visitproc visit, void *arg)
 {
     /* Array's tp_traverse */
-
-    Py_ssize_t i;
-
-    for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
+    PyHamtNode_Array *self = _PyHamtNode_Array_CAST(op);
+    for (Py_ssize_t i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
         Py_VISIT(self->a_array[i]);
     }
-
     return 0;
 }
 
 static void
-hamt_node_array_dealloc(PyHamtNode_Array *self)
+hamt_node_array_dealloc(PyObject *self)
 {
     /* Array's tp_dealloc */
-
-    Py_ssize_t i;
-
     PyObject_GC_UnTrack(self);
     Py_TRASHCAN_BEGIN(self, hamt_node_array_dealloc)
-
-    for (i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
-        Py_XDECREF(self->a_array[i]);
+    PyHamtNode_Array *obj = _PyHamtNode_Array_CAST(self);
+    for (Py_ssize_t i = 0; i < HAMT_ARRAY_NODE_SIZE; i++) {
+        Py_XDECREF(obj->a_array[i]);
     }
-
-    Py_TYPE(self)->tp_free((PyObject *)self);
+    Py_TYPE(self)->tp_free(self);
     Py_TRASHCAN_END
 }
 
@@ -2605,6 +2592,8 @@ static PyObject *
 hamt_dump(PyHamtObject *self);
 #endif
 
+#define _PyHamtObject_CAST(op)      ((PyHamtObject *)(op))
+
 
 static PyObject *
 hamt_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
@@ -2613,24 +2602,27 @@ hamt_tp_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 }
 
 static int
-hamt_tp_clear(PyHamtObject *self)
+hamt_tp_clear(PyObject *op)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     Py_CLEAR(self->h_root);
     return 0;
 }
 
 
 static int
-hamt_tp_traverse(PyHamtObject *self, visitproc visit, void *arg)
+hamt_tp_traverse(PyObject *op, visitproc visit, void *arg)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     Py_VISIT(self->h_root);
     return 0;
 }
 
 static void
-hamt_tp_dealloc(PyHamtObject *self)
+hamt_tp_dealloc(PyObject *self)
 {
-    if (self == _empty_hamt) {
+    PyHamtObject *obj = _PyHamtObject_CAST(self);
+    if (obj == _empty_hamt) {
         /* The empty one is statically allocated. */
 #ifdef Py_DEBUG
         _Py_FatalRefcountError("deallocating the empty hamt singleton");
@@ -2640,8 +2632,8 @@ hamt_tp_dealloc(PyHamtObject *self)
     }
 
     PyObject_GC_UnTrack(self);
-    if (self->h_weakreflist != NULL) {
-        PyObject_ClearWeakRefs((PyObject*)self);
+    if (obj->h_weakreflist != NULL) {
+        PyObject_ClearWeakRefs(self);
     }
     (void)hamt_tp_clear(self);
     Py_TYPE(self)->tp_free(self);
@@ -2673,16 +2665,18 @@ hamt_tp_richcompare(PyObject *v, PyObject *w, int op)
 }
 
 static int
-hamt_tp_contains(PyHamtObject *self, PyObject *key)
+hamt_tp_contains(PyObject *op, PyObject *key)
 {
     PyObject *val;
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_Find(self, key, &val);
 }
 
 static PyObject *
-hamt_tp_subscript(PyHamtObject *self, PyObject *key)
+hamt_tp_subscript(PyObject *op, PyObject *key)
 {
     PyObject *val;
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     hamt_find_t res = hamt_find(self, key, &val);
     switch (res) {
         case F_ERROR:
@@ -2698,19 +2692,21 @@ hamt_tp_subscript(PyHamtObject *self, PyObject *key)
 }
 
 static Py_ssize_t
-hamt_tp_len(PyHamtObject *self)
+hamt_tp_len(PyObject *op)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_Len(self);
 }
 
 static PyObject *
-hamt_tp_iter(PyHamtObject *self)
+hamt_tp_iter(PyObject *op)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_NewIterKeys(self);
 }
 
 static PyObject *
-hamt_py_set(PyHamtObject *self, PyObject *args)
+hamt_py_set(PyObject *op, PyObject *args)
 {
     PyObject *key;
     PyObject *val;
@@ -2719,11 +2715,12 @@ hamt_py_set(PyHamtObject *self, PyObject *args)
         return NULL;
     }
 
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return (PyObject *)_PyHamt_Assoc(self, key, val);
 }
 
 static PyObject *
-hamt_py_get(PyHamtObject *self, PyObject *args)
+hamt_py_get(PyObject *op, PyObject *args)
 {
     PyObject *key;
     PyObject *def = NULL;
@@ -2733,6 +2730,7 @@ hamt_py_get(PyHamtObject *self, PyObject *args)
     }
 
     PyObject *val = NULL;
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     hamt_find_t res = hamt_find(self, key, &val);
     switch (res) {
         case F_ERROR:
@@ -2750,67 +2748,63 @@ hamt_py_get(PyHamtObject *self, PyObject *args)
 }
 
 static PyObject *
-hamt_py_delete(PyHamtObject *self, PyObject *key)
+hamt_py_delete(PyObject *op, PyObject *key)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return (PyObject *)_PyHamt_Without(self, key);
 }
 
 static PyObject *
-hamt_py_items(PyHamtObject *self, PyObject *args)
+hamt_py_items(PyObject *op, PyObject *args)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_NewIterItems(self);
 }
 
 static PyObject *
-hamt_py_values(PyHamtObject *self, PyObject *args)
+hamt_py_values(PyObject *op, PyObject *args)
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_NewIterValues(self);
 }
 
 static PyObject *
-hamt_py_keys(PyHamtObject *self, PyObject *Py_UNUSED(args))
+hamt_py_keys(PyObject *op, PyObject *Py_UNUSED(args))
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return _PyHamt_NewIterKeys(self);
 }
 
 #ifdef Py_DEBUG
 static PyObject *
-hamt_py_dump(PyHamtObject *self, PyObject *Py_UNUSED(args))
+hamt_py_dump(PyObject *op, PyObject *Py_UNUSED(args))
 {
+    PyHamtObject *self = _PyHamtObject_CAST(op);
     return hamt_dump(self);
 }
 #endif
 
 
 static PyMethodDef PyHamt_methods[] = {
-    {"set", _PyCFunction_CAST(hamt_py_set), METH_VARARGS, NULL},
-    {"get", _PyCFunction_CAST(hamt_py_get), METH_VARARGS, NULL},
-    {"delete", _PyCFunction_CAST(hamt_py_delete), METH_O, NULL},
-    {"items", _PyCFunction_CAST(hamt_py_items), METH_NOARGS, NULL},
-    {"keys", _PyCFunction_CAST(hamt_py_keys), METH_NOARGS, NULL},
-    {"values", _PyCFunction_CAST(hamt_py_values), METH_NOARGS, NULL},
+    {"set", hamt_py_set, METH_VARARGS, NULL},
+    {"get", hamt_py_get, METH_VARARGS, NULL},
+    {"delete", hamt_py_delete, METH_O, NULL},
+    {"items", hamt_py_items, METH_NOARGS, NULL},
+    {"keys", hamt_py_keys, METH_NOARGS, NULL},
+    {"values", hamt_py_values, METH_NOARGS, NULL},
 #ifdef Py_DEBUG
-    {"__dump__", _PyCFunction_CAST(hamt_py_dump), METH_NOARGS, NULL},
+    {"__dump__", hamt_py_dump, METH_NOARGS, NULL},
 #endif
     {NULL, NULL}
 };
 
 static PySequenceMethods PyHamt_as_sequence = {
-    0,                                /* sq_length */
-    0,                                /* sq_concat */
-    0,                                /* sq_repeat */
-    0,                                /* sq_item */
-    0,                                /* sq_slice */
-    0,                                /* sq_ass_item */
-    0,                                /* sq_ass_slice */
-    (objobjproc)hamt_tp_contains,     /* sq_contains */
-    0,                                /* sq_inplace_concat */
-    0,                                /* sq_inplace_repeat */
+    .sq_contains = hamt_tp_contains,
 };
 
 static PyMappingMethods PyHamt_as_mapping = {
-    (lenfunc)hamt_tp_len,             /* mp_length */
-    (binaryfunc)hamt_tp_subscript,    /* mp_subscript */
+    .mp_length = hamt_tp_len,
+    .mp_subscript = hamt_tp_subscript,
 };
 
 PyTypeObject _PyHamt_Type = {
@@ -2820,13 +2814,13 @@ PyTypeObject _PyHamt_Type = {
     .tp_methods = PyHamt_methods,
     .tp_as_mapping = &PyHamt_as_mapping,
     .tp_as_sequence = &PyHamt_as_sequence,
-    .tp_iter = (getiterfunc)hamt_tp_iter,
-    .tp_dealloc = (destructor)hamt_tp_dealloc,
+    .tp_iter = hamt_tp_iter,
+    .tp_dealloc = hamt_tp_dealloc,
     .tp_getattro = PyObject_GenericGetAttr,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
     .tp_richcompare = hamt_tp_richcompare,
-    .tp_traverse = (traverseproc)hamt_tp_traverse,
-    .tp_clear = (inquiry)hamt_tp_clear,
+    .tp_traverse = hamt_tp_traverse,
+    .tp_clear = hamt_tp_clear,
     .tp_new = hamt_tp_new,
     .tp_weaklistoffset = offsetof(PyHamtObject, h_weakreflist),
     .tp_hash = PyObject_HashNotImplemented,
@@ -2841,10 +2835,10 @@ PyTypeObject _PyHamt_ArrayNode_Type = {
     "hamt_array_node",
     sizeof(PyHamtNode_Array),
     0,
-    .tp_dealloc = (destructor)hamt_node_array_dealloc,
+    .tp_dealloc = hamt_node_array_dealloc,
     .tp_getattro = PyObject_GenericGetAttr,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
-    .tp_traverse = (traverseproc)hamt_node_array_traverse,
+    .tp_traverse = hamt_node_array_traverse,
     .tp_free = PyObject_GC_Del,
     .tp_hash = PyObject_HashNotImplemented,
 };
@@ -2854,10 +2848,10 @@ PyTypeObject _PyHamt_BitmapNode_Type = {
     "hamt_bitmap_node",
     sizeof(PyHamtNode_Bitmap) - sizeof(PyObject *),
     sizeof(PyObject *),
-    .tp_dealloc = (destructor)hamt_node_bitmap_dealloc,
+    .tp_dealloc = hamt_node_bitmap_dealloc,
     .tp_getattro = PyObject_GenericGetAttr,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
-    .tp_traverse = (traverseproc)hamt_node_bitmap_traverse,
+    .tp_traverse = hamt_node_bitmap_traverse,
     .tp_free = PyObject_GC_Del,
     .tp_hash = PyObject_HashNotImplemented,
 };
@@ -2867,10 +2861,10 @@ PyTypeObject _PyHamt_CollisionNode_Type = {
     "hamt_collision_node",
     sizeof(PyHamtNode_Collision) - sizeof(PyObject *),
     sizeof(PyObject *),
-    .tp_dealloc = (destructor)hamt_node_collision_dealloc,
+    .tp_dealloc = hamt_node_collision_dealloc,
     .tp_getattro = PyObject_GenericGetAttr,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
-    .tp_traverse = (traverseproc)hamt_node_collision_traverse,
+    .tp_traverse = hamt_node_collision_traverse,
     .tp_free = PyObject_GC_Del,
     .tp_hash = PyObject_HashNotImplemented,
 };