basic/strv: introduce strv_extend_strv_consume()

3 files changed, 104 insertions, 5 deletions

diff --git a/src/basic/strv.c b/src/basic/strv.c
index 0c3733a02b..542883ea3f 100644
--- a/src/basic/strv.c
+++ b/src/basic/strv.c
@@ -229,20 +229,18 @@ char** strv_new_internal(const char *x, ...) {
 
 int strv_extend_strv(char ***a, char * const *b, bool filter_duplicates) {
         size_t p, q, i = 0;
-        char **t;
 
         assert(a);
 
-        if (strv_isempty(b))
+        q = strv_length(b);
+        if (q == 0)
                 return 0;
 
         p = strv_length(*a);
-        q = strv_length(b);
-
         if (p >= SIZE_MAX - q)
                 return -ENOMEM;
 
-        t = reallocarray(*a, GREEDY_ALLOC_ROUND_UP(p + q + 1), sizeof(char *));
+        char **t = reallocarray(*a, GREEDY_ALLOC_ROUND_UP(p + q + 1), sizeof(char *));
         if (!t)
                 return -ENOMEM;
 
@@ -271,9 +269,67 @@ rollback:
         return -ENOMEM;
 }
 
+int strv_extend_strv_consume(char ***a, char **b, bool filter_duplicates) {
+        _cleanup_strv_free_ char **b_consume = b;
+        size_t p, q, i;
+
+        assert(a);
+
+        q = strv_length(b);
+        if (q == 0)
+                return 0;
+
+        p = strv_length(*a);
+        if (p == 0) {
+                strv_free_and_replace(*a, b_consume);
+
+                if (filter_duplicates)
+                        strv_uniq(*a);
+
+                return strv_length(*a);
+        }
+
+        if (p >= SIZE_MAX - q)
+                return -ENOMEM;
+
+        char **t = reallocarray(*a, GREEDY_ALLOC_ROUND_UP(p + q + 1), sizeof(char *));
+        if (!t)
+                return -ENOMEM;
+
+        t[p] = NULL;
+        *a = t;
+
+        if (!filter_duplicates) {
+                *mempcpy_typesafe(t + p, b, q) = NULL;
+                i = q;
+        } else {
+                i = 0;
+
+                STRV_FOREACH(s, b) {
+                        if (strv_contains(t, *s)) {
+                                free(*s);
+                                continue;
+                        }
+
+                        t[p+i] = *s;
+
+                        i++;
+                        t[p+i] = NULL;
+                }
+        }
+
+        assert(i <= q);
+
+        b_consume = mfree(b_consume);
+
+        return (int) i;
+}
+
 int strv_extend_strv_biconcat(char ***a, const char *prefix, const char* const *b, const char *suffix) {
         int r;
 
+        assert(a);
+
         STRV_FOREACH(s, b) {
                 char *v;
 
diff --git a/src/basic/strv.h b/src/basic/strv.h
index 1ef3ea5faf..b9580f7634 100644
--- a/src/basic/strv.h
+++ b/src/basic/strv.h
@@ -44,10 +44,13 @@ int strv_copy_unless_empty(char * const *l, char ***ret);
 size_t strv_length(char * const *l) _pure_;
 
 int strv_extend_strv(char ***a, char * const *b, bool filter_duplicates);
+int strv_extend_strv_consume(char ***a, char **b, bool filter_duplicates);
+
 int strv_extend_strv_biconcat(char ***a, const char *prefix, const char* const *b, const char *suffix);
 static inline int strv_extend_strv_concat(char ***a, const char* const *b, const char *suffix) {
         return strv_extend_strv_biconcat(a, NULL, b, suffix);
 }
+
 int strv_prepend(char ***l, const char *value);
 
 /* _with_size() are lower-level functions where the size can be provided externally,
diff --git a/src/test/test-strv.c b/src/test/test-strv.c
index f7b6ee2bb2..a35a2b10dc 100644
--- a/src/test/test-strv.c
+++ b/src/test/test-strv.c
@@ -660,6 +660,46 @@ TEST(strv_extend_strv) {
         assert_se(strv_length(n) == 4);
 }
 
+TEST(strv_extend_strv_consume) {
+        _cleanup_strv_free_ char **a = NULL, **b = NULL, **c = NULL, **n = NULL;
+        const char *s1, *s2, *s3;
+
+        ASSERT_NOT_NULL(a = strv_new("abc", "def", "ghi"));
+        ASSERT_NOT_NULL(b = strv_new("jkl", "mno", "abc", "pqr"));
+
+        s1 = b[0];
+        s2 = b[1];
+        s3 = b[3];
+
+        ASSERT_EQ(strv_extend_strv_consume(&a, TAKE_PTR(b), true), 3);
+
+        assert_se(s1 == a[3]);
+        assert_se(s2 == a[4]);
+        assert_se(s3 == a[5]);
+
+        ASSERT_STREQ(a[0], "abc");
+        ASSERT_STREQ(a[1], "def");
+        ASSERT_STREQ(a[2], "ghi");
+        ASSERT_STREQ(a[3], "jkl");
+        ASSERT_STREQ(a[4], "mno");
+        ASSERT_STREQ(a[5], "pqr");
+        ASSERT_EQ(strv_length(a), (size_t) 6);
+
+        ASSERT_NOT_NULL(c = strv_new("jkl", "mno"));
+
+        s1 = c[0];
+        s2 = c[1];
+
+        ASSERT_EQ(strv_extend_strv_consume(&n, TAKE_PTR(c), false), 2);
+
+        assert_se(s1 == n[0]);
+        assert_se(s2 == n[1]);
+
+        ASSERT_STREQ(n[0], "jkl");
+        ASSERT_STREQ(n[1], "mno");
+        ASSERT_EQ(strv_length(n), (size_t) 2);
+}
+
 TEST(strv_extend_with_size) {
         _cleanup_strv_free_ char **a = NULL;
         size_t n = SIZE_MAX;