bpo-40609: Add destroy functions to _Py_hashtable (pythonGH-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

Modules/_tracemalloc.c

@@ -279,12 +279,13 @@ hashtable_compare_pointer_t(const void *key1, const void *key2)
 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);
 }
 
 
@@ -519,59 +520,41 @@ static int
 tracemalloc_use_domain_cb(_Py_hashtable_t *old_traces,
                            _Py_hashtable_entry_t *entry, void *user_data)
 {
-    uintptr_t ptr;
-    pointer_t key;
-    _Py_hashtable_t *new_traces = (_Py_hashtable_t *)user_data;
-    const void *pdata = _Py_HASHTABLE_ENTRY_PDATA(old_traces, entry);
+    return hashtable_new(sizeof(trace_t),
+                         _Py_hashtable_hash_ptr,
+                         _Py_hashtable_compare_direct,
+                         NULL);
+}
 
     _Py_HASHTABLE_ENTRY_READ_KEY(old_traces, entry, ptr);
     key.ptr = ptr;
     key.domain = DEFAULT_DOMAIN;
 
-    return _Py_hashtable_set(new_traces,
-                             sizeof(key), &key,
-                             old_traces->data_size, pdata);
+static void
+tracemalloc_destroy_domain_table(_Py_hashtable_t *domains,
+                               _Py_hashtable_entry_t *entry)
+{
+    _Py_hashtable_t *traces;
+    _Py_HASHTABLE_ENTRY_READ_DATA(domains, entry, traces);
+    _Py_hashtable_destroy(traces);
 }
 
 
-/* Convert tracemalloc_traces from compact key (uintptr_t) to pointer_t key.
- * Return 0 on success, -1 on error. */
-static int
-tracemalloc_use_domain(void)
+static _Py_hashtable_t*
+tracemalloc_create_domains_table(void)
 {
-    _Py_hashtable_t *new_traces = NULL;
-
-    assert(!tracemalloc_config.use_domain);
-
-    new_traces = hashtable_new(sizeof(pointer_t),
-                               sizeof(trace_t),
-                               hashtable_hash_pointer_t,
-                               hashtable_compare_pointer_t);
-    if (new_traces == NULL) {
-        return -1;
-    }
-
-    if (_Py_hashtable_foreach(tracemalloc_traces, tracemalloc_use_domain_cb,
-                              new_traces) < 0)
-    {
-        _Py_hashtable_destroy(new_traces);
-        return -1;
-    }
-
-    _Py_hashtable_destroy(tracemalloc_traces);
-    tracemalloc_traces = new_traces;
-
-    tracemalloc_config.use_domain = 1;
-
-    return 0;
+    return hashtable_new(sizeof(_Py_hashtable_t *),
+                         hashtable_hash_uint,
+                         _Py_hashtable_compare_direct,
+                         tracemalloc_destroy_domain_table);
 }
 
 
 static void
 tracemalloc_remove_trace(_PyTraceMalloc_domain_t domain, uintptr_t ptr)
 {
-    trace_t trace;
-    int removed;
+    _Py_hashtable_destroy(domains);
+}
 
     assert(tracemalloc_config.tracing);
 
@@ -1002,11 +985,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);
 
     if (tracemalloc_config.use_domain) {
         tracemalloc_traces = hashtable_new(sizeof(pointer_t),
@@ -1324,14 +1309,30 @@ tracemalloc_get_traces_fill(_Py_hashtable_t *traces, _Py_hashtable_entry_t *entr
 
 
 static int
+tracemalloc_get_traces_domain(_Py_hashtable_t *domains,
+                              _Py_hashtable_entry_t *entry,
+                              void *user_data)
+{
+    get_traces_t *get_traces = user_data;
+
+    unsigned int domain = (unsigned int)FROM_PTR(entry->key);
+    _Py_hashtable_t *traces;
+    _Py_HASHTABLE_ENTRY_READ_DATA(domains, entry, traces);
+
+    get_traces->domain = domain;
+    return _Py_hashtable_foreach(traces,
+                                 tracemalloc_get_traces_fill,
+                                 get_traces);
+}
+
+
+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;
 }
 
 
@@ -1364,7 +1365,8 @@ py_tracemalloc_get_traces(PyObject *self, PyObject *obj)
     get_traces.tracebacks = hashtable_new(sizeof(traceback_t *),
                                           sizeof(PyObject *),
                                           _Py_hashtable_hash_ptr,
-                                          _Py_hashtable_compare_direct);
+                                          _Py_hashtable_compare_direct,
+                                          tracemalloc_pyobject_decref_cb);
     if (get_traces.tracebacks == NULL) {
         PyErr_NoMemory();
         goto error;
@@ -1393,8 +1395,6 @@ py_tracemalloc_get_traces(PyObject *self, PyObject *obj)
 
 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) {

Modules/hashtable.h

@@ -29,21 +29,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)
 
@@ -56,6 +56,9 @@ typedef struct _Py_hashtable_t _Py_hashtable_t;
 
 typedef Py_uhash_t (*_Py_hashtable_hash_func) (const void *pkey);
 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,
@@ -81,6 +84,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;
 };
 
@@ -102,6 +107,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);

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_Modules/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

@@ -570,13 +570,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;
@@ -585,20 +593,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);
     }
 }