UCS/DATASTRUCT: Add ucs_array_for_each_index

Signed-off-by: Guy Ealey Morag <gealeymorag@nvidia.com>

Author: guy-ealey-morag

.clang-format

@@ -72,6 +72,7 @@ ForEachMacros: ['UCS_STATIC_BITMAP_FOR_EACH_BIT',
                 'UCS_STATIC_BITMAP_FOR_EACH_BIT',
                 'FOR_EACH_ENTITY',
                 'ucs_carray_for_each',
+                'ucs_carray_for_each_index',
                 'kh_foreach',
                 'kh_foreach_key',
                 'kh_foreach_value',
@@ -81,6 +82,7 @@ ForEachMacros: ['UCS_STATIC_BITMAP_FOR_EACH_BIT',
                 'ucp_proto_perf_segment_foreach',
                 'ucp_proto_perf_segment_foreach_range',
                 'ucs_array_for_each',
+                'ucs_array_for_each_index',
                 'UCS_BITMAP_FOR_EACH_BIT',
                 'ucs_for_each_bit',
                 'ucs_for_each_submask',

src/ucp/proto/proto_select.c

@@ -75,9 +75,7 @@ static ucs_status_t ucp_proto_thresholds_next_range(
     ucs_dynamic_bitmap_reset_all(proto_mask);
     ucs_dynamic_bitmap_init(&disabled_proto_mask);
 
-    for (proto_idx = 0; proto_idx < ucs_array_length(&proto_init->protocols);
-         ++proto_idx) {
-        proto = &ucs_array_elem(&proto_init->protocols, proto_idx);
+    ucs_array_for_each_index(proto, proto_idx, &proto_init->protocols) {
         range = ucp_proto_flat_perf_find_lb(proto->flat_perf, msg_length);
         if (range == NULL) {
             ucs_trace("skipping proto %s for msg_length %zu",

src/ucs/datastruct/array.h

@@ -1,5 +1,5 @@
 /**
- * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.
+ * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -456,6 +456,18 @@ ucs_array_old_buffer_set_null(void **old_buffer_p)
                         ucs_array_length(_array))
 
 
+/**
+ * Iterate over array elements and track the current index
+ *
+ * @param _elem    Pointer variable to the current array element
+ * @param _idx     Variable containing the current array index
+ * @param _array   Array to iterate over
+ */
+#define ucs_array_for_each_index(_elem, _idx, _array) \
+    ucs_carray_for_each_index(_elem, _idx, ucs_array_begin(_array), \
+                              ucs_array_length(_array))
+
+
 /* Internal helper function */
 ucs_status_t ucs_array_grow(void **buffer_p, size_t *capacity_p,
                             size_t min_capacity, size_t max_capacity,

src/ucs/datastruct/callbackq.c

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2016. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 * Copyright (C) ARM Ltd. 2016-2017.  ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
@@ -295,9 +295,7 @@ static void ucs_callbackq_spill_elems_purge(ucs_callbackq_t *cbq)
      * always be equal to dst_idx, so nothing will be actually copied/moved.
      */
     dst_idx = 0;
-    for (src_idx = 0; src_idx < ucs_array_length(&priv->spill_elems);
-         ++src_idx) {
-        src_elem = &ucs_array_elem(&priv->spill_elems, src_idx);
+    ucs_array_for_each_index(src_elem, src_idx, &priv->spill_elems) {
         if (src_elem->id != UCS_CALLBACKQ_ID_NULL) {
             ucs_assert(dst_idx <= src_idx);
             if (dst_idx != src_idx) {

src/ucs/sys/compiler_def.h

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2017. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 * Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED.
 * Copyright (C) Arm, Ltd. 2021. ALL RIGHTS RESERVED.
 * Copyright (C) Advanced Micro Devices, Inc. 2024. ALL RIGHTS RESERVED.
@@ -250,6 +250,12 @@
 #define ucs_carray_for_each(_elem, _array, _length) \
     for ((_elem) = (_array); (_elem) < ((_array) + (_length)); ++(_elem))
 
+/*
+ * Iterate over all elements of a C-array and track the current index
+ */
+#define ucs_carray_for_each_index(_elem, _idx, _array, _length) \
+    for ((_idx) = 0, (_elem) = (_array); (_idx) < (_length); ++(_idx), ++(_elem))
+
 /*
  * Swap two variables values
  */

test/gtest/ucs/test_datatype.cc

@@ -1,5 +1,5 @@
 /**
-* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2014. ALL RIGHTS RESERVED.
+* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
 * Copyright (C) UT-Battelle, LLC. 2014. ALL RIGHTS RESERVED.
 * Copyright (C) Huawei Technologies Co., Ltd. 2020.  ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
@@ -1115,6 +1115,50 @@ class test_array : public test_datatype {
     void cleanup_array_of_linked_lists(test_list_links_array_t *test_array);
 };
 
+UCS_TEST_F(test_array, carray_for_each_index) {
+    const int values[] = {3, 5, 7};
+    const int *elem;
+    size_t idx;
+
+    ucs_carray_for_each_index(elem, idx, values, 0) {
+        ADD_FAILURE() << "iterated over an empty array";
+    }
+    EXPECT_EQ(0, idx);
+
+    ucs_carray_for_each_index(elem, idx, values,
+                              ucs_static_array_size(values)) {
+        EXPECT_EQ(&values[idx], elem);
+        EXPECT_EQ(values[idx], *elem);
+    }
+    EXPECT_EQ(ucs_static_array_size(values), idx);
+}
+
+UCS_TEST_F(test_array, array_for_each_index) {
+    constexpr size_t NUM_ELEMENTS = 3;
+    test_1int_t test_array;
+    int *elem;
+    size_t idx;
+
+    ucs_array_init_dynamic(&test_array);
+
+    ucs_array_for_each_index(elem, idx, &test_array) {
+        ADD_FAILURE() << "iterated over an empty array";
+    }
+    EXPECT_EQ(0, idx);
+
+    for (size_t i = 0; i < NUM_ELEMENTS; ++i) {
+        *ucs_array_append(&test_array, FAIL()) = static_cast<int>(i * i + 1);
+    }
+
+    ucs_array_for_each_index(elem, idx, &test_array) {
+        EXPECT_EQ(&ucs_array_elem(&test_array, idx), elem);
+        EXPECT_EQ(idx * idx + 1, static_cast<size_t>(*elem));
+    }
+    EXPECT_EQ(NUM_ELEMENTS, idx);
+
+    ucs_array_cleanup_dynamic(&test_array);
+}
+
 /* generate a list of numbers in a certain size */
 void test_array::generate_linked_list(int size, simple_elem_t *head)
 {