Add set/hashmap helpers for non-trivial freeing and use where straighforward

A macro is needed because otherwise we couldn't ensure type safety.
Some simple tests are included.
No functional change intended.

6 files changed, 86 insertions, 16 deletions

diff --git a/src/basic/hashmap.h b/src/basic/hashmap.h
index c1089652d..57809f50b 100644
--- a/src/basic/hashmap.h
+++ b/src/basic/hashmap.h
@@ -328,6 +328,29 @@ static inline void *ordered_hashmap_first(OrderedHashmap *h) {
         return internal_hashmap_first(HASHMAP_BASE(h));
 }
 
+#define hashmap_clear_with_destructor(_s, _f)                   \
+        ({                                                      \
+                void *_item;                                    \
+                while ((_item = hashmap_steal_first(_s)))       \
+                        _f(_item);                              \
+        })
+#define hashmap_free_with_destructor(_s, _f)                    \
+        ({                                                      \
+                hashmap_clear_with_destructor(_s, _f);          \
+                hashmap_free(_s);                               \
+        })
+#define ordered_hashmap_clear_with_destructor(_s, _f)                   \
+        ({                                                              \
+                void *_item;                                            \
+                while ((_item = ordered_hashmap_steal_first(_s)))       \
+                        _f(_item);                                      \
+        })
+#define ordered_hashmap_free_with_destructor(_s, _f)                    \
+        ({                                                              \
+                ordered_hashmap_clear_with_destructor(_s, _f);          \
+                ordered_hashmap_free(_s);                               \
+        })
+
 /* no hashmap_next */
 void *ordered_hashmap_next(OrderedHashmap *h, const void *key);
 
diff --git a/src/basic/set.h b/src/basic/set.h
index 08b7c4a72..b424f3423 100644
--- a/src/basic/set.h
+++ b/src/basic/set.h
@@ -109,6 +109,18 @@ static inline void *set_steal_first(Set *s) {
         return internal_hashmap_steal_first(HASHMAP_BASE(s));
 }
 
+#define set_clear_with_destructor(_s, _f)               \
+        ({                                              \
+                void *_item;                            \
+                while ((_item = set_steal_first(_s)))   \
+                        _f(_item);                      \
+        })
+#define set_free_with_destructor(_s, _f)                \
+        ({                                              \
+                set_clear_with_destructor(_s, _f);      \
+                set_free(_s);                           \
+        })
+
 /* no set_steal_first_key */
 /* no set_first_key */
 
diff --git a/src/libelogind/sd-bus/bus-track.c b/src/libelogind/sd-bus/bus-track.c
index 4acaf2479..9d9025275 100644
--- a/src/libelogind/sd-bus/bus-track.c
+++ b/src/libelogind/sd-bus/bus-track.c
@@ -183,8 +183,6 @@ _public_ sd_bus_track* sd_bus_track_ref(sd_bus_track *track) {
 }
 
 _public_ sd_bus_track* sd_bus_track_unref(sd_bus_track *track) {
-        struct track_item *i;
-
         if (!track)
                 return NULL;
 
@@ -195,14 +193,11 @@ _public_ sd_bus_track* sd_bus_track_unref(sd_bus_track *track) {
                 return NULL;
         }
 
-        while ((i = hashmap_steal_first(track->names)))
-                track_item_free(i);
-
         if (track->in_list)
                 LIST_REMOVE(tracks, track->bus->tracks, track);
 
         bus_track_remove_from_queue(track);
-        hashmap_free(track->names);
+        hashmap_free_with_destructor(track->names, track_item_free);
         sd_bus_unref(track->bus);
         return mfree(track);
 }
@@ -428,8 +423,6 @@ void bus_track_dispatch(sd_bus_track *track) {
 }
 
 void bus_track_close(sd_bus_track *track) {
-        struct track_item *i;
-
         assert(track);
 
         /* Called whenever our bus connected is closed. If so, and our track object is non-empty, dispatch it
@@ -447,8 +440,7 @@ void bus_track_close(sd_bus_track *track) {
                 return;
 
         /* Let's flush out all names */
-        while ((i = hashmap_steal_first(track->names)))
-                track_item_free(i);
+        hashmap_clear_with_destructor(track->names, track_item_free);
 
         /* Invoke handler */
         if (track->handler)
diff --git a/src/shared/bus-util.c b/src/shared/bus-util.c
index 0fb28c24d..81e1db34f 100644
--- a/src/shared/bus-util.c
+++ b/src/shared/bus-util.c
@@ -555,12 +555,7 @@ int bus_verify_polkit_async(
 
 void bus_verify_polkit_async_registry_free(Hashmap *registry) {
 #if ENABLE_POLKIT
-        AsyncPolkitQuery *q;
-
-        while ((q = hashmap_steal_first(registry)))
-                async_polkit_query_free(q);
-
-        hashmap_free(registry);
+        hashmap_free_with_destructor(registry, async_polkit_query_free);
 #endif
 }
 
diff --git a/src/test/test-hashmap.c b/src/test/test-hashmap.c
index 83cea360e..2615c98eb 100644
--- a/src/test/test-hashmap.c
+++ b/src/test/test-hashmap.c
@@ -37,6 +37,29 @@ static void test_ordered_hashmap_next(void) {
         assert_se(!ordered_hashmap_next(m, INT_TO_PTR(3)));
 }
 
+typedef struct Item {
+        int seen;
+} Item;
+static void item_seen(Item *item) {
+        item->seen++;
+}
+
+static void test_hashmap_free_with_destructor(void) {
+        Hashmap *m;
+        struct Item items[4] = {};
+        unsigned i;
+
+        assert_se(m = hashmap_new(NULL));
+        for (i = 0; i < ELEMENTSOF(items) - 1; i++)
+                assert_se(hashmap_put(m, INT_TO_PTR(i), items + i) == 1);
+
+        m = hashmap_free_with_destructor(m, item_seen);
+        assert_se(items[0].seen == 1);
+        assert_se(items[1].seen == 1);
+        assert_se(items[2].seen == 1);
+        assert_se(items[3].seen == 0);
+}
+
 static void test_uint64_compare_func(void) {
         const uint64_t a = 0x100, b = 0x101;
 
@@ -61,6 +84,7 @@ int main(int argc, const char *argv[]) {
         test_ordered_hashmap_funcs();
 
         test_ordered_hashmap_next();
+        test_hashmap_free_with_destructor();
         test_uint64_compare_func();
         test_trivial_compare_func();
         test_string_compare_func();
diff --git a/src/test/test-set.c b/src/test/test-set.c
index 3fab350cf..21540fa1d 100644
--- a/src/test/test-set.c
+++ b/src/test/test-set.c
@@ -39,6 +39,29 @@ static void test_set_steal_first(void) {
         assert_se(set_isempty(m));
 }
 
+typedef struct Item {
+        int seen;
+} Item;
+static void item_seen(Item *item) {
+        item->seen++;
+}
+
+static void test_set_free_with_destructor(void) {
+        Set *m;
+        struct Item items[4] = {};
+        unsigned i;
+
+        assert_se(m = set_new(NULL));
+        for (i = 0; i < ELEMENTSOF(items) - 1; i++)
+                assert_se(set_put(m, items + i) == 1);
+
+        m = set_free_with_destructor(m, item_seen);
+        assert_se(items[0].seen == 1);
+        assert_se(items[1].seen == 1);
+        assert_se(items[2].seen == 1);
+        assert_se(items[3].seen == 0);
+}
+
 static void test_set_put(void) {
         _cleanup_set_free_ Set *m = NULL;
 
@@ -101,6 +124,7 @@ static void test_set_make(void) {
 
 int main(int argc, const char *argv[]) {
         test_set_steal_first();
+        test_set_free_with_destructor();
         test_set_put();
         test_set_make();