diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 34403b11..4f53d2db 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 int dsuFind(int* parent, int tree) { +static inline int dsuFind(int* parent, int tree) { int root = tree; while (parent[root] != root) { root = parent[root]; @@ -101,13 +101,13 @@ static int dsuFind(int* parent, int tree) { // initialize all trees as inactive -static void dsuInit(int* parent, int ntree) { +static inline void dsuInit(int* parent, int ntree) { mju_fillInt(parent, -1, ntree); } // activate and union two incident trees; -1 denotes a static endpoint -static void dsuUnion(int* parent, int tree1, int tree2) { +static inline void dsuUnion(int* parent, int tree1, int tree2) { if (tree1 == -1 && tree2 == -1) { mjERROR("self-incidence of the static tree"); // SHOULD NOT OCCUR return; @@ -119,6 +119,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) { if (parent[tree1] == -1) parent[tree1] = tree1; if (parent[tree2] == -1) parent[tree2] = tree2; + if (parent[tree1] == parent[tree2]) return; + int root1 = dsuFind(parent, tree1); int root2 = dsuFind(parent, tree2); if (root1 < root2) { @@ -130,7 +132,8 @@ static void dsuUnion(int* parent, int tree1, int tree2) { // assign deterministic island ids in ascending canonical-root order -static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree, int* nidof) { +static inline int 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++) { @@ -139,14 +142,41 @@ static int dsuAssign(int* island, int* parent, const int* tree_dofnum, int ntree continue; } - int root = dsuFind(parent, tree); - island[tree] = root == tree ? nisland++ : island[root]; + if (parent[tree] == tree) { + island[tree] = nisland++; + } else { + // 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]]; + } *nidof += tree_dofnum[tree]; } return nisland; } + +// 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 @@ -353,7 +383,7 @@ static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) { // activate and union all trees with direct incidence in a constraint -static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) { +static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, int* efc_tree) { int nefc = d->nefc; int efc_type = -1; int efc_id = -1; @@ -361,9 +391,10 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) // iterate over constraints and union incident trees for (int i=0; i < nefc; i++) { // row i is still in the same constraint: skip it - if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { + if (i > 0 && efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) { // unless it is a flex equality, where the tree pattern changes per dof if (!isFlexEquality(m, efc_type, efc_id)) { + efc_tree[i] = efc_tree[i-1]; continue; } } @@ -378,6 +409,7 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) int tree1 = treeNext(m, d, i, &iter); if (tree1 != -2) { int tree2 = treeNext(m, d, i, &iter); + efc_tree[i] = tree1 == -1 ? tree2 : tree1; // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { @@ -393,35 +425,41 @@ static void unionConstraintTrees(const mjModel* m, const mjData* d, int* parent) mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR } } -} -// assign each constraint from its first non-negative incident tree -static void assignConstraintIslands(const mjModel* m, mjData* d, const int* tree_island) { - int efc_type = -1; - int efc_id = -1; - - for (int i=0; i < d->nefc; i++) { - // reuse assignment for repeated scalar rows, except flex equality rows - if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i] && - !isFlexEquality(m, efc_type, efc_id)) { - d->efc_island[i] = d->efc_island[i-1]; + // Flex stiffness couples all vertices (nodes for interpolated flexes) without a constraint + // row representing the coupling. Union the awake dynamic trees of each stiffness-active flex. + for (int f=0; f < m->nflex; f++) { + if (m->flex_rigid[f] || m->flex_dim[f] < 2) { + continue; + } + int sadr = m->flex_stiffnessadr[f]; + if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) { continue; } - efc_type = d->efc_type[i]; - efc_id = d->efc_id[i]; - mjTreeIter iter; - treeIterInit(m, d, i, &iter); - - int tree; - do { - tree = treeNext(m, d, i, &iter); - } while (tree == -1); - - if (tree == -2) { - mjERROR("no dynamic tree found for constraint %d", i); // SHOULD NOT OCCUR + int num, adr; + const int* bodyid; + if (m->flex_interp[f]) { + num = m->flex_nodenum[f]; + adr = m->flex_nodeadr[f]; + bodyid = m->flex_nodebodyid; } else { - d->efc_island[i] = tree_island[tree]; + num = m->flex_vertnum[f]; + adr = m->flex_vertadr[f]; + bodyid = m->flex_vertbodyid; + } + + int tree1 = -1; + for (int j=0; j < num; j++) { + int tree2 = m->body_treeid[bodyid[adr+j]]; + if (tree2 < 0 || tree2 == tree1 || !d->tree_awake[tree2]) { + continue; + } + if (tree1 < 0) { + tree1 = tree2; + } else { + dsuUnion(parent, tree1, tree2); + } } } } @@ -443,9 +481,10 @@ void mj_island(const mjModel* m, mjData* d) { mj_markStack(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); - unionConstraintTrees(m, d, parent); + 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); @@ -558,14 +597,12 @@ void mj_island(const mjModel* m, mjData* d) { // ------------------------------------- constraints --------------------------------------------- - // compute efc_island from first non-negative tree of each constraint - assignConstraintIslands(m, d, tree_island); - // compute efc_island, island_{ne,nf,nefc} mju_zeroInt(d->island_ne, nisland); mju_zeroInt(d->island_nf, nisland); mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { + d->efc_island[i] = tree_island[efc_tree[i]]; int island = d->efc_island[i]; d->island_nefc[island]++; switch (d->efc_type[i]) { diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 140bcdf6..7e239b40 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -23,6 +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); + // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, diff --git a/test/benchmark/CMakeLists.txt b/test/benchmark/CMakeLists.txt index b9aff09f..6821aa03 100644 --- a/test/benchmark/CMakeLists.txt +++ b/test/benchmark/CMakeLists.txt @@ -79,6 +79,12 @@ mujoco_test( ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers ) +mujoco_test( + island_benchmark_test + MAIN_TARGET benchmark::benchmark_main + ADDITIONAL_LINK_LIBRARIES benchmark::benchmark absl::core_headers +) + mujoco_test( engine_util_sparse_benchmark_test MAIN_TARGET benchmark::benchmark_main diff --git a/test/benchmark/island_benchmark_test.cc b/test/benchmark/island_benchmark_test.cc new file mode 100644 index 00000000..21a922b8 --- /dev/null +++ b/test/benchmark/island_benchmark_test.cc @@ -0,0 +1,374 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Benchmarks island discovery on deterministic geodesic Rips graphs. The +// corpus spans the connectivity transition of points sampled on S^2 and adds +// MuJoCo-relevant static and repeated incidences. Corpus construction and +// validation are deliberately outside the timed region. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "src/engine/engine_island.h" + +namespace mujoco { +namespace { + +struct Edge { + int first; + int second; +}; + +struct GraphCase { + std::string name; + int node_count; + std::vector incidences; + std::vector expected_partition; + int active_nodes; + int expected_components; + int pre_bridge_components; + bool bridge_added; + std::uint64_t expected_checksum; +}; + +struct Point { + double x; + double y; + double z; +}; + +std::uint64_t SplitMix64(std::uint64_t& state) { + state += 0x9e3779b97f4a7c15ULL; + std::uint64_t value = state; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); +} + +double Uniform01(std::uint64_t& state) { + return static_cast(SplitMix64(state) >> 11) * 0x1.0p-53; +} + +std::vector SampleSphere(int count, std::uint64_t seed) { + constexpr double kTwoPi = 6.283185307179586476925286766559; + std::vector points; + points.reserve(count); + for (int i = 0; i < count; ++i) { + const double z = 2.0 * Uniform01(seed) - 1.0; + const double angle = kTwoPi * Uniform01(seed); + const double radial = std::sqrt(std::max(0.0, 1.0 - z*z)); + points.push_back({radial * std::cos(angle), radial * std::sin(angle), z}); + } + return points; +} + +double Dot(const Point& a, const Point& b) { + return a.x*b.x + a.y*b.y + a.z*b.z; +} + +std::vector RipsEdges(const std::vector& points, double target_degree) { + const double probability = target_degree / (points.size() - 1); + const double radius = 2.0 * std::asin(std::sqrt(probability)); + const double minimum_dot = std::cos(radius); + std::vector edges; + for (int i = 0; i < static_cast(points.size()); ++i) { + for (int j = i + 1; j < static_cast(points.size()); ++j) { + if (Dot(points[i], points[j]) >= minimum_dot) { + edges.push_back({i, j}); + } + } + } + return edges; +} + +std::vector CanonicalPartition(int node_count, const std::vector& edges) { + std::vector> adjacency(node_count); + std::vector active(node_count, false); + for (const Edge& edge : edges) { + active[edge.first] = true; + active[edge.second] = true; + if (edge.first != edge.second) { + adjacency[edge.first].push_back(edge.second); + adjacency[edge.second].push_back(edge.first); + } + } + + std::vector partition(node_count, -1); + std::queue pending; + for (int start = 0; start < node_count; ++start) { + if (!active[start] || partition[start] != -1) { + continue; + } + partition[start] = start; + pending.push(start); + while (!pending.empty()) { + const int node = pending.front(); + pending.pop(); + for (int neighbor : adjacency[node]) { + if (partition[neighbor] == -1) { + partition[neighbor] = start; + pending.push(neighbor); + } + } + } + } + return partition; +} + +int CountComponents(const std::vector& partition) { + int count = 0; + for (int node = 0; node < static_cast(partition.size()); ++node) { + count += partition[node] == node; + } + return count; +} + +std::uint64_t PartitionChecksum(const std::vector& partition) { + std::uint64_t hash = 1469598103934665603ULL; + for (int value : partition) { + hash ^= static_cast(value); + hash *= 1099511628211ULL; + } + return hash; +} + +std::vector CanonicalizeLabels(const std::vector& labels) { + std::vector minimum(labels.size(), std::numeric_limits::max()); + for (int node = 0; node < static_cast(labels.size()); ++node) { + if (labels[node] >= 0) { + minimum[labels[node]] = std::min(minimum[labels[node]], node); + } + } + std::vector canonical(labels.size(), -1); + for (int node = 0; node < static_cast(labels.size()); ++node) { + if (labels[node] >= 0) { + canonical[node] = minimum[labels[node]]; + } + } + return canonical; +} + +void DeterministicShuffle(std::vector& edges, std::uint64_t seed) { + for (std::size_t i = edges.size(); i > 1; --i) { + const std::size_t j = SplitMix64(seed) % i; + std::swap(edges[i - 1], edges[j]); + } +} + +bool AddCriticalBridge(const std::vector& points, std::vector& edges) { + const std::vector partition = CanonicalPartition(points.size(), edges); + if (CountComponents(partition) < 2) { + return false; + } + + double best_dot = -2.0; + Edge bridge{-1, -1}; + for (int i = 0; i < static_cast(points.size()); ++i) { + for (int j = i + 1; j < static_cast(points.size()); ++j) { + if (partition[i] >= 0 && partition[j] >= 0 && partition[i] != partition[j] && + Dot(points[i], points[j]) > best_dot) { + best_dot = Dot(points[i], points[j]); + bridge = {i, j}; + } + } + } + if (bridge.first >= 0) { + edges.push_back(bridge); + return true; + } + return false; +} + +GraphCase MakeCase(std::string name, int node_count, double target_degree, + std::uint64_t seed, bool critical_bridge, bool static_rows, + bool repeated_rows) { + const std::vector points = SampleSphere(node_count, seed); + std::vector edges = RipsEdges(points, target_degree); + const int pre_bridge_components = + critical_bridge ? CountComponents(CanonicalPartition(node_count, edges)) : -1; + bool bridge_added = false; + if (critical_bridge) { + bridge_added = AddCriticalBridge(points, edges); + } + + const std::vector unique_edges = edges; + if (static_rows) { + for (int node = 0; node < node_count; node += 17) { + edges.push_back({node, node}); + } + } + if (repeated_rows) { + for (std::size_t i = 0; i < unique_edges.size(); i += 11) { + edges.push_back(unique_edges[i]); + edges.push_back({unique_edges[i].second, unique_edges[i].first}); + } + } + DeterministicShuffle(edges, seed ^ 0xd1b54a32d192ed03ULL); + + std::vector expected = CanonicalPartition(node_count, edges); + const int active_nodes = std::count_if(expected.begin(), expected.end(), + [](int component) { return component >= 0; }); + const int components = CountComponents(expected); + const std::uint64_t checksum = PartitionChecksum(expected); + return {std::move(name), node_count, std::move(edges), std::move(expected), active_nodes, + components, pre_bridge_components, bridge_added, checksum}; +} + +const std::vector& Corpus() { + static const std::vector corpus = { + MakeCase("StableSparse_S2Rips_64", 64, 2.0, 0x33960001ULL, false, false, false), + MakeCase("CriticalBridge_S2Rips_256", 256, 0.75 * std::log(256.0), + 0x33960002ULL, true, false, false), + MakeCase("SupercriticalDense_S2Rips_256", 256, 2.0 * std::ceil(std::log(256.0)), + 0x33960003ULL, false, false, false), + MakeCase("GroundedStaticRepeated_S2Rips_256", 256, + 2.0 * std::ceil(std::log(256.0)), 0x33960004ULL, false, true, true), + MakeCase("StableRepeated_S2Rips_1024", 1024, 2.0, 0x33960005ULL, + false, false, true), + MakeCase("CriticalLarge_S2Rips_1024", 1024, std::ceil(std::log(1024.0)), + 0x33960006ULL, true, false, false), + }; + return corpus; +} + +struct FloodFillWorkspace { + explicit FloodFillWorkspace(int node_count) + : adjacency(node_count * node_count), rownnz(node_count), rowadr(node_count), + colind(node_count * node_count), stack(node_count * node_count + node_count), + island(node_count) {} + + std::vector adjacency; + std::vector rownnz; + std::vector rowadr; + std::vector colind; + std::vector stack; + std::vector island; +}; + +int RunFloodFill(const GraphCase& graph, FloodFillWorkspace& work) { + const int n = graph.node_count; + std::fill(work.adjacency.begin(), work.adjacency.end(), 0); + std::fill(work.rownnz.begin(), work.rownnz.end(), 0); + for (const Edge& edge : graph.incidences) { + work.adjacency[edge.first*n + edge.second] = 1; + work.adjacency[edge.second*n + edge.first] = 1; + } + + int address = 0; + for (int row = 0; row < n; ++row) { + work.rowadr[row] = address; + for (int column = 0; column < n; ++column) { + if (work.adjacency[row*n + column]) { + work.colind[address++] = column; + ++work.rownnz[row]; + } + } + } + return mj_floodFill(work.island.data(), n, work.rownnz.data(), work.rowadr.data(), + work.colind.data(), work.stack.data()); +} + +struct DsuWorkspace { + explicit DsuWorkspace(int node_count) + : parent(node_count), island(node_count), dof_count(node_count, 1) {} + + std::vector parent; + std::vector island; + std::vector dof_count; +}; + +int RunDsu(const GraphCase& graph, DsuWorkspace& work) { + _mjPRIVATE_dsuInit(work.parent.data(), graph.node_count); + for (const Edge& edge : graph.incidences) { + _mjPRIVATE_dsuUnion(work.parent.data(), edge.first, edge.second); + } + int dof_count = 0; + return _mjPRIVATE_dsuAssign(work.island.data(), work.parent.data(), work.dof_count.data(), + graph.node_count, &dof_count); +} + +bool Validate(const GraphCase& graph) { + FloodFillWorkspace flood(graph.node_count); + DsuWorkspace dsu(graph.node_count); + const int flood_components = RunFloodFill(graph, flood); + const int dsu_components = RunDsu(graph, dsu); + const bool bridge_valid = graph.pre_bridge_components < 0 || + (graph.bridge_added && + graph.pre_bridge_components == graph.expected_components + 1); + return bridge_valid && flood_components == graph.expected_components && + dsu_components == graph.expected_components && + CanonicalizeLabels(flood.island) == graph.expected_partition && + CanonicalizeLabels(dsu.island) == graph.expected_partition && + PartitionChecksum(graph.expected_partition) == graph.expected_checksum; +} + +void BM_FloodFill(benchmark::State& state, const GraphCase* graph) { + if (!Validate(*graph)) { + state.SkipWithError("invalid S2-Rips graph fixture"); + return; + } + FloodFillWorkspace work(graph->node_count); + state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + + " active=" + std::to_string(graph->active_nodes) + + " components=" + std::to_string(graph->expected_components) + + " pre_bridge=" + std::to_string(graph->pre_bridge_components) + + " bridge_added=" + std::to_string(graph->bridge_added) + + " checksum=" + std::to_string(graph->expected_checksum)); + for (auto _ : state) { + int components = RunFloodFill(*graph, work); + benchmark::DoNotOptimize(components); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * graph->incidences.size()); +} + +void BM_Dsu(benchmark::State& state, const GraphCase* graph) { + if (!Validate(*graph)) { + state.SkipWithError("invalid S2-Rips graph fixture"); + return; + } + DsuWorkspace work(graph->node_count); + state.SetLabel("edges=" + std::to_string(graph->incidences.size()) + + " active=" + std::to_string(graph->active_nodes) + + " components=" + std::to_string(graph->expected_components) + + " pre_bridge=" + std::to_string(graph->pre_bridge_components) + + " bridge_added=" + std::to_string(graph->bridge_added) + + " checksum=" + std::to_string(graph->expected_checksum)); + for (auto _ : state) { + int components = RunDsu(*graph, work); + benchmark::DoNotOptimize(components); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * graph->incidences.size()); +} + +const bool kRegistered = [] { + for (const GraphCase& graph : Corpus()) { + benchmark::RegisterBenchmark(("Island/FloodFill/" + graph.name).c_str(), BM_FloodFill, &graph); + benchmark::RegisterBenchmark(("Island/DSU/" + graph.name).c_str(), BM_Dsu, &graph); + } + return true; +}(); + +} // namespace +} // namespace mujoco diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index 3a4462ab..be1f2832 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -16,6 +16,8 @@ #include "src/engine/engine_island.h" +#include +#include #include #include @@ -35,6 +37,310 @@ using ::testing::NotNull; using ::testing::Pointwise; using IslandTest = MujocoTest; +TEST_F(IslandTest, DsuInitHandlesEmptyAndNonemptyRanges) { + int parent[] = {8, 6, 7, 5}; + + _mjPRIVATE_dsuInit(parent, 0); + EXPECT_THAT(parent, ElementsAre(8, 6, 7, 5)); + + _mjPRIVATE_dsuInit(parent, 4); + EXPECT_THAT(parent, ElementsAre(-1, -1, -1, -1)); +} + +TEST_F(IslandTest, DsuFindReturnsCanonicalRootAndCompressesPath) { + int parent[] = {0, 0, 1, 2, 3}; + + EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 0), 0); + EXPECT_EQ(_mjPRIVATE_dsuFind(parent, 4), 0); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 0)); +} + +TEST_F(IslandTest, DsuUnionActivatesEndpointsAndUsesMinimumRoot) { + 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); + + 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_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); +} + +TEST_F(IslandTest, DsuUnionRedundantAndReversedEdgesAreIdempotent) { + int parent[] = {-1, -1, -1, -1}; + _mjPRIVATE_dsuUnion(parent, 3, 1); + _mjPRIVATE_dsuUnion(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); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); +} + +TEST_F(IslandTest, DsuUnionFastPathActivatesBeforeTestingParents) { + int self_parent[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(self_parent, 1, 1); + EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1)); + + int static_first[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(static_first, -1, 2); + EXPECT_THAT(static_first, ElementsAre(-1, -1, 2)); + + int static_second[] = {-1, -1, -1}; + _mjPRIVATE_dsuUnion(static_second, 0, -1); + EXPECT_THAT(static_second, ElementsAre(0, -1, -1)); +} + +TEST_F(IslandTest, DsuUnionFastPathDistinguishesParentsFromRoots) { + int distinct_parent[] = {0, 0, 2, 2}; + _mjPRIVATE_dsuUnion(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); + EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3)); + + int long_paths[] = {0, 0, 1, 3, 3, 4}; + _mjPRIVATE_dsuUnion(long_paths, 2, 5); + EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuUnionFastPathPreservesCyclesDuplicatesAndForest) { + 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); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 3, 3, 3)); + + _mjPRIVATE_dsuUnion(parent, 5, 0); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuUnionRejectsStaticSelfIncidence) { + int parent[] = {-1, 1, 1, 3}; + + EXPECT_EQ(MjuErrorMessageFrom(_mjPRIVATE_dsuUnion)(parent, -1, -1), + "self-incidence of the static tree"); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 3)); +} + +TEST_F(IslandTest, DsuAssignHandlesEmptyAndInactiveInputs) { + int island[] = {71}; + int parent[] = {72}; + const int tree_dofnum[] = {73}; + int nidof = -1; + + EXPECT_EQ(_mjPRIVATE_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(nidof, 0); + EXPECT_THAT(island, ElementsAre(-1)); + EXPECT_THAT(parent, ElementsAre(-1)); +} + +TEST_F(IslandTest, DsuAssignLabelsComponentsAndCountsOnlyActiveDofs) { + int parent[] = {-1, 1, 1, 2, 4, 4, 6}; + const int tree_dofnum[] = {1000, 0, 3, 5, 7, 11, 13}; + 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(nidof, 39); + EXPECT_THAT(island, ElementsAre(-1, 0, 0, 0, 1, 1, 2)); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1, 4, 4, 6)); +} + +TEST_F(IslandTest, DsuAssignCompressesAscendingMultiHopForest) { + int parent[] = {-1, 1, 1, 2, 4, 4, 5, 7, 7, 8}; + const int tree_dofnum[] = {99, 0, 2, 3, 0, 5, 7, 11, 0, 13}; + 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(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)); +} + +TEST_F(IslandTest, DsuAssignCompresses4096NodeAdversarialChain) { + constexpr int kTreeCount = 4096; + std::vector parent(kTreeCount); + std::vector island(kTreeCount, -2); + std::vector tree_dofnum(kTreeCount); + parent[0] = 0; + int expected_nidof = 0; + for (int tree = 1; tree < kTreeCount; ++tree) { + parent[tree] = tree - 1; + tree_dofnum[tree] = tree % 5; + expected_nidof += tree_dofnum[tree]; + } + + int nidof = -1; + EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + kTreeCount, &nidof), + 1); + EXPECT_EQ(nidof, expected_nidof); + for (int tree = 0; tree < kTreeCount; ++tree) { + EXPECT_EQ(island[tree], 0); + EXPECT_EQ(parent[tree], 0); + } +} + +TEST_F(IslandTest, DsuHandlesLongConnectedBoundaryCase) { + constexpr int kTreeCount = 4096; + std::vector parent(kTreeCount); + std::vector island(kTreeCount, -2); + std::vector tree_dofnum(kTreeCount); + _mjPRIVATE_dsuInit(parent.data(), kTreeCount); + + int expected_nidof = 0; + for (int tree = kTreeCount - 1; tree > 0; --tree) { + _mjPRIVATE_dsuUnion(parent.data(), tree, tree - 1); + } + for (int tree = 0; tree < kTreeCount; ++tree) { + tree_dofnum[tree] = tree % 7; + expected_nidof += tree_dofnum[tree]; + } + + int nidof = -1; + EXPECT_EQ(_mjPRIVATE_dsuAssign(island.data(), parent.data(), tree_dofnum.data(), + kTreeCount, &nidof), + 1); + EXPECT_EQ(nidof, expected_nidof); + for (int tree = 0; tree < kTreeCount; ++tree) { + EXPECT_EQ(island[tree], 0); + EXPECT_EQ(parent[tree], 0); + } +} + +TEST_F(IslandTest, DsuRandomizedDifferentialAgainstGraphTraversal) { + constexpr uint32_t kSeed = 0x5eed3396u; + constexpr int kTrials = 2000; + uint32_t state = kSeed; + auto next = [&state]() { + state = state * 1664525u + 1013904223u; + return state; + }; + + for (int trial = 0; trial < kTrials; ++trial) { + const int ntree = 1 + next() % 64; + const int nedge = next() % 192; + std::vector> edges; + edges.reserve(nedge); + for (int edge = 0; edge < nedge; ++edge) { + int tree1; + int tree2; + switch (next() % 8) { + case 0: + tree1 = -1; + tree2 = next() % ntree; + break; + case 1: + tree1 = next() % ntree; + tree2 = -1; + break; + case 2: + tree1 = next() % ntree; + tree2 = tree1; + break; + case 3: + if (!edges.empty()) { + const auto& previous = edges[next() % edges.size()]; + tree1 = previous[0]; + tree2 = previous[1]; + break; + } + [[fallthrough]]; + case 4: + if (!edges.empty()) { + const auto& previous = edges[next() % edges.size()]; + tree1 = previous[1]; + tree2 = previous[0]; + break; + } + [[fallthrough]]; + default: + tree1 = next() % ntree; + tree2 = next() % ntree; + break; + } + edges.push_back({tree1, tree2}); + } + + std::vector parent(ntree); + _mjPRIVATE_dsuInit(parent.data(), ntree); + for (const auto& edge : edges) { + _mjPRIVATE_dsuUnion(parent.data(), edge[0], edge[1]); + } + + std::vector active(ntree); + std::vector> adjacent(ntree); + for (const auto& edge : edges) { + if (edge[0] >= 0) active[edge[0]] = 1; + if (edge[1] >= 0) active[edge[1]] = 1; + if (edge[0] >= 0 && edge[1] >= 0) { + adjacent[edge[0]].push_back(edge[1]); + adjacent[edge[1]].push_back(edge[0]); + } + } + + std::vector expected_island(ntree, -1); + std::vector expected_parent(ntree, -1); + int expected_nisland = 0; + for (int start = 0; start < ntree; ++start) { + if (!active[start] || expected_island[start] != -1) continue; + std::vector pending = {start}; + std::vector component; + expected_island[start] = expected_nisland; + while (!pending.empty()) { + const int tree = pending.back(); + pending.pop_back(); + component.push_back(tree); + for (int neighbor : adjacent[tree]) { + if (expected_island[neighbor] == -1) { + expected_island[neighbor] = expected_nisland; + pending.push_back(neighbor); + } + } + } + for (int tree : component) expected_parent[tree] = start; + ++expected_nisland; + } + + std::vector tree_dofnum(ntree); + int expected_nidof = 0; + for (int tree = 0; tree < ntree; ++tree) { + tree_dofnum[tree] = next() % 8; + if (active[tree]) expected_nidof += tree_dofnum[tree]; + } + std::vector island(ntree, -2); + int nidof = -1; + const int nisland = _mjPRIVATE_dsuAssign( + island.data(), parent.data(), tree_dofnum.data(), ntree, &nidof); + + SCOPED_TRACE(::testing::Message() << "seed=" << kSeed << " trial=" << trial + << " ntree=" << ntree << " nedge=" << nedge); + EXPECT_EQ(nisland, expected_nisland); + EXPECT_EQ(nidof, expected_nidof); + EXPECT_EQ(island, expected_island); + EXPECT_EQ(parent, expected_parent); + } +} + TEST_F(IslandTest, FloodFillSingleton) { // adjacency matrix for the graph 0 1 2 // U U @@ -532,6 +838,13 @@ TEST_F(IslandTest, IslandEfc) { EXPECT_EQ(data->nf, 2); EXPECT_EQ(data->nl, 1); EXPECT_EQ(data->nefc, 30); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(1, 0, 0, 6)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 1, 1, 0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(6, 17, 1, 6)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), + ElementsAre(0, 3, 3, 3, 3, 3, 3, 1, 2, 0, + 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1)); mj_deleteData(data); mj_deleteModel(model);