bpo-40609: Add destroy functions to _Py_hashtable (GH-20062)

Add key_destroy_func and value_destroy_func parameters to
_Py_hashtable_new_full().

marshal.c and _tracemalloc.c use these destroy functions.

Author: vstinner

Include/internal/pycore_hashtable.h

@@ -34,21 +34,21 @@ typedef struct {
     /* data (data_size bytes) follows */
 } _Py_hashtable_entry_t;
 
-#define _Py_HASHTABLE_ENTRY_PDATA(TABLE, ENTRY) \
+#define _Py_HASHTABLE_ENTRY_PDATA(ENTRY) \
         ((const void *)((char *)(ENTRY) \
                         + sizeof(_Py_hashtable_entry_t)))
 
 #define _Py_HASHTABLE_ENTRY_READ_DATA(TABLE, ENTRY, DATA) \
     do { \
         assert(sizeof(DATA) == (TABLE)->data_size); \
-        memcpy(&(DATA), _Py_HASHTABLE_ENTRY_PDATA(TABLE, (ENTRY)), \
+        memcpy(&(DATA), _Py_HASHTABLE_ENTRY_PDATA((ENTRY)), \
                   sizeof(DATA)); \
     } while (0)
 
 #define _Py_HASHTABLE_ENTRY_WRITE_DATA(TABLE, ENTRY, DATA) \
     do { \
         assert(sizeof(DATA) == (TABLE)->data_size); \
-        memcpy((void *)_Py_HASHTABLE_ENTRY_PDATA((TABLE), (ENTRY)), \
+        memcpy((void *)_Py_HASHTABLE_ENTRY_PDATA(ENTRY), \
                   &(DATA), sizeof(DATA)); \
     } while (0)
 
@@ -61,6 +61,9 @@ typedef struct _Py_hashtable_t _Py_hashtable_t;
 
 typedef Py_uhash_t (*_Py_hashtable_hash_func) (const void *key);
 typedef int (*_Py_hashtable_compare_func) (const void *key1, const void *key2);
+typedef void (*_Py_hashtable_destroy_func) (void *key);
+typedef void (*_Py_hashtable_value_destroy_func) (_Py_hashtable_t *ht,
+                                                  _Py_hashtable_entry_t *entry);
 typedef _Py_hashtable_entry_t* (*_Py_hashtable_get_entry_func)(_Py_hashtable_t *ht,
                                                                const void *key);
 typedef int (*_Py_hashtable_get_func) (_Py_hashtable_t *ht,
@@ -86,6 +89,8 @@ struct _Py_hashtable_t {
     _Py_hashtable_get_entry_func get_entry_func;
     _Py_hashtable_hash_func hash_func;
     _Py_hashtable_compare_func compare_func;
+    _Py_hashtable_destroy_func key_destroy_func;
+    _Py_hashtable_value_destroy_func value_destroy_func;
     _Py_hashtable_allocator_t alloc;
 };
 
@@ -107,6 +112,8 @@ PyAPI_FUNC(_Py_hashtable_t *) _Py_hashtable_new_full(
     size_t init_size,
     _Py_hashtable_hash_func hash_func,
     _Py_hashtable_compare_func compare_func,
+    _Py_hashtable_destroy_func key_destroy_func,
+    _Py_hashtable_value_destroy_func value_destroy_func,
     _Py_hashtable_allocator_t *allocator);
 
 PyAPI_FUNC(void) _Py_hashtable_destroy(_Py_hashtable_t *ht);

Modules/_tracemalloc.c

@@ -238,12 +238,13 @@ hashtable_hash_uint(const void *key_raw)
 static _Py_hashtable_t *
 hashtable_new(size_t data_size,
               _Py_hashtable_hash_func hash_func,
-              _Py_hashtable_compare_func compare_func)
+              _Py_hashtable_compare_func compare_func,
+              _Py_hashtable_value_destroy_func value_destroy_fun)
 {
     _Py_hashtable_allocator_t hashtable_alloc = {malloc, free};
     return _Py_hashtable_new_full(data_size, 0,
                                   hash_func, compare_func,
-                                  &hashtable_alloc);
+                                  NULL, value_destroy_fun, &hashtable_alloc);
 }
 
 
@@ -471,35 +472,34 @@ tracemalloc_create_traces_table(void)
 {
     return hashtable_new(sizeof(trace_t),
                          _Py_hashtable_hash_ptr,
-                         _Py_hashtable_compare_direct);
+                         _Py_hashtable_compare_direct,
+                         NULL);
 }
 
 
-static _Py_hashtable_t*
-tracemalloc_create_domains_table(void)
+static void
+tracemalloc_destroy_domain_table(_Py_hashtable_t *domains,
+                               _Py_hashtable_entry_t *entry)
 {
-    return hashtable_new(sizeof(_Py_hashtable_t *),
-                         hashtable_hash_uint,
-                         _Py_hashtable_compare_direct);
+    _Py_hashtable_t *traces;
+    _Py_HASHTABLE_ENTRY_READ_DATA(domains, entry, traces);
+    _Py_hashtable_destroy(traces);
 }
 
 
-static int
-tracemalloc_destroy_domains_cb(_Py_hashtable_t *domains,
-                               _Py_hashtable_entry_t *entry,
-                               void *user_data)
+static _Py_hashtable_t*
+tracemalloc_create_domains_table(void)
 {
-    _Py_hashtable_t *traces;
-    _Py_HASHTABLE_ENTRY_READ_DATA(domains, entry, traces);
-    _Py_hashtable_destroy(traces);
-    return 0;
+    return hashtable_new(sizeof(_Py_hashtable_t *),
+                         hashtable_hash_uint,
+                         _Py_hashtable_compare_direct,
+                         tracemalloc_destroy_domain_table);
 }
 
 
 static void
 tracemalloc_destroy_domains(_Py_hashtable_t *domains)
 {
-    _Py_hashtable_foreach(domains, tracemalloc_destroy_domains_cb, NULL);
     _Py_hashtable_destroy(domains);
 }
 
@@ -924,11 +924,13 @@ tracemalloc_init(void)
 
     tracemalloc_filenames = hashtable_new(0,
                                           hashtable_hash_pyobject,
-                                          hashtable_compare_unicode);
+                                          hashtable_compare_unicode,
+                                          NULL);
 
     tracemalloc_tracebacks = hashtable_new(0,
                                            hashtable_hash_traceback,
-                                           hashtable_compare_traceback);
+                                           hashtable_compare_traceback,
+                                          NULL);
 
     tracemalloc_traces = tracemalloc_create_traces_table();
     tracemalloc_domains = tracemalloc_create_domains_table();
@@ -1285,15 +1287,13 @@ tracemalloc_get_traces_domain(_Py_hashtable_t *domains,
 }
 
 
-static int
+static void
 tracemalloc_pyobject_decref_cb(_Py_hashtable_t *tracebacks,
-                               _Py_hashtable_entry_t *entry,
-                               void *user_data)
+                               _Py_hashtable_entry_t *entry)
 {
     PyObject *obj;
     _Py_HASHTABLE_ENTRY_READ_DATA(tracebacks, entry, obj);
     Py_DECREF(obj);
-    return 0;
 }
 
 
@@ -1329,7 +1329,8 @@ _tracemalloc__get_traces_impl(PyObject *module)
        of (filename, lineno) tuples */
     get_traces.tracebacks = hashtable_new(sizeof(PyObject *),
                                           _Py_hashtable_hash_ptr,
-                                          _Py_hashtable_compare_direct);
+                                          _Py_hashtable_compare_direct,
+                                          tracemalloc_pyobject_decref_cb);
     if (get_traces.tracebacks == NULL) {
         goto no_memory;
     }
@@ -1381,8 +1382,6 @@ _tracemalloc__get_traces_impl(PyObject *module)
 
 finally:
     if (get_traces.tracebacks != NULL) {
-        _Py_hashtable_foreach(get_traces.tracebacks,
-                              tracemalloc_pyobject_decref_cb, NULL);
         _Py_hashtable_destroy(get_traces.tracebacks);
     }
     if (get_traces.traces != NULL) {

Python/hashtable.c

@@ -64,14 +64,14 @@
 #define ENTRY_READ_PDATA(TABLE, ENTRY, DATA_SIZE, PDATA) \
     do { \
         assert((DATA_SIZE) == (TABLE)->data_size); \
-        memcpy((PDATA), _Py_HASHTABLE_ENTRY_PDATA(TABLE, (ENTRY)), \
+        memcpy((PDATA), _Py_HASHTABLE_ENTRY_PDATA(ENTRY), \
                   (DATA_SIZE)); \
     } while (0)
 
 #define ENTRY_WRITE_PDATA(TABLE, ENTRY, DATA_SIZE, PDATA) \
     do { \
         assert((DATA_SIZE) == (TABLE)->data_size); \
-        memcpy((void *)_Py_HASHTABLE_ENTRY_PDATA((TABLE), (ENTRY)), \
+        memcpy((void *)_Py_HASHTABLE_ENTRY_PDATA(ENTRY), \
                   (PDATA), (DATA_SIZE)); \
     } while (0)
 
@@ -432,6 +432,8 @@ _Py_hashtable_t *
 _Py_hashtable_new_full(size_t data_size, size_t init_size,
                        _Py_hashtable_hash_func hash_func,
                        _Py_hashtable_compare_func compare_func,
+                       _Py_hashtable_destroy_func key_destroy_func,
+                       _Py_hashtable_value_destroy_func value_destroy_func,
                        _Py_hashtable_allocator_t *allocator)
 {
     _Py_hashtable_t *ht;
@@ -466,6 +468,8 @@ _Py_hashtable_new_full(size_t data_size, size_t init_size,
     ht->get_entry_func = _Py_hashtable_get_entry_generic;
     ht->hash_func = hash_func;
     ht->compare_func = compare_func;
+    ht->key_destroy_func = key_destroy_func;
+    ht->value_destroy_func = value_destroy_func;
     ht->alloc = alloc;
     if (ht->hash_func == _Py_hashtable_hash_ptr
         && ht->compare_func == _Py_hashtable_compare_direct)
@@ -484,7 +488,7 @@ _Py_hashtable_new(size_t data_size,
 {
     return _Py_hashtable_new_full(data_size, HASHTABLE_MIN_SIZE,
                                   hash_func, compare_func,
-                                  NULL);
+                                  NULL, NULL, NULL);
 }
 
 
@@ -506,16 +510,27 @@ _Py_hashtable_clear(_Py_hashtable_t *ht)
 }
 
 
+static void
+_Py_hashtable_destroy_entry(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry)
+{
+    if (ht->key_destroy_func) {
+        ht->key_destroy_func(entry->key);
+    }
+    if (ht->value_destroy_func) {
+        ht->value_destroy_func(ht, entry);
+    }
+    ht->alloc.free(entry);
+}
+
+
 void
 _Py_hashtable_destroy(_Py_hashtable_t *ht)
 {
-    size_t i;
-
-    for (i = 0; i < ht->num_buckets; i++) {
-        _Py_slist_item_t *entry = ht->buckets[i].head;
+    for (size_t i = 0; i < ht->num_buckets; i++) {
+        _Py_hashtable_entry_t *entry = TABLE_HEAD(ht, i);
         while (entry) {
-            _Py_slist_item_t *entry_next = entry->next;
-            ht->alloc.free(entry);
+            _Py_hashtable_entry_t *entry_next = ENTRY_NEXT(entry);
+            _Py_hashtable_destroy_entry(ht, entry);
             entry = entry_next;
         }
     }
@@ -537,6 +552,8 @@ _Py_hashtable_copy(_Py_hashtable_t *src)
     dst = _Py_hashtable_new_full(data_size, src->num_buckets,
                                  src->hash_func,
                                  src->compare_func,
+                                 src->key_destroy_func,
+                                 src->value_destroy_func,
                                  &src->alloc);
     if (dst == NULL)
         return NULL;
@@ -545,7 +562,7 @@ _Py_hashtable_copy(_Py_hashtable_t *src)
         entry = TABLE_HEAD(src, bucket);
         for (; entry; entry = ENTRY_NEXT(entry)) {
             const void *key = entry->key;
-            const void *pdata = _Py_HASHTABLE_ENTRY_PDATA(src, entry);
+            const void *pdata = _Py_HASHTABLE_ENTRY_PDATA(entry);
             err = _Py_hashtable_set(dst, key, data_size, pdata);
             if (err) {
                 _Py_hashtable_destroy(dst);

Python/marshal.c

@@ -545,13 +545,21 @@ w_complex_object(PyObject *v, char flag, WFILE *p)
     }
 }
 
+static void
+w_decref_entry(void *key)
+{
+    PyObject *entry_key = (PyObject *)key;
+    Py_XDECREF(entry_key);
+}
+
 static int
 w_init_refs(WFILE *wf, int version)
 {
     if (version >= 3) {
-        wf->hashtable = _Py_hashtable_new(sizeof(int),
-                                          _Py_hashtable_hash_ptr,
-                                          _Py_hashtable_compare_direct);
+        wf->hashtable = _Py_hashtable_new_full(sizeof(int), 0,
+                                               _Py_hashtable_hash_ptr,
+                                               _Py_hashtable_compare_direct,
+                                               w_decref_entry, NULL, NULL);
         if (wf->hashtable == NULL) {
             PyErr_NoMemory();
             return -1;
@@ -560,20 +568,10 @@ w_init_refs(WFILE *wf, int version)
     return 0;
 }
 
-static int
-w_decref_entry(_Py_hashtable_t *ht, _Py_hashtable_entry_t *entry,
-               void *Py_UNUSED(data))
-{
-    PyObject *entry_key = (PyObject *)entry->key;
-    Py_XDECREF(entry_key);
-    return 0;
-}
-
 static void
 w_clear_refs(WFILE *wf)
 {
     if (wf->hashtable != NULL) {
-        _Py_hashtable_foreach(wf->hashtable, w_decref_entry, NULL);
         _Py_hashtable_destroy(wf->hashtable);
     }
 }