chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Export disjoint-set island helpers directly

Browse Source 
Signed-off-by: teerthsharma <teerths57@gmail.com>
This commit is contained in:
teerthsharma
2026-07-20 08:59:54 +05:30
committed by teerthsharma
parent d9c8bcbc8f
commit 8d9f230514
3 changed files with 67 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_island.c
+13 -33
View File
@@ -84,7 +84,7 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) {
//-------------------------- flood-fill and graph construction ------------------------------------ //-------------------------- flood-fill and graph construction ------------------------------------
// find the canonical root of an active tree and compress its path // 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; int root = tree;
while (parent[root] != root) { while (parent[root] != root) {
root = parent[root]; root = parent[root];
@@ -101,13 +101,13 @@ static inline int dsuFind(int* parent, int tree) {
// initialize all trees as inactive // 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); mju_fillInt(parent, -1, ntree);
} }
// activate and union two incident trees; -1 denotes a static endpoint // 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) { if (tree1 == -1 && tree2 == -1) {
mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR
return; return;
@@ -121,8 +121,8 @@ static inline void dsuUnion(int* parent, int tree1, int tree2) {
if (parent[tree1] == parent[tree2]) return; if (parent[tree1] == parent[tree2]) return;
int root1 = dsuFind(parent, tree1); int root1 = mj_dsuRoot(parent, tree1);
int root2 = dsuFind(parent, tree2); int root2 = mj_dsuRoot(parent, tree2);
if (root1 < root2) { if (root1 < root2) {
parent[root2] = root1; parent[root2] = root1;
} else if (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 // assign deterministic island ids in ascending canonical-root order
static inline int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int mj_dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree,
int* nidof) { int* nidof) {
int nisland = 0; int nisland = 0;
*nidof = 0; *nidof = 0;
for (int tree=0; tree < ntree; tree++) { 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) { if (parent[tree] == tree) {
island[tree] = nisland++; island[tree] = nisland++;
} else { } 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. // ascending order, this predecessor has already been compressed and assigned an island.
parent[tree] = parent[parent[tree]]; parent[tree] = parent[parent[tree]];
island[tree] = island[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 // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
// arguments: // arguments:
// island (nr) - island index assigned to vertex, -1 if vertex has no edges // 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 // activate a singleton or union all trees in a multi-tree constraint
if (tree2 == -2) { if (tree2 == -2) {
dsuUnion(parent, tree1, -1); mj_dsuMerge(parent, tree1, -1);
} else { } else {
while (tree2 != -2) { while (tree2 != -2) {
dsuUnion(parent, tree1, tree2); mj_dsuMerge(parent, tree1, tree2);
tree1 = tree2; tree1 = tree2;
tree2 = treeNext(m, d, i, &iter); tree2 = treeNext(m, d, i, &iter);
} }
@@ -458,7 +438,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent,
if (tree1 < 0) { if (tree1 < 0) {
tree1 = tree2; tree1 = tree2;
} else { } 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 // union direct tree incidence and assign deterministic components
int* efc_tree = mjSTACKALLOC(d, nefc, int); int* efc_tree = mjSTACKALLOC(d, nefc, int);
int* parent = mjSTACKALLOC(d, ntree, int); int* parent = mjSTACKALLOC(d, ntree, int);
dsuInit(parent, ntree); mj_dsuInit(parent, ntree);
unionConstraintTrees(m, d, parent, efc_tree); unionConstraintTrees(m, d, parent, efc_tree);
int* tree_island = mjSTACKALLOC(d, ntree, int); int* tree_island = mjSTACKALLOC(d, ntree, int);
int nidof; 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 // no islands found: quick return
if (!d->nisland) { if (!d->nisland) {
src/engine/engine_island.h
+6 -5
View File
@@ -23,11 +23,12 @@
extern "C" { extern "C" {
#endif #endif
MJAPI int _mjPRIVATE_dsuFind(int* parent, int tree); // disjoint-set roots are minimum tree indices; mj_dsuRoot requires parent[tree] >= 0
MJAPI void _mjPRIVATE_dsuInit(int* parent, int ntree); MJAPI void mj_dsuInit(int* parent, int ntree);
MJAPI void _mjPRIVATE_dsuUnion(int* parent, int tree1, int tree2); MJAPI void mj_dsuMerge(int* parent, int tree1, int tree2);
MJAPI int _mjPRIVATE_dsuAssign(int* island, int* parent, MJAPI int mj_dsuRoot(int* parent, int tree);
const int* tree_dofnum, int ntree, int* nidof); 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 // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix
test/engine/engine_island_test.cc
+48 -48
View File
@@ -40,96 +40,96 @@ using IslandTest = MujocoTest;
TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) { TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) {
int parent[] = {8, 6, 7, 5}; int parent[] = {8, 6, 7, 5};
_mjPRIVATE_dsuInit(parent, 0); mj_dsuInit(parent, 0);
EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5)); EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5));
_mjPRIVATE_dsuInit(parent, 4); mj_dsuInit(parent, 4);
EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1)); EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1));
} }
TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) { TEST_F(IslandTest, DsuRootReturnsCanonicalRootAndCompressesPath) {
int parent[] = {0, 0, 1, 2, 3}; int parent[] = {0, 0, 1, 2, 3};
EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0); EXPECT_EQ(mj_dsuRoot(parent, 0), 0);
EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0); EXPECT_EQ(mj_dsuRoot(parent, 4), 0);
EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 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}; int parent[] = {-1, -1, -1, -1, -1, -1};
_mjPRIVATE_dsuUnion(parent, -1, 4); mj_dsuMerge(parent, -1, 4);
_mjPRIVATE_dsuUnion(parent, 3, -1); mj_dsuMerge(parent, 3, -1);
_mjPRIVATE_dsuUnion(parent, 5, 2); mj_dsuMerge(parent, 5, 2);
_mjPRIVATE_dsuUnion(parent, 4, 5); mj_dsuMerge(parent, 4, 5);
_mjPRIVATE_dsuUnion(parent, 3, 4); mj_dsuMerge(parent, 3, 4);
EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
for (int tree = 2; tree < 6; ++tree) { 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)); EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2));
} }
TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) { TEST_F(IslandTest, DsuMergeRedundantAndReversedEdgesAreIdempotent) {
int parent[] = {-1, -1, -1, -1}; int parent[] = {-1, -1, -1, -1};
_mjPRIVATE_dsuUnion(parent, 3, 1); mj_dsuMerge(parent, 3, 1);
_mjPRIVATE_dsuUnion(parent, 2, 1); mj_dsuMerge(parent, 2, 1);
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
_mjPRIVATE_dsuUnion(parent, 1, 3); mj_dsuMerge(parent, 1, 3);
_mjPRIVATE_dsuUnion(parent, 3, 1); mj_dsuMerge(parent, 3, 1);
_mjPRIVATE_dsuUnion(parent, 2, 2); mj_dsuMerge(parent, 2, 2);
_mjPRIVATE_dsuUnion(parent, -1, 2); mj_dsuMerge(parent, -1, 2);
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1));
} }
TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) { TEST_F(IslandTest, DsuMergeFastPathActivatesBeforeTestingParents) {
int self_parent[] = {-1, -1, -1}; 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)); EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1));
int static_first[] = {-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)); EXPECT_THAT(static_first, ElementsAre(-1, -1, 2));
int static_second[] = {-1, -1, -1}; 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)); EXPECT_THAT(static_second, ElementsAre(0, -1, -1));
} }
TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) { TEST_F(IslandTest, DsuMergeFastPathDistinguishesParentsFromRoots) {
int distinct_parent[] = {0, 0, 2, 2}; 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)); EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2));
int shared_parent[] = {0, 0, 0, 3}; 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)); EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3));
int long_paths[] = {0, 0, 1, 3, 3, 4}; 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)); 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}; int parent[] = {-1, -1, -1, -1, -1, -1};
_mjPRIVATE_dsuUnion(parent, 0, 1); mj_dsuMerge(parent, 0, 1);
_mjPRIVATE_dsuUnion(parent, 1, 2); mj_dsuMerge(parent, 1, 2);
_mjPRIVATE_dsuUnion(parent, 2, 0); mj_dsuMerge(parent, 2, 0);
_mjPRIVATE_dsuUnion(parent, 0, 2); mj_dsuMerge(parent, 0, 2);
_mjPRIVATE_dsuUnion(parent, 3, 4); mj_dsuMerge(parent, 3, 4);
_mjPRIVATE_dsuUnion(parent, 4, 5); mj_dsuMerge(parent, 4, 5);
EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3)); 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)); EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3));
} }
TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) { TEST_F(IslandTest, DsuMergeRejectsStaticSelfIncidence) {
int parent[] = {-1, 1, 1, 3}; 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"); "self-incidence of the static tree");
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3)); EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3));
} }
@@ -140,13 +140,13 @@ TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) {
const int tree_dofnum[] = {73}; const int tree_dofnum[] = {73};
int nidof = -1; 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_EQ(nidof, 0);
EXPECT_THAT(island, ElementsAre(71)); EXPECT_THAT(island, ElementsAre(71));
EXPECT_THAT(parent, ElementsAre(72)); EXPECT_THAT(parent, ElementsAre(72));
parent[0] = -1; 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_EQ(nidof, 0);
EXPECT_THAT(island, ElementsAre(-1)); EXPECT_THAT(island, ElementsAre(-1));
EXPECT_THAT(parent, 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 island[] = {9, 9, 9, 9, 9, 9, 9};
int nidof = -1; 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_EQ(nidof, 39);
EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2)); EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2));
EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6)); 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 island[] = {9, 9, 9, 9, 9, 9, 9, 9, 9, 9};
int nidof = -1; 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_EQ(nidof, 41);
EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 1, 2, 2, 2)); 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)); 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; 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), kTreeCount, &nidof),
1); 1);
EXPECT_EQ(nidof, expected_nidof); EXPECT_EQ(nidof, expected_nidof);
@@ -205,11 +205,11 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) {
std::vector<int> parent(kTreeCount); std::vector<int> parent(kTreeCount);
std::vector<int> island(kTreeCount, -2); std::vector<int> island(kTreeCount, -2);
std::vector<int> tree_dofnum(kTreeCount); std::vector<int> tree_dofnum(kTreeCount);
_mjPRIVATE_dsuInit(parent.data(), kTreeCount); mj_dsuInit(parent.data(), kTreeCount);
int expected_nidof = 0; int expected_nidof = 0;
for (int tree = kTreeCount - 1; tree > 0; --tree) { 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) { for (int tree = 0; tree < kTreeCount; ++tree) {
tree_dofnum[tree] = tree % 7; tree_dofnum[tree] = tree % 7;
@@ -217,7 +217,7 @@ TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) {
} }
int nidof = -1; 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), kTreeCount, &nidof),
1); 1);
EXPECT_EQ(nidof, expected_nidof); EXPECT_EQ(nidof, expected_nidof);
@@ -282,9 +282,9 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) {
} }
std::vector<int> parent(ntree); std::vector<int> parent(ntree);
_mjPRIVATE_dsuInit(parent.data(), ntree); mj_dsuInit(parent.data(), ntree);
for (const auto& edge : edges) { for (const auto& edge : edges) {
_mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]); mj_dsuMerge(parent.data(), edge[0], edge[1]);
} }
std::vector<int> active(ntree); std::vector<int> active(ntree);
@@ -329,7 +329,7 @@ TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) {
} }
std::vector<int> island(ntree, -2); std::vector<int> island(ntree, -2);
int nidof = -1; int nidof = -1;
const int nisland = _mjPRIVATE_dsuAssign( const int nisland = mj_dsuAssign(
island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof); island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof);
SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial