diff --git a/librz/include/rz_vector.h b/librz/include/rz_vector.h index a319b926a6..f3aa6d6d75 100644 --- a/librz/include/rz_vector.h +++ b/librz/include/rz_vector.h @@ -144,6 +144,7 @@ RZ_API void *rz_vector_insert(RzVector *vec, size_t index, void *x); RZ_API void *rz_vector_insert_range(RzVector *vec, size_t index, RZ_NULLABLE void *first, size_t count); RZ_API void *rz_vector_insert_sorted(RZ_NONNULL RzVector *vec, RZ_NONNULL void *elem, RzVectorComparator cmp, void *user); +RZ_API size_t rz_vector_find_sorted(RZ_NONNULL RzVector *vec, RZ_NONNULL void *elem, RzVectorComparator cmp, void *user); // like rz_vector_remove_at for the last element RZ_API void rz_vector_pop(RzVector *vec, void *into); diff --git a/librz/util/vector.c b/librz/util/vector.c index 4e6935142e..7680d79343 100644 --- a/librz/util/vector.c +++ b/librz/util/vector.c @@ -291,6 +291,59 @@ RZ_API void *rz_vector_insert_sorted(RZ_NONNULL RzVector *vec, RZ_NONNULL void * return vec->reverse_sorted ? rz_vector_push_front(vec, elem) : rz_vector_push(vec, elem); } +static bool bin_search_range(RZ_NONNULL RzVector *vec, RZ_NONNULL void *elem, RzVectorComparator cmp, void *user, RZ_OUT size_t *i) { + size_t vlen = rz_vector_len(vec); + if (vlen == 0) { + return false; + } + + int inc = vec->reverse_sorted ? -1 : 1; + ssize_t low = vec->reverse_sorted ? vlen - 1 : 0; + ssize_t hi = vec->reverse_sorted ? 0 : vlen - 1; + + do { + size_t mid = (low + hi) >> 1; + if (cmp(elem, rz_vector_index_ptr(vec, mid), user) == 0) { + *i = mid; + return true; + } + if (low == hi) { + break; + } + if (cmp(elem, rz_vector_index_ptr(vec, mid), user) > 0) { + low = mid + inc; + } + if (cmp(elem, rz_vector_index_ptr(vec, mid), user) < 0) { + hi = mid - inc; + } + } while (vec->reverse_sorted ? hi <= low : low <= hi); + + return false; +} + +/** + * \brief Finds an element in the sorted vector via binary search. + * NOTE: This function assumes the vector is already sorted! + * If it isn't the result is undefined! + * + * \param vec A sorted vector to find the element in. + * \param elem Pointer to the element to find in the vector. + * \param cmp The comparator for the elements. + * \param user The user data passed to the comparator. + * + * \return Index into the vector where the element is located. + * Or SZT_MAX in case of failure or if no element was found. + */ +RZ_API size_t rz_vector_find_sorted(RZ_NONNULL RzVector *vec, RZ_NONNULL void *elem, RzVectorComparator cmp, void *user) { + rz_return_val_if_fail(vec && elem, SZT_MAX); + + size_t i; + if (!bin_search_range(vec, elem, cmp, user, &i)) { + return SZT_MAX; + } + return i; +} + RZ_API void rz_vector_pop(RzVector *vec, void *into) { if (rz_vector_empty(vec)) { return; diff --git a/test/unit/test_vector.c b/test/unit/test_vector.c index 964608fa5c..c67781481d 100644 --- a/test/unit/test_vector.c +++ b/test/unit/test_vector.c @@ -305,6 +305,62 @@ static bool test_vector_insert_sorted(void) { mu_end; } +static bool test_vector_find_sorted(void) { + RzVector *v = rz_vector_new(sizeof(ut64), NULL, NULL); + for (size_t i = 1; i < 13; i++) { + rz_vector_push(v, &i); + } + ut64 i = UT64_MAX; + rz_vector_push(v, &i); + + rz_vector_sort(v, (RzVectorComparator)uint_cmp, false, NULL); + + i = 5; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 4, "Not found"); + i = 6; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 5, "Not found"); + i = 7; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 6, "Not found"); + i = 8; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 7, "Not found"); + i = 12; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 11, "Not found"); + i = UT64_MAX; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 12, "Not found"); + + i = 0; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + i = 13; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + i = UT64_MAX - 1; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + + rz_vector_sort(v, (RzVectorComparator)uint_cmp, true, NULL); + + i = 5; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 8, "Not found"); + i = 6; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 7, "Not found"); + i = 7; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 6, "Not found"); + i = 8; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 5, "Not found"); + i = 12; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 1, "Not found"); + i = UT64_MAX; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), 0, "Not found"); + + i = 0; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + i = 13; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + i = UT64_MAX - 1; + mu_assert_eq(rz_vector_find_sorted(v, &i, (RzVectorComparator)uint_cmp, NULL), SZT_MAX, "Not failed"); + + rz_vector_free(v); + mu_end; +} + static bool test_vector_empty(void) { RzVector v; rz_vector_init(&v, 1, NULL, NULL); @@ -1562,6 +1618,7 @@ static int all_tests(void) { mu_run_test(test_vector_insert); mu_run_test(test_vector_insert_range); mu_run_test(test_vector_insert_sorted); + mu_run_test(test_vector_find_sorted); mu_run_test(test_vector_pop); mu_run_test(test_vector_pop_front); mu_run_test(test_vector_push);