From 8d9f230514e7c5231a34cc40d5e9269380549f70 Mon Sep 17 00:00:00 2001 From: teerthsharma Date: Mon, 20 Jul 2026 08:59:54 +0530 Subject: [PATCH] Export disjoint-set island helpers directly Signed-off-by: teerthsharma --- src/engine/engine_island.c | 46 +++++---------- src/engine/engine_island.h | 11 ++-- test/engine/engine_island_test.cc | 96 +++++++++++++++---------------- 3 files changed, 67 insertions(+), 86 deletions(-) diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 4f53d2db..ae986cd6 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -84,7 +84,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) { //-------------------------- flood-fill and graph construction ------------------------------------ // find the canonical root of an active tree and compress its path -static inline int dsuFind(int* parent, int tree) { +int mj_dsuRoot(int* parent, int tree) { int root = tree; while (parent[root] != root) { root = parent[root]; @@ -101,13 +101,13 @@ static inline int dsuFind(int* parent, int tree) { // initialize all trees as inactive -static inline void dsuInit(int* parent, int ntree) { +void mj_dsuInit(int* parent, int ntree) { mju_fillInt(parent, -1, ntree); } // activate and union two incident trees; -1 denotes a static endpoint -static inline void dsuUnion(int* parent, int tree1, int tree2) { +void mj_dsuMerge(int* parent, int tree1, int tree2) { if (tree1 == -1 && tree2 == -1) { mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR return; @@ -121,8 +121,8 @@ static inline void dsuUnion(int* parent, int tree1, int tree2) { if (parent[tree1] == parent[tree2]) return; - int root1 = dsuFind(parent, tree1); - int root2 = dsuFind(parent, tree2); + int root1 = mj_dsuRoot(parent, tree1); + int root2 = mj_dsuRoot(parent, tree2); if (root1 < root2) { parent[root2] = root1; } else if (root2 < root1) { @@ -132,8 +132,8 @@ static inline void dsuUnion(int* parent, int tree1, int tree2) { // assign deterministic island ids in ascending canonical-root order -static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, - int* nidof) { +int mj_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, + int* nidof) { int nisland = 0; *nidof = 0; for (int tree=0; tree < ntree; tree++) { @@ -145,7 +145,7 @@ static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, in if (parent[tree] == tree) { island[tree] = nisland++; } else { - // Union always links the larger root to the smaller root. Since trees are visited in + // union always links the larger root to the smaller root. Since trees are visited in // ascending order, this predecessor has already been compressed and assigned an island. parent[tree] = parent[parent[tree]]; island[tree] = island[parent[tree]]; @@ -157,26 +157,6 @@ static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, in } -// exported private wrappers for direct unit tests and benchmarks -int _mjPRIVATE_dsuFind(int* parent, int tree) { - return dsuFind(parent, tree); -} - - -void _mjPRIVATE_dsuInit(int* parent, int ntree) { - dsuInit(parent, ntree); -} - - -void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2) { - dsuUnion(parent, tree1, tree2); -} - - -int _mjPRIVATE_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { - return dsuAssign(island, parent, tree_dofnum, ntree, nidof); -} - // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: // island (nr) - island index assigned to vertex, -1 if vertex has no edges @@ -413,10 +393,10 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { - dsuUnion(parent, tree1, -1); + mj_dsuMerge(parent, tree1, -1); } else { while (tree2 != -2) { - dsuUnion(parent, tree1, tree2); + mj_dsuMerge(parent, tree1, tree2); tree1 = tree2; tree2 = treeNext(m, d, i, &iter); } @@ -458,7 +438,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, if (tree1 < 0) { tree1 = tree2; } else { - dsuUnion(parent, tree1, tree2); + mj_dsuMerge(parent, tree1, tree2); } } } @@ -483,11 +463,11 @@ void mj_island(const mjModel* m, mjData* d) { // union direct tree incidence and assign deterministic components int* efc_tree = mjSTACKALLOC(d, nefc, int); int* parent = mjSTACKALLOC(d, ntree, int); - dsuInit(parent, ntree); + mj_dsuInit(parent, ntree); unionConstraintTrees(m, d, parent, efc_tree); int* tree_island = mjSTACKALLOC(d, ntree, int); int nidof; - d->nisland = dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); + d->nisland = mj_dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); // no islands found: quick return if (!d->nisland) { diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 7e239b40..ca48ba20 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -23,11 +23,12 @@ extern "C" { #endif -MJAPI int _mjPRIVATE_dsuFind(int* parent, int tree); -MJAPI void _mjPRIVATE_dsuInit(int* parent, int ntree); -MJAPI void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2); -MJAPI int _mjPRIVATE_dsuAssign(int* island, int* parent, - const int* tree_dofnum, int ntree, int* nidof); +// disjoint-set roots are minimum tree indices; mj_dsuRoot requires parent[tree] >= 0 +MJAPI void mj_dsuInit(int* parent, int ntree); +MJAPI void mj_dsuMerge(int* parent, int tree1, int tree2); +MJAPI int mj_dsuRoot(int* parent, int tree); +MJAPI int mj_dsuAssign(int* island, int* parent, + const int* tree_dofnum, int ntree, int* nidof); // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index be1f2832..241a053c 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -40,96 +40,96 @@ using IslandTest = MujocoTest; TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) { int parent[] = {8, 6, 7, 5}; - _mjPRIVATE_dsuInit(parent, 0); + mj_dsuInit(parent, 0); EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5)); - _mjPRIVATE_dsuInit(parent, 4); + mj_dsuInit(parent, 4); EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1)); } -TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) { +TEST_F(IslandTest, DsuRootReturnsCanonicalRootAndCompressesPath) { int parent[] = {0, 0, 1, 2, 3}; - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0); - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0); + EXPECT_EQ(mj_dsuRoot(parent, 0), 0); + EXPECT_EQ(mj_dsuRoot(parent, 4), 0); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0)); } -TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) { +TEST_F(IslandTest, DsuMergeActivatesEndpointsAndUsesMinimumRoot) { int parent[] = {-1, -1, -1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, -1, 4); - _mjPRIVATE_dsuUnion(parent, 3, -1); - _mjPRIVATE_dsuUnion(parent, 5, 2); - _mjPRIVATE_dsuUnion(parent, 4, 5); - _mjPRIVATE_dsuUnion(parent, 3, 4); + mj_dsuMerge(parent, -1, 4); + mj_dsuMerge(parent, 3, -1); + mj_dsuMerge(parent, 5, 2); + mj_dsuMerge(parent, 4, 5); + mj_dsuMerge(parent, 3, 4); EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); for (int tree = 2; tree < 6; ++tree) { - EXPECT_EQ(_mjPRIVATE_dsuFind(parent, tree), 2); + EXPECT_EQ(mj_dsuRoot(parent, tree), 2); } EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); } -TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) { +TEST_F(IslandTest, DsuMergeRedundantAndReversedEdgesAreIdempotent) { int parent[] = {-1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, 3, 1); - _mjPRIVATE_dsuUnion(parent, 2, 1); + mj_dsuMerge(parent, 3, 1); + mj_dsuMerge(parent, 2, 1); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); - _mjPRIVATE_dsuUnion(parent, 1, 3); - _mjPRIVATE_dsuUnion(parent, 3, 1); - _mjPRIVATE_dsuUnion(parent, 2, 2); - _mjPRIVATE_dsuUnion(parent, -1, 2); + mj_dsuMerge(parent, 1, 3); + mj_dsuMerge(parent, 3, 1); + mj_dsuMerge(parent, 2, 2); + mj_dsuMerge(parent, -1, 2); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); } -TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) { +TEST_F(IslandTest, DsuMergeFastPathActivatesBeforeTestingParents) { int self_parent[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(self_parent, 1, 1); + mj_dsuMerge(self_parent, 1, 1); EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1)); int static_first[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(static_first, -1, 2); + mj_dsuMerge(static_first, -1, 2); EXPECT_THAT(static_first, ElementsAre(-1, -1, 2)); int static_second[] = {-1, -1, -1}; - _mjPRIVATE_dsuUnion(static_second, 0, -1); + mj_dsuMerge(static_second, 0, -1); EXPECT_THAT(static_second, ElementsAre(0, -1, -1)); } -TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) { +TEST_F(IslandTest, DsuMergeFastPathDistinguishesParentsFromRoots) { int distinct_parent[] = {0, 0, 2, 2}; - _mjPRIVATE_dsuUnion(distinct_parent, 1, 3); + mj_dsuMerge(distinct_parent, 1, 3); EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2)); int shared_parent[] = {0, 0, 0, 3}; - _mjPRIVATE_dsuUnion(shared_parent, 1, 2); + mj_dsuMerge(shared_parent, 1, 2); EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3)); int long_paths[] = {0, 0, 1, 3, 3, 4}; - _mjPRIVATE_dsuUnion(long_paths, 2, 5); + mj_dsuMerge(long_paths, 2, 5); EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3)); } -TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) { +TEST_F(IslandTest, DsuMergeFastPathPreservesCyclesDuplicatesAndForest) { int parent[] = {-1, -1, -1, -1, -1, -1}; - _mjPRIVATE_dsuUnion(parent, 0, 1); - _mjPRIVATE_dsuUnion(parent, 1, 2); - _mjPRIVATE_dsuUnion(parent, 2, 0); - _mjPRIVATE_dsuUnion(parent, 0, 2); - _mjPRIVATE_dsuUnion(parent, 3, 4); - _mjPRIVATE_dsuUnion(parent, 4, 5); + mj_dsuMerge(parent, 0, 1); + mj_dsuMerge(parent, 1, 2); + mj_dsuMerge(parent, 2, 0); + mj_dsuMerge(parent, 0, 2); + mj_dsuMerge(parent, 3, 4); + mj_dsuMerge(parent, 4, 5); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3)); - _mjPRIVATE_dsuUnion(parent, 5, 0); + mj_dsuMerge(parent, 5, 0); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3)); } -TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) { +TEST_F(IslandTest, DsuMergeRejectsStaticSelfIncidence) { int parent[] = {-1, 1, 1, 3}; - EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1), + EXPECT_EQ(MjuErrorMessageFrom(mj_dsuMerge)(parent, -1, -1), "self-incidence of the static tree"); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3)); } @@ -140,13 +140,13 @@ TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) { const int tree_dofnum[] = {73}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 0, &nidof), 0); EXPECT_EQ(nidof, 0); EXPECT_THAT(island, ElementsAre(71)); EXPECT_THAT(parent, ElementsAre(72)); parent[0] = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 1, &nidof), 0); EXPECT_EQ(nidof, 0); EXPECT_THAT(island, ElementsAre(-1)); EXPECT_THAT(parent, ElementsAre(-1)); @@ -158,7 +158,7 @@ TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) { int island[] = {9, 9, 9, 9, 9, 9, 9}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 7, &nidof), 3); EXPECT_EQ(nidof, 39); EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6)); @@ -170,7 +170,7 @@ TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) { int island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9}; int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3); + EXPECT_EQ(mj_dsuAssign(island, parent, tree_dofnum, 10, &nidof), 3); EXPECT_EQ(nidof, 41); EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 4, 7, 7, 7)); @@ -190,7 +190,7 @@ TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) { } int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + EXPECT_EQ(mj_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), kTreeCount, &nidof), 1); EXPECT_EQ(nidof, expected_nidof); @@ -205,11 +205,11 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { std::vector parent(kTreeCount); std::vector island(kTreeCount, -2); std::vector tree_dofnum(kTreeCount); - _mjPRIVATE_dsuInit(parent.data(), kTreeCount); + mj_dsuInit(parent.data(), kTreeCount); int expected_nidof = 0; for (int tree = kTreeCount - 1; tree > 0; --tree) { - _mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1); + mj_dsuMerge(parent.data(), tree, tree - 1); } for (int tree = 0; tree < kTreeCount; ++tree) { tree_dofnum[tree] = tree % 7; @@ -217,7 +217,7 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { } int nidof = -1; - EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + EXPECT_EQ(mj_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), kTreeCount, &nidof), 1); EXPECT_EQ(nidof, expected_nidof); @@ -282,9 +282,9 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { } std::vector parent(ntree); - _mjPRIVATE_dsuInit(parent.data(), ntree); + mj_dsuInit(parent.data(), ntree); for (const auto& edge : edges) { - _mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]); + mj_dsuMerge(parent.data(), edge[0], edge[1]); } std::vector active(ntree); @@ -329,7 +329,7 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { } std::vector island(ntree, -2); int nidof = -1; - const int nisland = _mjPRIVATE_dsuAssign( + const int nisland = mj_dsuAssign( island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof); SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial