set: introduce set_put_strndup()

Note, if `n != SIZE_MAX`, we cannot check the existence of the specified
string in the set without duplicating the string. And, set_consume() also
checks the existence of the string. Hence, it is not necessary to call
set_contains() if `n != SIZE_MAX`.

3 files changed, 44 insertions, 8 deletions

diff --git a/src/basic/hashmap.c b/src/basic/hashmap.c
index e33d6c3073..62380b0e4a 100644
--- a/src/basic/hashmap.c
+++ b/src/basic/hashmap.c
@@ -1842,7 +1842,7 @@ int _hashmap_put_strdup_full(Hashmap **h, const struct hash_ops *hash_ops, const
         return r;
 }
 
-int _set_put_strdup_full(Set **s, const struct hash_ops *hash_ops, const char *p  HASHMAP_DEBUG_PARAMS) {
+int _set_put_strndup_full(Set **s, const struct hash_ops *hash_ops, const char *p, size_t n  HASHMAP_DEBUG_PARAMS) {
         char *c;
         int r;
 
@@ -1853,10 +1853,13 @@ int _set_put_strdup_full(Set **s, const struct hash_ops *hash_ops, const char *p
         if (r < 0)
                 return r;
 
-        if (set_contains(*s, (char*) p))
-                return 0;
+        if (n == SIZE_MAX) {
+                if (set_contains(*s, (char*) p))
+                        return 0;
 
-        c = strdup(p);
+                c = strdup(p);
+        } else
+                c = strndup(p, n);
         if (!c)
                 return -ENOMEM;
 
@@ -1869,7 +1872,7 @@ int _set_put_strdupv_full(Set **s, const struct hash_ops *hash_ops, char **l  HA
         assert(s);
 
         STRV_FOREACH(i, l) {
-                r = _set_put_strdup_full(s, hash_ops, *i  HASHMAP_DEBUG_PASS_ARGS);
+                r = _set_put_strndup_full(s, hash_ops, *i, SIZE_MAX  HASHMAP_DEBUG_PASS_ARGS);
                 if (r < 0)
                         return r;
 
diff --git a/src/basic/set.h b/src/basic/set.h
index 243a747e98..52cf63e2dd 100644
--- a/src/basic/set.h
+++ b/src/basic/set.h
@@ -127,9 +127,12 @@ int _set_ensure_consume(Set **s, const struct hash_ops *hash_ops, void *key  HAS
 
 int set_consume(Set *s, void *value);
 
-int _set_put_strdup_full(Set **s, const struct hash_ops *hash_ops, const char *p  HASHMAP_DEBUG_PARAMS);
-#define set_put_strdup_full(s, hash_ops, p) _set_put_strdup_full(s, hash_ops, p  HASHMAP_DEBUG_SRC_ARGS)
-#define set_put_strdup(s, p) set_put_strdup_full(s, &string_hash_ops_free, p)
+int _set_put_strndup_full(Set **s, const struct hash_ops *hash_ops, const char *p, size_t n  HASHMAP_DEBUG_PARAMS);
+#define set_put_strndup_full(s, hash_ops, p, n) _set_put_strndup_full(s, hash_ops, p, n  HASHMAP_DEBUG_SRC_ARGS)
+#define set_put_strdup_full(s, hash_ops, p) set_put_strndup_full(s, hash_ops, p, SIZE_MAX)
+#define set_put_strndup(s, p, n) set_put_strndup_full(s, &string_hash_ops_free, p, n)
+#define set_put_strdup(s, p) set_put_strndup(s, p, SIZE_MAX)
+
 int _set_put_strdupv_full(Set **s, const struct hash_ops *hash_ops, char **l  HASHMAP_DEBUG_PARAMS);
 #define set_put_strdupv_full(s, hash_ops, l) _set_put_strdupv_full(s, hash_ops, l  HASHMAP_DEBUG_SRC_ARGS)
 #define set_put_strdupv(s, l) set_put_strdupv_full(s, &string_hash_ops_free, l)
diff --git a/src/test/test-set.c b/src/test/test-set.c
index 5c5c35f3a2..0fc9dffe23 100644
--- a/src/test/test-set.c
+++ b/src/test/test-set.c
@@ -90,6 +90,27 @@ TEST(set_put) {
         assert_se(strv_length(t) == 3);
 }
 
+TEST(set_put_strndup) {
+        _cleanup_set_free_ Set *m = NULL;
+
+        assert_se(set_put_strndup(&m, "12345", 0) == 1);
+        assert_se(set_put_strndup(&m, "12345", 1) == 1);
+        assert_se(set_put_strndup(&m, "12345", 2) == 1);
+        assert_se(set_put_strndup(&m, "12345", 3) == 1);
+        assert_se(set_put_strndup(&m, "12345", 4) == 1);
+        assert_se(set_put_strndup(&m, "12345", 5) == 1);
+        assert_se(set_put_strndup(&m, "12345", 6) == 0);
+
+        assert_se(set_contains(m, ""));
+        assert_se(set_contains(m, "1"));
+        assert_se(set_contains(m, "12"));
+        assert_se(set_contains(m, "123"));
+        assert_se(set_contains(m, "1234"));
+        assert_se(set_contains(m, "12345"));
+
+        assert_se(set_size(m) == 6);
+}
+
 TEST(set_put_strdup) {
         _cleanup_set_free_ Set *m = NULL;
 
@@ -98,6 +119,10 @@ TEST(set_put_strdup) {
         assert_se(set_put_strdup(&m, "bbb") == 1);
         assert_se(set_put_strdup(&m, "bbb") == 0);
         assert_se(set_put_strdup(&m, "aaa") == 0);
+
+        assert_se(set_contains(m, "aaa"));
+        assert_se(set_contains(m, "bbb"));
+
         assert_se(set_size(m) == 2);
 }
 
@@ -106,6 +131,11 @@ TEST(set_put_strdupv) {
 
         assert_se(set_put_strdupv(&m, STRV_MAKE("aaa", "aaa", "bbb", "bbb", "aaa")) == 2);
         assert_se(set_put_strdupv(&m, STRV_MAKE("aaa", "aaa", "bbb", "bbb", "ccc")) == 1);
+
+        assert_se(set_contains(m, "aaa"));
+        assert_se(set_contains(m, "bbb"));
+        assert_se(set_contains(m, "ccc"));
+
         assert_se(set_size(m) == 3);
 }