Skip to content

UCS/DATASTRUCT: Interval tree fixes #11213

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yosefe merged 3 commits into from
Mar 3, 2026
Merged

UCS/DATASTRUCT: Interval tree fixes #11213

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0cb50c8
UCS/DATASTRUCT: Interval tree fixes
shasson5 Feb 25, 2026
14a6e56
UCS/DATASTURCT: CR fixes
shasson5 Mar 1, 2026
6224ab6
UCS/DATASTRUCT: CR fixes
shasson5 Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 6 additions & 25 deletions src/ucp/core/ucp_am.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,18 +28,6 @@
#define ucp_am_rdesc_frag_tree(_rdesc) \
(UCS_PTR_BYTE_OFFSET(_rdesc, -sizeof(ucs_interval_tree_t)))


static void *ucp_am_frag_tree_alloc_node(size_t size, void *arg)
{
ucs_mpool_t *mp = (ucs_mpool_t*)arg;
return ucs_mpool_get(mp);
}

static void ucp_am_frag_tree_free_node(void *node, void *arg)
{
ucs_mpool_put(node);
}

static ucs_mpool_ops_t ucp_am_frag_tree_mpool_ops = {
.chunk_alloc = ucs_mpool_chunk_malloc,
.chunk_release = ucs_mpool_chunk_free,
Expand Down Expand Up @@ -1415,12 +1403,9 @@ ucp_am_copy_data_fragment(ucp_recv_desc_t *first_rdesc, void *data,
UCS_PROFILE_NAMED_CALL("am_memcpy_recv", ucs_memcpy_relaxed,
UCS_PTR_BYTE_OFFSET(first_rdesc + 1, offset),
data, length, UCS_ARCH_MEMCPY_NT_SOURCE, length);

/* Message length can be 0 if only header was sent */
if (ucs_likely(length > 0)) {
ucs_interval_tree_insert(ucp_am_rdesc_frag_tree(first_rdesc), offset,
offset + length - 1);
}

ucs_interval_tree_insert(ucp_am_rdesc_frag_tree(first_rdesc), offset,
offset + length - 1);
}

static UCS_F_ALWAYS_INLINE uint64_t
Expand Down Expand Up @@ -1463,8 +1448,7 @@ ucp_am_handle_unfinished(ucp_worker_h worker, ucp_recv_desc_t *rdesc,
first_ftr = (ucp_am_first_ftr_t*)(first_rdesc + 1);
seg_end = first_rdesc->payload_offset + first_ftr->total_size - 1;

if ((first_ftr->total_size > 0) &&
!ucs_interval_tree_is_equal_range(ucp_am_rdesc_frag_tree(
if (!ucs_interval_tree_is_equal_range(ucp_am_rdesc_frag_tree(
first_rdesc),
first_rdesc->payload_offset,
seg_end)) {
Expand Down Expand Up @@ -1557,7 +1541,6 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_am_long_first_handler,
size_t total_length, padding;
uint64_t recv_flags;
void *user_hdr, *buffer;
ucs_interval_tree_ops_t frag_tree_ops;

first_payload_length = am_length - sizeof(*first_ftr);
first_ftr = UCS_PTR_BYTE_OFFSET(am_data, first_payload_length);
Expand Down Expand Up @@ -1617,11 +1600,9 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_am_long_first_handler,
worker->am.alignment);

first_rdesc->payload_offset = UCP_AM_FIRST_FRAG_META_LEN + padding;
frag_tree_ops.alloc_node = ucp_am_frag_tree_alloc_node;
frag_tree_ops.free_node = ucp_am_frag_tree_free_node;
frag_tree_ops.arg = &worker->am.frag_tree_mpool;
/* Initialize frag_tree in the allocated storage before rdesc */
ucs_interval_tree_init(ucp_am_rdesc_frag_tree(first_rdesc), &frag_tree_ops);
ucs_interval_tree_init(ucp_am_rdesc_frag_tree(first_rdesc),
&worker->am.frag_tree_mpool);

/* Copy first fragment and base headers before the data, it will be needed
* for middle fragments processing. */
Expand Down
22 changes: 13 additions & 9 deletions src/ucs/datastruct/interval_tree.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@
#endif

#include "interval_tree.h"
#include <ucs/debug/assert.h>
#include <ucs/sys/math.h>
#include <ucs/datastruct/mpool.h>


/**
Expand All @@ -23,7 +23,7 @@ ucs_interval_tree_node_create(ucs_interval_tree_t *tree, uint64_t start,
{
ucs_interval_node_t *node;

node = tree->ops.alloc_node(sizeof(*node), tree->ops.arg);
node = ucs_mpool_get(tree->mpool);
Comment thread
yosefe marked this conversation as resolved.
Resolved
if (node == NULL) {
return NULL;
}
Expand All @@ -35,11 +35,11 @@ ucs_interval_tree_node_create(ucs_interval_tree_t *tree, uint64_t start,
return node;
}

void ucs_interval_tree_init(ucs_interval_tree_t *tree,
const ucs_interval_tree_ops_t *ops)
void ucs_interval_tree_init(ucs_interval_tree_t *tree, ucs_mpool_t *mpool)
{
tree->ops = *ops;
tree->root = NULL;
tree->mpool = mpool;
tree->root = NULL;
tree->single_node = 0;
}

static void ucs_interval_tree_cleanup_recursive(ucs_interval_tree_t *tree,
Expand All @@ -51,7 +51,7 @@ static void ucs_interval_tree_cleanup_recursive(ucs_interval_tree_t *tree,

ucs_interval_tree_cleanup_recursive(tree, node->left);
ucs_interval_tree_cleanup_recursive(tree, node->right);
tree->ops.free_node(node, tree->ops.arg);
ucs_mpool_put(node);
}

void ucs_interval_tree_cleanup(ucs_interval_tree_t *tree)
Expand Down Expand Up @@ -92,11 +92,11 @@ ucs_interval_tree_remove_node(ucs_interval_tree_t *tree,
/* Found target - handle deletion based on number of children */
if (root->left == NULL) {
temp = root->right;
tree->ops.free_node(root, tree->ops.arg);
ucs_mpool_put(root);
return temp;
} else if (root->right == NULL) {
temp = root->left;
tree->ops.free_node(root, tree->ops.arg);
ucs_mpool_put(root);
return temp;
}

Expand Down Expand Up @@ -205,5 +205,9 @@ ucs_status_t ucs_interval_tree_insert_slow(ucs_interval_tree_t *tree,

/* Insert the merged node into the tree */
tree->root = ucs_interval_tree_insert_node(tree->root, new_node);

/* Update single_node flag: tree is single node if root has no children */
tree->single_node = (tree->root->left == NULL) &&

@yosefe yosefe Feb 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe have tree->num_nodes and do ++ every time we insert to it and -- when remove?

(tree->root->right == NULL);
return UCS_OK;
}
58 changes: 15 additions & 43 deletions src/ucs/datastruct/interval_tree.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,41 +10,11 @@
#include <ucs/debug/assert.h>
#include <ucs/sys/math.h>
#include <ucs/type/status.h>
#include <ucs/datastruct/mpool.h>
#include <stdint.h>
#include <stddef.h>


/**
* Callback for allocating an interval tree node
*
* @param [in] size Size of the node to allocate
* @param [in] arg User-defined argument passed during tree initialization

* @return Pointer to allocated node, or NULL on failure
*/
typedef void *(*ucs_interval_tree_alloc_node_func_t)(size_t size, void *arg);


/**
* Callback for deallocating an interval tree node
*
* @param [in] node Node to deallocate
* @param [in] arg User-defined argument passed during tree initialization
*/
typedef void (*ucs_interval_tree_free_node_func_t)(void *node, void *arg);


/**
* Interval tree node allocation/deallocation operations
*/
typedef struct {
ucs_interval_tree_alloc_node_func_t alloc_node; /* Node allocation callback */
ucs_interval_tree_free_node_func_t free_node; /* Node deallocation callback */
void *arg; /* User-defined argument for
* callbacks */
} ucs_interval_tree_ops_t;


/**
* Interval tree node structure
*/
Expand All @@ -57,19 +27,19 @@ typedef struct ucs_interval_node {


typedef struct {
ucs_interval_node_t *root; /**< Root node of the tree */
ucs_interval_tree_ops_t ops; /**< Memory management callbacks */
ucs_interval_node_t *root; /**< Root node of the tree */
ucs_mpool_t *mpool; /**< Memory pool for node allocation */
int single_node; /**< Cached flag: 1 if tree has only root node */

@yosefe yosefe Feb 26, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: fix indent

} ucs_interval_tree_t;


/**
* Initialize an interval tree
*
* @param [in] tree Interval tree to initialize
* @param [in] ops Memory management operations (alloc, free, arg)
* @param [in] tree Interval tree to initialize
* @param [in] mpool Memory pool for node allocation
*/
void ucs_interval_tree_init(ucs_interval_tree_t *tree,
const ucs_interval_tree_ops_t *ops);
void ucs_interval_tree_init(ucs_interval_tree_t *tree, ucs_mpool_t *mpool);


/**
Expand Down Expand Up @@ -108,13 +78,15 @@ ucs_interval_tree_is_single_node(const ucs_interval_tree_t *tree)
static UCS_F_ALWAYS_INLINE ucs_status_t ucs_interval_tree_insert(
ucs_interval_tree_t *tree, uint64_t start, uint64_t end)
{
ucs_assert(start <= end);
ucs_interval_node_t *root = tree->root;

ucs_assertv(start <= end + 1, "start=%lu, end=%lu", start, end);
Comment thread
yosefe marked this conversation as resolved.
Resolved

/* Fast path: if tree has only root and new interval overlaps/touches it, extend it */
if (ucs_interval_tree_is_single_node(tree) &&
(start <= (tree->root->end + 1)) && (tree->root->start <= (end + 1))) {
tree->root->start = ucs_min(tree->root->start, start);
tree->root->end = ucs_max(tree->root->end, end);
if (ucs_likely(tree->single_node) && ucs_likely(start <= (root->end + 1)) &&
ucs_likely(root->start <= (end + 1))) {
root->start = ucs_min(root->start, start);
root->end = ucs_max(root->end, end);
return UCS_OK;
}

Expand All @@ -135,7 +107,7 @@ static UCS_F_ALWAYS_INLINE int
ucs_interval_tree_is_equal_range(const ucs_interval_tree_t *tree,
uint64_t start, uint64_t end)
{
ucs_assert(start <= end);
ucs_assertv(start <= end + 1, "start=%lu, end=%lu", start, end);
Comment thread
yosefe marked this conversation as resolved.
Resolved

return ucs_interval_tree_is_single_node(tree) &&
(tree->root->start == start) && (tree->root->end == end);
Expand Down
42 changes: 36 additions & 6 deletions test/gtest/ucs/test_interval_tree.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

extern "C" {
#include <ucs/datastruct/interval_tree.h>
#include <ucs/datastruct/mpool.h>
}

class test_interval_tree : public ucs::test {
Expand All @@ -20,12 +21,25 @@ class test_interval_tree : public ucs::test {
void init()
{
ucs::test::init();
ucs_interval_tree_init(&m_tree, &m_ops);

/* Initialize memory pool for interval tree nodes */
ucs_mpool_params_t mp_params;
ucs_mpool_params_reset(&mp_params);
mp_params.elem_size = sizeof(ucs_interval_node_t);
mp_params.elems_per_chunk = 256;
mp_params.ops = &m_mpool_ops;
mp_params.name = "test_interval_tree_nodes";

ucs_status_t status = ucs_mpool_init(&mp_params, &m_mpool);
ASSERT_UCS_OK(status);

ucs_interval_tree_init(&m_tree, &m_mpool);
}

void cleanup()
{
ucs_interval_tree_cleanup(&m_tree);
ucs_mpool_cleanup(&m_mpool, 0);
ucs::test::cleanup();
}

Expand Down Expand Up @@ -56,10 +70,18 @@ class test_interval_tree : public ucs::test {
}

private:
ucs_interval_tree_t m_tree = {NULL, {NULL}};
ucs_interval_tree_ops_t m_ops = {(ucs_interval_tree_alloc_node_func_t)malloc,
(ucs_interval_tree_free_node_func_t)free,
NULL};
static ucs_mpool_ops_t m_mpool_ops;

ucs_interval_tree_t m_tree;
ucs_mpool_t m_mpool;
};

ucs_mpool_ops_t test_interval_tree::m_mpool_ops = {
.chunk_alloc = ucs_mpool_chunk_malloc,
.chunk_release = ucs_mpool_chunk_free,
.obj_init = NULL,
.obj_cleanup = NULL,
.obj_str = NULL
};

UCS_TEST_F(test_interval_tree, single_interval) {
Expand Down Expand Up @@ -111,14 +133,22 @@ UCS_TEST_F(test_interval_tree, out_of_order_insertion) {
EXPECT_TRUE(is_fully_covered({0, 60}));
}

UCS_TEST_F(test_interval_tree, zero_length_insert) {
UCS_TEST_F(test_interval_tree, length_one_insert) {
insert_intervals({{10, 10}});
EXPECT_TRUE(is_fully_covered({10, 10}));

insert_intervals({{10, 20}, {15, 15}});
EXPECT_TRUE(is_fully_covered({10, 20}));
}

UCS_TEST_F(test_interval_tree, length_zero_insert) {
insert_intervals({{10, 9}});
EXPECT_TRUE(is_fully_covered({10, 9}));

insert_intervals({{10, 20}, {20, 19}});
EXPECT_TRUE(is_fully_covered({10, 20}));
}

UCS_TEST_F(test_interval_tree, large_intervals) {
insert_intervals({{1000000, 2000000}, {2000002, 3000000}});
EXPECT_FALSE(is_any_covered(
Expand Down