alist: Add a way to get the next element

Add a new function which returns the next element after the one
provided, if it exists in the list.

Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/include/alist.h b/include/alist.h
index 2c78ede..97523af 100644
--- a/include/alist.h
+++ b/include/alist.h
@@ -72,6 +72,21 @@
 }
 
 /**
+ * alist_calc_index() - Calculate the index of an item in the list
+ *
+ * The returned element number will be -1 if the list is empty or the pointer
+ * pointers to before the list starts.
+ *
+ * If the pointer points to after the last item, the calculated element-number
+ * will be returned, even though it is greater than lst->count
+ *
+ * @lst: alist to check
+ * @ptr: pointer to check
+ * Return: element number of the pointer
+ */
+int alist_calc_index(const struct alist *lst, const void *ptr);
+
+/**
  * alist_err() - Check if the alist is still valid
  *
  * @lst: List to check
@@ -190,6 +205,25 @@
 #define alist_add(_lst, _obj)	\
 	((typeof(_obj) *)alist_add_ptr(_lst, &(_obj)))
 
+/** get next entry as a constant */
+#define alist_next(_lst, _objp)	\
+	((const typeof(_objp))alist_next_ptrd(_lst, _objp))
+
+/** get next entry, which can be written to */
+#define alist_nextw(_lst, _objp)	\
+	((typeof(_objp))alist_next_ptrd(_lst, _objp))
+
+/**
+ * alist_next_ptrd() - Get a pointer to the next list element
+ *
+ * This returns NULL if the requested element is beyond lst->count
+ *
+ * @lst: List to check
+ * @ptr: Pointer to current element (must be valid)
+ * Return: Pointer to next element, or NULL if @ptr is the last
+ */
+const void *alist_next_ptrd(const struct alist *lst, const void *ptr);
+
 /**
  * alist_init() - Set up a new object list
  *
diff --git a/lib/alist.c b/lib/alist.c
index b7928ca..7730fe0 100644
--- a/lib/alist.c
+++ b/lib/alist.c
@@ -106,6 +106,27 @@
 	return lst->data + index * lst->obj_size;
 }
 
+int alist_calc_index(const struct alist *lst, const void *ptr)
+{
+	uint index;
+
+	if (!lst->count || ptr < lst->data)
+		return -1;
+
+	index = (ptr - lst->data) / lst->obj_size;
+
+	return index;
+}
+
+const void *alist_next_ptrd(const struct alist *lst, const void *ptr)
+{
+	int index = alist_calc_index(lst, ptr);
+
+	assert(index != -1);
+
+	return alist_get_ptr(lst, index + 1);
+}
+
 void *alist_ensure_ptr(struct alist *lst, uint index)
 {
 	uint minsize = index + 1;
diff --git a/test/lib/alist.c b/test/lib/alist.c
index d41845c..96092af 100644
--- a/test/lib/alist.c
+++ b/test/lib/alist.c
@@ -240,3 +240,55 @@
 	return 0;
 }
 LIB_TEST(lib_test_alist_add, 0);
+
+/* Test alist_next()  */
+static int lib_test_alist_next(struct unit_test_state *uts)
+{
+	const struct my_struct *ptr;
+	struct my_struct data, *ptr2;
+	struct alist lst;
+	ulong start;
+
+	start = ut_check_free();
+
+	ut_assert(alist_init_struct(&lst, struct my_struct));
+	data.val = 123;
+	data.other_val = 0;
+	alist_add(&lst, data);
+
+	data.val = 321;
+	alist_add(&lst, data);
+
+	data.val = 789;
+	alist_add(&lst, data);
+
+	ptr = alist_get(&lst, 0, struct my_struct);
+	ut_assertnonnull(ptr);
+	ut_asserteq(123, ptr->val);
+
+	ptr = alist_next(&lst, ptr);
+	ut_assertnonnull(ptr);
+	ut_asserteq(321, ptr->val);
+
+	ptr2 = (struct my_struct *)ptr;
+	ptr2 = alist_nextw(&lst, ptr2);
+	ut_assertnonnull(ptr2);
+
+	ptr = alist_next(&lst, ptr);
+	ut_assertnonnull(ptr);
+	ut_asserteq(789, ptr->val);
+	ut_asserteq_ptr(ptr, ptr2);
+	ptr2->val = 89;
+	ut_asserteq(89, ptr->val);
+
+	ptr = alist_next(&lst, ptr);
+	ut_assertnull(ptr);
+
+	alist_uninit(&lst);
+
+	/* Check for memory leaks */
+	ut_assertok(ut_check_delta(start));
+
+	return 0;
+}
+LIB_TEST(lib_test_alist_next, 0);