diff --git a/doc/changelog.rst b/doc/changelog.rst index a8d3af11..f1ad20dd 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -48,6 +48,8 @@ Engine as originally guarded by ``ngravcomp``. Since the engine uses these integers as flags (zero vs. non-zero), the new flags are honest boolean properties, writeable from the Python bindings at runtime. The field ``ngravcomp`` is deprecated and will be removed in a future release. +- Replaced quadratic scratch in DFS flood-fill island discovery with a linear-memory Union-Find (disjoint set). + Contribution by :github:user:`teerthsharma`. .. admonition:: Breaking API changes :class: attention diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 47501298..499b9e32 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -16,7 +16,6 @@ #include #include -#include #include #include @@ -82,7 +81,74 @@ static int arenaAllocIsland(const mjModel* m, mjData* d) { } -//-------------------------- flood-fill and graph construction ------------------------------------ +//-------------------------- union-find and flood-fill -------------------------------------------- + +// find the canonical root of an active tree and compress its path +int mj_dsuRoot(int* parent, int tree) { + int root = tree; + while (parent[root] != root) { + root = parent[root]; + } + + while (parent[tree] != tree) { + int next = parent[tree]; + parent[tree] = root; + tree = next; + } + + return root; +} + + +// activate and union two incident trees; -1 denotes a static endpoint +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; + } + + if (tree1 == -1) tree1 = tree2; + if (tree2 == -1) tree2 = tree1; + + if (parent[tree1] == -1) parent[tree1] = tree1; + if (parent[tree2] == -1) parent[tree2] = tree2; + + if (parent[tree1] == parent[tree2]) return; + + int root1 = mj_dsuRoot(parent, tree1); + int root2 = mj_dsuRoot(parent, tree2); + if (root1 < root2) { + parent[root2] = root1; + } else if (root2 < root1) { + parent[root1] = root2; + } +} + + +// assign deterministic island ids in ascending canonical-root order +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++) { + if (parent[tree] == -1) { + island[tree] = -1; + continue; + } + + 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; +} + // find disjoint subgraphs ("islands") given sparse symmetric adjacency matrix // arguments: @@ -280,60 +346,28 @@ static void treeIterInit(const mjModel* m, const mjData* d, int i, mjTreeIter* i } -// add 0, 1 or 2 edges to uncompressed CSR adjacency matrix -// increment rownnz using tree_tree to de-dupe; return number of edges added -static int addEdge(int* rownnz, int* colind, mjtByte* tree_tree, int ntree, int tree1, int tree2) { - if (tree1 == -1 && tree2 == -1) { - mjERROR("self-edge of the static tree"); // SHOULD NOT OCCUR - return 0; - } - - // handle static trees (treat as self-edge) - if (tree1 == -1) tree1 = tree2; - if (tree2 == -1) tree2 = tree1; - - // skip if edge already present - if (tree_tree[tree1*ntree + tree2]) { - return 0; - } - - // add edge - tree_tree[tree1*ntree + tree2] = 1; - colind[tree1*ntree + rownnz[tree1]++] = tree2; // uncompressed format, rowadr is known - - // add flipped edge (off-diagonal) - if (tree1 != tree2) { - tree_tree[tree2*ntree + tree1] = 1; - colind[tree2*ntree + rownnz[tree2]++] = tree1; // uncompressed format, rowadr is known - return 2; - } - - return 1; +// return whether repeated scalar rows of this constraint require separate tree scans +static int isFlexEquality(const mjModel* m, int efc_type, int efc_id) { + return efc_type == mjCNSTR_EQUALITY && + (m->eq_type[efc_id] == mjEQ_FLEX || + m->eq_type[efc_id] == mjEQ_FLEXVERT || + m->eq_type[efc_id] == mjEQ_FLEXSTRAIN); } -// find tree-tree edges (column indices), return total number of edges -// efc_tree: first nonegative tree index of each constraint -static int findEdges(const mjModel* m, const mjData* d, - int* rownnz, int* colind, mjtByte* tree_tree, int* efc_tree, int ntree) { +// activate and union all trees with direct incidence in a constraint +static const char* unionConstraintTrees(const mjModel* m, const mjData* d, int* parent, + int* efc_tree, int* err_i) { int nefc = d->nefc; - int nnz = 0; int efc_type = -1; int efc_id = -1; - // clear row nonzeros - mju_zeroInt(rownnz, ntree); - - // iterate over constraints, compute tree-tree edges, assign efc_tree + // 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]) { + // row i is still in the same constraint: skip it + 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 (!(efc_type == mjCNSTR_EQUALITY && - (m->eq_type[efc_id] == mjEQ_FLEX || - m->eq_type[efc_id] == mjEQ_FLEXVERT || - m->eq_type[efc_id] == mjEQ_FLEXSTRAIN))) { - // copy tree assignment from previous constraint and continue + if (!isFlexEquality(m, efc_type, efc_id)) { efc_tree[i] = efc_tree[i-1]; continue; } @@ -349,25 +383,26 @@ static int findEdges(const mjModel* m, const mjData* d, int tree1 = treeNext(m, d, i, &iter); if (tree1 != -2) { int tree2 = treeNext(m, d, i, &iter); - // assign tree to constraint, one of (tree1, tree2) must be non-negative efc_tree[i] = tree1 >= 0 ? tree1 : tree2; if (efc_tree[i] < 0) { - mjERROR("constraint %d is between two static bodies", i); // SHOULD NOT OCCUR + *err_i = i; + return "constraint %d is between two static bodies"; } - // add one edge or continue to search for more edges + // activate a singleton or union all trees in a multi-tree constraint if (tree2 == -2) { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, -1); + mj_dsuMerge(parent, tree1, -1); } else { while (tree2 != -2) { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, tree2); + mj_dsuMerge(parent, tree1, tree2); tree1 = tree2; tree2 = treeNext(m, d, i, &iter); } } } else { - mjERROR("no tree found for constraint %d", i); // SHOULD NOT OCCUR + *err_i = i; + return "no tree found for constraint %d"; } } @@ -387,6 +422,7 @@ static int findEdges(const mjModel* m, const mjData* d, if (m->flex_bendingadr[f] < 0 && (sadr < 0 || m->flex_stiffness[sadr] == 0)) { continue; } + int num, adr; const int* bodyid; if (m->flex_interp[f]) { @@ -398,21 +434,22 @@ static int findEdges(const mjModel* m, const mjData* d, adr = m->flex_vertadr[f]; bodyid = m->flex_vertbodyid; } + int tree1 = -1; for (int j=0; j < num; j++) { - int treeid = m->body_treeid[bodyid[adr+j]]; - if (treeid < 0 || treeid == tree1 || !d->tree_awake[treeid]) { + int tree2 = m->body_treeid[bodyid[adr+j]]; + if (tree2 < 0 || tree2 == tree1 || !d->tree_awake[tree2]) { continue; } if (tree1 < 0) { - tree1 = treeid; + tree1 = tree2; } else { - nnz += addEdge(rownnz, colind, tree_tree, ntree, tree1, treeid); + mj_dsuMerge(parent, tree1, tree2); } } } - return nnz; + return NULL; } @@ -431,29 +468,19 @@ void mj_island(const mjModel* m, mjData* d) { mj_markStack(d); - // dense tree-tree adjacency matrix - int ntree2 = ntree * ntree; - mjtByte* tree_tree = mjSTACKALLOC(d, ntree2, mjtByte); - memset(tree_tree, 0, ntree2); - - // CSR representation of tree-tree adjacency matrix (uncompressed) - int* colind = mjSTACKALLOC(d, ntree2, int); - int* rownnz = mjSTACKALLOC(d, ntree, int); - int* rowadr = mjSTACKALLOC(d, ntree, int); - for (int r=0; r < ntree; r++) { - rowadr[r] = r * ntree; - } - - // first non-negative tree index of each constraint, used later for computing efc_island + // union direct tree incidence and assign deterministic components int* efc_tree = mjSTACKALLOC(d, nefc, int); - - // compute tree-tree adjacency matrix: fill rownnz and colind - int nnz = findEdges(m, d, rownnz, colind, tree_tree, efc_tree, ntree); - - // discover islands + int* parent = mjSTACKALLOC(d, ntree, int); + mju_fillInt(parent, -1, ntree); + int err_i = -1; + const char* err_msg = unionConstraintTrees(m, d, parent, efc_tree, &err_i); + if (err_msg) { + mj_freeStack(d); + mjERROR(err_msg, err_i); + } int* tree_island = mjSTACKALLOC(d, ntree, int); - int* stack = mjSTACKALLOC(d, nnz, int); - d->nisland = mj_floodFill(tree_island, ntree, rownnz, rowadr, colind, stack); + int nidof; + d->nisland = mj_dsuAssign(tree_island, parent, m->tree_dofnum, ntree, &nidof); // no islands found: quick return if (!d->nisland) { @@ -462,13 +489,6 @@ void mj_island(const mjModel* m, mjData* d) { return; } - // count nidof: total number of dofs in islands - int nidof = 0; - for (int i=0; i < ntree; i++) { - if (tree_island[i] >= 0) { - nidof += m->tree_dofnum[i]; - } - } d->nidof = nidof; // allocate island arrays on arena @@ -575,8 +595,8 @@ void mj_island(const mjModel* m, mjData* d) { mju_zeroInt(d->island_nf, nisland); mju_zeroInt(d->island_nefc, nisland); for (int i=0; i < nefc; i++) { - int island = tree_island[efc_tree[i]]; - d->efc_island[i] = island; + d->efc_island[i] = tree_island[efc_tree[i]]; + int island = d->efc_island[i]; d->island_nefc[island]++; switch (d->efc_type[i]) { case mjCNSTR_EQUALITY: @@ -605,17 +625,15 @@ void mj_island(const mjModel* m, mjData* d) { int ic = d->island_iefcadr[island] + island_nefc2[island]++; d->map_efc2iefc[c] = ic; d->map_iefc2efc[ic] = c; + d->iefc_type[ic] = d->efc_type[c]; + d->iefc_id[ic] = d->efc_id[c]; + d->iefc_frictionloss[ic] = d->efc_frictionloss[c]; + d->iefc_D[ic] = d->efc_D[c]; + d->iefc_R[ic] = d->efc_R[c]; } // SHOULD NOT OCCUR if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount"); - // copy position-dependent efc vectors required by solver - mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc); - mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc); - mju_gather(d->iefc_frictionloss, d->efc_frictionloss, d->map_iefc2efc, nefc); - mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, nefc); - mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, nefc); - mj_freeStack(d); } diff --git a/src/engine/engine_island.h b/src/engine/engine_island.h index 140bcdf6..dd6f92f6 100644 --- a/src/engine/engine_island.h +++ b/src/engine/engine_island.h @@ -23,6 +23,11 @@ extern "C" { #endif +// disjoint-set roots are minimum tree indices; mj_dsuRoot requires parent[tree] >= 0 +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 MJAPI int mj_floodFill(int* island, int nr, const int* rownnz, const int* rowadr, const int* colind, diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc index f520c6df..9a7d9fd2 100644 --- a/test/engine/engine_island_test.cc +++ b/test/engine/engine_island_test.cc @@ -16,6 +16,9 @@ #include "src/engine/engine_island.h" +#include +#include +#include #include #include @@ -35,6 +38,301 @@ using ::testing::NotNull; using ::testing::Pointwise; using IslandTest = MujocoTest; + + +TEST_F(IslandTest, DsuRootReturnsCanonicalRootAndCompressesPath) { + int parent[] = {0, 0, 1, 2, 3}; + + 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, DsuMergeActivatesEndpointsAndUsesMinimumRoot) { + int parent[] = {-1, -1, -1, -1, -1, -1}; + + 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(mj_dsuRoot(parent, tree), 2); + } + EXPECT_THAT(parent, ElementsAre(-1, -1, 2, 2, 2, 2)); +} + +TEST_F(IslandTest, DsuMergeRedundantAndReversedEdgesAreIdempotent) { + int parent[] = {-1, -1, -1, -1}; + mj_dsuMerge(parent, 3, 1); + mj_dsuMerge(parent, 2, 1); + EXPECT_THAT(parent, ElementsAre(-1, 1, 1, 1)); + + 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, DsuMergeFastPathActivatesBeforeTestingParents) { + int self_parent[] = {-1, -1, -1}; + mj_dsuMerge(self_parent, 1, 1); + EXPECT_THAT(self_parent, ElementsAre(-1, 1, -1)); + + int static_first[] = {-1, -1, -1}; + mj_dsuMerge(static_first, -1, 2); + EXPECT_THAT(static_first, ElementsAre(-1, -1, 2)); + + int static_second[] = {-1, -1, -1}; + mj_dsuMerge(static_second, 0, -1); + EXPECT_THAT(static_second, ElementsAre(0, -1, -1)); +} + +TEST_F(IslandTest, DsuMergeFastPathDistinguishesParentsFromRoots) { + int distinct_parent[] = {0, 0, 2, 2}; + mj_dsuMerge(distinct_parent, 1, 3); + EXPECT_THAT(distinct_parent, ElementsAre(0, 0, 0, 2)); + + int shared_parent[] = {0, 0, 0, 3}; + mj_dsuMerge(shared_parent, 1, 2); + EXPECT_THAT(shared_parent, ElementsAre(0, 0, 0, 3)); + + int long_paths[] = {0, 0, 1, 3, 3, 4}; + mj_dsuMerge(long_paths, 2, 5); + EXPECT_THAT(long_paths, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuMergeFastPathPreservesCyclesDuplicatesAndForest) { + int parent[] = {-1, -1, -1, -1, -1, -1}; + 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)); + + mj_dsuMerge(parent, 5, 0); + EXPECT_THAT(parent, ElementsAre(0, 0, 0, 0, 3, 3)); +} + +TEST_F(IslandTest, DsuMergeRejectsStaticSelfIncidence) { + int parent[] = {-1, 1, 1, 3}; + + EXPECT_EQ(MjuErrorMessageFrom(mj_dsuMerge)(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(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(mj_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(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)); +} + +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(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)); +} + +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(mj_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, -1); + std::vector island(kTreeCount, -2); + std::vector tree_dofnum(kTreeCount); + + int expected_nidof = 0; + for (int tree = kTreeCount - 1; tree > 0; --tree) { + mj_dsuMerge(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(mj_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, -1); + for (const auto& edge : edges) { + mj_dsuMerge(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 = mj_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 @@ -156,6 +454,230 @@ TEST_F(IslandTest, FloodFill3b) { EXPECT_THAT(island, ElementsAre(0, 1, 1, -1, 0, 0, 0)); } +TEST_F(IslandTest, ProductionStaticFirstAndRepeatedRows) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, 3); + ASSERT_EQ(model->nv, 3); + + // The first equality incidence is static first, then dynamic tree 0. + ASSERT_EQ(model->eq_objtype[0], mjOBJ_SITE); + int body1 = model->site_bodyid[model->eq_obj1id[0]]; + int body2 = model->site_bodyid[model->eq_obj2id[0]]; + EXPECT_EQ(model->body_treeid[body1], -1); + EXPECT_EQ(model->body_treeid[body2], 0); + + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 6); + EXPECT_EQ(data->nisland, 2); + EXPECT_EQ(data->nidof, 3); + EXPECT_EQ(data->ne, 6); + EXPECT_EQ(data->nf, 0); + EXPECT_THAT( + AsVector(data->efc_type, data->nefc), + ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, + mjCNSTR_EQUALITY, mjCNSTR_EQUALITY, mjCNSTR_EQUALITY)); + EXPECT_THAT(AsVector(data->efc_id, data->nefc), + ElementsAre(0, 0, 0, 1, 1, 1)); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 1, 1)); + EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland), + ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree), + ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), ElementsAre(0, 1, 1)); + EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_dof2idof, model->nv), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->map_idof2dof, model->nv), ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), + ElementsAre(0, 0, 0, 1, 1, 1)); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(3, 3)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0, 0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(3, 3)); + EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0, 3)); + EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc), + ElementsAre(0, 1, 2, 3, 4, 5)); + EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), + ElementsAre(0, 1, 2, 3, 4, 5)); +} + +TEST_F(IslandTest, ReportsConstraintBetweenTwoStaticBodies) { + static constexpr char xml[] = R"( + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + ASSERT_GT(data->nefc, 0); + ASSERT_EQ(data->efc_type[0], mjCNSTR_FRICTION_DOF); + + model->dof_treeid[data->efc_id[0]] = -1; + + EXPECT_EQ(MjuErrorMessageFrom(mj_island)(model.get(), data.get()), + "constraint 0 is between two static bodies"); +} + +TEST_F(IslandTest, ProductionFlexEqualityRescansRows) { + static constexpr char xml[] = R"( + + + + + + + + + +)"; + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, 3); + ASSERT_EQ(model->nv, 9); + ASSERT_EQ(model->neq, 1); + ASSERT_EQ(model->eq_type[0], mjEQ_FLEX); + ASSERT_TRUE(mj_isSparse(model.get())); + + MjDataPtr data = MakeData(model); + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 2); + auto row_trees = [&](int row) { + std::vector trees; + for (int j=0; j < data->efc_J_rownnz[row]; j++) { + int dof = data->efc_J_colind[data->efc_J_rowadr[row] + j]; + int tree = model->dof_treeid[dof]; + if (trees.empty() || trees.back() != tree) { + trees.push_back(tree); + } + } + return trees; + }; + + // Rows share one flex equality id but have different tree incidence. + EXPECT_THAT(row_trees(0), ElementsAre(0, 1)); + EXPECT_THAT(row_trees(1), ElementsAre(1, 2)); + EXPECT_THAT(AsVector(data->efc_type, data->nefc), + ElementsAre(mjCNSTR_EQUALITY, mjCNSTR_EQUALITY)); + EXPECT_THAT(AsVector(data->efc_id, data->nefc), ElementsAre(0, 0)); + EXPECT_EQ(data->nisland, 1); + EXPECT_EQ(data->nidof, 9); + EXPECT_EQ(data->ne, 2); + EXPECT_EQ(data->nf, 0); + EXPECT_THAT(AsVector(data->tree_island, model->ntree), ElementsAre(0, 0, 0)); + EXPECT_THAT(AsVector(data->island_ntree, data->nisland), ElementsAre(3)); + EXPECT_THAT(AsVector(data->island_itreeadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_itree2tree, model->ntree), + ElementsAre(0, 1, 2)); + EXPECT_THAT(AsVector(data->dof_island, model->nv), + ElementsAre(0, 0, 0, 0, 0, 0, 0, 0, 0)); + EXPECT_THAT(AsVector(data->island_nv, data->nisland), ElementsAre(9)); + EXPECT_THAT(AsVector(data->island_idofadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->island_dofadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_dof2idof, model->nv), + ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8)); + EXPECT_THAT(AsVector(data->map_idof2dof, model->nv), + ElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8)); + EXPECT_THAT(AsVector(data->efc_island, data->nefc), ElementsAre(0, 0)); + EXPECT_THAT(AsVector(data->island_ne, data->nisland), ElementsAre(2)); + EXPECT_THAT(AsVector(data->island_nf, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->island_nefc, data->nisland), ElementsAre(2)); + EXPECT_THAT(AsVector(data->island_iefcadr, data->nisland), ElementsAre(0)); + EXPECT_THAT(AsVector(data->map_efc2iefc, data->nefc), ElementsAre(0, 1)); + EXPECT_THAT(AsVector(data->map_iefc2efc, data->nefc), ElementsAre(0, 1)); +} + +TEST_F(IslandTest, BoundedArenaSupports1024Trees) { + constexpr int kTreeCount = 1024; + constexpr size_t kArenaBytes = 2 * 1024 * 1024; + std::string xml = R"( + + + + +)"; + xml.reserve(160 * kTreeCount); + for (int i=0; i < kTreeCount; i++) { + std::string name = std::to_string(i); + xml += ""; + xml += ""; + xml += ""; + if (i < 2) { + xml += ""; + } + xml += ""; + } + xml += R"( + + + + + +)"; + + char error[1024] = {}; + MjModelPtr model = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(model.get(), NotNull()) << error; + ASSERT_EQ(model->ntree, kTreeCount); + ASSERT_EQ(model->narena, kArenaBytes); + MjDataPtr data = MakeData(model); + ASSERT_THAT(data.get(), NotNull()); + + mj_fwdPosition(model.get(), data.get()); + + ASSERT_EQ(data->nefc, 3); + ASSERT_EQ(data->nisland, 1); + ASSERT_THAT(data->tree_island, NotNull()); + EXPECT_EQ(data->tree_island[0], 0); + EXPECT_EQ(data->tree_island[1], 0); + for (int tree=2; tree < kTreeCount; tree++) { + EXPECT_EQ(data->tree_island[tree], -1); + } + EXPECT_LE(data->maxuse_arena, kArenaBytes); +} + static const char* const kAbacusPath = "engine/testdata/island/abacus.xml"; TEST_F(IslandTest, Abacus) { @@ -342,6 +864,16 @@ 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);