From a27c8bfe6ef3a00c13b1a604f5693837123d33e2 Mon Sep 17 00:00:00 2001
From: Rongsong Shen <rshen@shenrs.eu>
Date: Fri, 26 Dec 2025 16:23:04 +0800
Subject: [PATCH] enhance zarray library function

---
 src/include/zarray.h            |   2 +
 src/lib/testsuite/test_zarray.c |  21 +++++++
 src/lib/zarray.c                | 106 +++++++++++++++++++++++---------
 3 files changed, 100 insertions(+), 29 deletions(-)

diff --git a/src/include/zarray.h b/src/include/zarray.h
index 916719ce..40422993 100644
--- a/src/include/zarray.h
+++ b/src/include/zarray.h
@@ -26,6 +26,8 @@ extern int   zarray_iter_ex(zarray_t*, zarray_iter_skip_fn_t, xlist_t*, void*);
 extern int   zarray_iter(zarray_t*, zarray_iter_fn_t, void*);
 extern int   zarray_del(zarray_t*, int);
 extern int   zarray_put(zarray_t*, int, void*);
+extern int   zarray_put_ex(zarray_t*, int, int*, void**);
+extern int   zarray_get_ex(zarray_t*, int, int*, void**);
 extern void* zarray_get(zarray_t*, int);
 extern void  zarray_fini(zarray_t*);
 
diff --git a/src/lib/testsuite/test_zarray.c b/src/lib/testsuite/test_zarray.c
index 13486227..294f4a11 100644
--- a/src/lib/testsuite/test_zarray.c
+++ b/src/lib/testsuite/test_zarray.c
@@ -71,6 +71,26 @@ static void test_zarray_put(void)
     zarray_fini(&array);
 }
 
+static void test_zarray_put_get_ex(void)
+{
+    zarray_t array;
+    int      idxs[2]   = { 1, 3 };
+    char*    s1        = "hello";
+    char*    s2        = "test";
+    void*    datas1[2] = { s1, s2 };
+    void*    datas2[2];
+    int      rc;
+
+    zarray_init(&array, 8, NULL);
+    rc = zarray_put_ex(&array, 2, &idxs[0], &datas1[0]);
+    TEST_ASSERT_TRUE(rc >= 0);
+    rc = zarray_get_ex(&array, 2, &idxs[0], &datas2[0]);
+    TEST_ASSERT_TRUE(rc >= 0);
+    TEST_ASSERT_TRUE((datas1[0] == datas2[0]) && (datas1[1] == datas2[1]));
+
+    zarray_fini(&array);
+}
+
 static int sum_(void* arg1, void* arg2)
 {
     int* pv = (int*)arg1;
@@ -106,4 +126,5 @@ void test_zarray(void)
     RUN_TEST(test_zarray_del);
     RUN_TEST(test_zarray_put);
     RUN_TEST(test_zarray_iter);
+    RUN_TEST(test_zarray_put_get_ex);
 }
diff --git a/src/lib/zarray.c b/src/lib/zarray.c
index b1645a4d..8e4e6c9b 100644
--- a/src/lib/zarray.c
+++ b/src/lib/zarray.c
@@ -221,53 +221,101 @@ int zarray_del(zarray_t* array, int idx)
     return 0;
 }
 
-int zarray_put(zarray_t* array, int idx, void* data)
+int zarray_put_ex(zarray_t* array, int count, int* idxs, void** data)
 {
+    int i, maxid;
     int rc = 0;
 
-    /* clang-format off */
-    if ((array == NULL) || 
-	(idx < 0) ||
-	(data == NULL)) {
-	return -EINVAL;
+    if ((array == NULL) || (count <= 0) || (idxs == NULL) || (data == NULL)) {
+        return -EINVAL;
+    }
+
+    maxid = 0;
+    for (i = 0; i < count; i++) {
+        if (idxs[i] < 0) {
+            return -EINVAL;
+        }
+
+        if (idxs[i] > maxid) {
+            maxid = idxs[i];
+        }
     }
-    /* clang-format on */
 
     mylock_lock(&(array->lock));
-    rc = zarray_ensure_items(array, idx + 1);
+    rc = zarray_ensure_items(array, maxid + 1);
     if (rc >= 0) {
-        void* ptr        = array->data[idx];
-        array->data[idx] = data;
-        rc               = 0;
-        if (ptr != NULL) {
-            if (array->killer != NULL) {
-                array->killer(ptr);
+        for (i = 0; i < count; i++) {
+            int   idx = idxs[i];
+            void* ptr = array->data[idx];
+
+            array->data[idx] = data[i];
+            if (ptr != NULL) {
+                if (array->killer != NULL) {
+                    array->killer(ptr);
+                }
+            } else {
+                array->count++;
             }
-        } else {
-            array->count++;
         }
     }
     mylock_unlock(&(array->lock));
 
-    return rc;
+    return 0;
+}
+
+int zarray_put(zarray_t* array, int idx, void* data)
+{
+    int   idxs[]  = { idx };
+    void* datas[] = { data };
+
+    return zarray_put_ex(array, 1, &idxs[0], &datas[0]);
+}
+
+int zarray_get_ex(zarray_t* array, int count, int* idxs, void** data)
+{
+    int i, maxid;
+
+    if ((array == NULL) || (count <= 0) || (idxs == NULL) || (data == NULL)) {
+        return -EINVAL;
+    }
+
+    maxid = 0;
+    for (i = 0; i < count; i++) {
+        if (idxs[i] < 0) {
+            return -EINVAL;
+        }
+
+        if (idxs[i] > maxid) {
+            maxid = idxs[i];
+        }
+    }
+
+    if (maxid >= array->caps) {
+        return -EINVAL;
+    }
+
+    mylock_lock(&(array->lock));
+    for (i = 0; i < count; i++) {
+        int idx = idxs[i];
+
+        data[i] = array->data[idx];
+    }
+    mylock_unlock(&(array->lock));
+
+    return 0;
 }
 
 void* zarray_get(zarray_t* array, int idx)
 {
-    void* ptr = NULL;
+    int   idxs[]  = { idx };
+    void* datas[] = { NULL };
+    int   rc      = 0;
 
-    /* clang-format off */
-    if ((array == NULL) || 
-        (idx < 0) || 
-        (idx >= array->caps)) {
-        SET_ERROR(EINVAL);
+    rc = zarray_get_ex(array, 1, &idxs[0], &datas[0]);
+    if (rc < 0) {
+        SET_ERROR(-rc);
         return NULL;
     }
-    /* clang-format on */
 
-    mylock_lock(&(array->lock));
-    ptr = array->data[idx];
-    mylock_unlock(&(array->lock));
-
-    return ptr;
+    return datas[0];
 }