[dartsim] Ensure Link and Model APIs continue to work after joint creation in DART

dartsim/src/Base.hh

@@ -49,12 +49,6 @@ namespace dartsim {
 ///    create a std::shared_ptr of the struct that wraps the corresponding DART
 ///    shared pointer.
 
-struct ModelInfo
-{
-  dart::dynamics::SkeletonPtr model;
-  dart::dynamics::SimpleFramePtr frame;
-  std::string canonicalLinkName;
-};
 
 struct LinkInfo
 {
@@ -65,6 +59,14 @@ struct LinkInfo
   std::string name;
 };
 
+struct ModelInfo
+{
+  dart::dynamics::SkeletonPtr model;
+  dart::dynamics::SimpleFramePtr frame;
+  std::string canonicalLinkName;
+  std::vector<std::shared_ptr<LinkInfo>> links {};
+};
+
 struct JointInfo
 {
   dart::dynamics::JointPtr joint;
@@ -100,14 +102,15 @@ struct EntityStorage
   /// \brief The key represents the parent ID. The value represents a vector of
   /// the objects' IDs. The key of the vector is the object's index within its
   /// container. This is used by World and Model objects, which don't know their
-  /// own indices within their containers.
+  /// own indices within their containers as well as Links, whose indices might
+  /// change when constructing joints.
   ///
   /// The container type for World is Engine.
   /// The container type for Model is World.
+  /// The container type for Link is Model.
   ///
-  /// Links and Joints are contained in Models, but Links and Joints know their
-  /// own indices within their Models, so we do not need to use this field for
-  /// either of those types.
+  /// Joints are contained in Models, but they know their own indices within
+  /// their Models, so we do not need to use this field for Joints
   IndexMap indexInContainerToID;
 
   /// \brief Map from an entity ID to its index within its container
@@ -289,18 +292,30 @@ class Base : public Implements3d<FeatureList<Feature>>
   }
 
   public: inline std::size_t AddLink(DartBodyNode *_bn,
-                                     const std::string &_fullName)
+        const std::string &_fullName, std::size_t _modelID)
   {
     const std::size_t id = this->GetNextEntity();
-    this->links.idToObject[id] = std::make_shared<LinkInfo>();
-    this->links.idToObject[id]->link = _bn;
+    auto linkInfo = std::make_shared<LinkInfo>();
+    this->links.idToObject[id] = linkInfo;
+    linkInfo->link = _bn;
     // The name of the BodyNode during creation is assumed to be the
     // Gazebo-specified name.
-    this->links.idToObject[id]->name = _bn->getName();
+    linkInfo->name = _bn->getName();
     this->links.objectToID[_bn] = id;
     this->frames[id] = _bn;
 
     this->linksByName[_fullName] = _bn;
+    this->models.at(_modelID)->links.push_back(linkInfo);
+
+    // Even though DART keeps track of the index of this BodyNode in the
+    // skeleton, the BodyNode may be moved to another skeleton when a joint is
+    // constructed. Thus, we store the original index here.
+    this->links.idToIndexInContainer[id] = _bn->getIndexInSkeleton();
+    std::vector<std::size_t> &indexInContainerToID =
+        this->links.indexInContainerToID[_modelID];
+    indexInContainerToID.push_back(id);
+
+    this->links.idToContainerID[id] = _modelID;
 
     return id;
   }

dartsim/src/Base_TEST.cc

@@ -68,7 +68,7 @@ TEST(BaseClass, RemoveModel)
     const std::string fullName = ::sdf::JoinName(
         world->getName(),
         ::sdf::JoinName(skel->getName(), pair.second->getName()));
-    base.AddLink(pair.second, fullName);
+    base.AddLink(pair.second, fullName, std::get<0>(res));
     base.AddJoint(pair.first);
     base.AddShape({sn, name + "_shape"});
 

dartsim/src/EntityManagementFeatures.cc

@@ -272,22 +272,25 @@ std::size_t EntityManagementFeatures::GetLinkCount(
     const Identity &_modelID) const
 {
   return this->ReferenceInterface<ModelInfo>(_modelID)
-      ->model->getNumBodyNodes();
+      ->links.size();
 }
 
 /////////////////////////////////////////////////
 Identity EntityManagementFeatures::GetLink(
     const Identity &_modelID, const std::size_t _linkIndex) const
 {
-  DartBodyNode *const bn =
-      this->ReferenceInterface<ModelInfo>(_modelID)->model->getBodyNode(
-          _linkIndex);
+  auto modelInfo = this->ReferenceInterface<ModelInfo>(_modelID);
+
+  if (_linkIndex >= modelInfo->links.size())
+    return this->GenerateInvalidId();
+
+  auto linkInfo = modelInfo->links[_linkIndex];
 
   // If the link doesn't exist in "links", it means the containing entity has
   // been removed.
-  if (this->links.HasEntity(bn))
+  if (this->links.HasEntity(linkInfo->link))
   {
-    const std::size_t linkID = this->links.IdentityOf(bn);
+    const std::size_t linkID = this->links.IdentityOf(linkInfo->link);
     return this->GenerateIdentity(linkID, this->links.at(linkID));
   }
   else
@@ -304,25 +307,29 @@ Identity EntityManagementFeatures::GetLink(
 Identity EntityManagementFeatures::GetLink(
     const Identity &_modelID, const std::string &_linkName) const
 {
-  DartBodyNode *const bn =
-      this->ReferenceInterface<ModelInfo>(_modelID)->model->getBodyNode(
-          _linkName);
-
-  // If the link doesn't exist in "links", it means the containing entity has
-  // been removed.
-  if (this->links.HasEntity(bn))
-  {
-    const std::size_t linkID = this->links.IdentityOf(bn);
-    return this->GenerateIdentity(linkID, this->links.at(linkID));
-  }
-  else
+  auto modelInfo = this->ReferenceInterface<ModelInfo>(_modelID);
+  for (const auto &linkInfo : modelInfo->links)
   {
-    // TODO(addisu) It's not clear what to do when `GetLink` is called on a
-    // model that has been removed. Right now we are returning an invalid
-    // identity, but that could cause a segfault if the use doesn't check if
-    // returned value before using it.
-    return this->GenerateInvalidId();
+    if (_linkName == linkInfo->name)
+    {
+      // If the link doesn't exist in "links", it means the containing entity
+      // has been removed.
+      if (this->links.HasEntity(linkInfo->link))
+      {
+        const std::size_t linkID = this->links.IdentityOf(linkInfo->link);
+        return this->GenerateIdentity(linkID, this->links.at(linkID));
+      }
+      else
+      {
+        // TODO(addisu) It's not clear what to do when `GetLink` is called on a
+        // model that has been removed. Right now we are returning an invalid
+        // identity, but that could cause a segfault if the use doesn't check if
+        // returned value before using it.
+        return this->GenerateInvalidId();
+      }
+    }
   }
+  return this->GenerateInvalidId();
 }
 
 /////////////////////////////////////////////////
@@ -393,22 +400,19 @@ const std::string &EntityManagementFeatures::GetLinkName(
 std::size_t EntityManagementFeatures::GetLinkIndex(
     const Identity &_linkID) const
 {
-  return this->ReferenceInterface<LinkInfo>(_linkID)
-      ->link->getIndexInSkeleton();
+  return this->links.idToIndexInContainer.at(_linkID);
 }
 
 /////////////////////////////////////////////////
 Identity EntityManagementFeatures::GetModelOfLink(
     const Identity &_linkID) const
 {
-  const DartSkeletonPtr &model =
-      this->ReferenceInterface<LinkInfo>(_linkID)->link->getSkeleton();
+  const std::size_t modelID = this->links.idToContainerID.at(_linkID);
 
   // If the model containing the link doesn't exist in "models", it means this
   // link belongs to a removed model.
-  if (this->models.HasEntity(model))
+  if (this->models.HasEntity(modelID))
   {
-    const std::size_t modelID = this->models.IdentityOf(model);
     return this->GenerateIdentity(modelID, this->models.at(modelID));
   }
   else
@@ -667,7 +671,7 @@ Identity EntityManagementFeatures::ConstructEmptyLink(
       world->getName(),
       ::sdf::JoinName(model->getName(), bn->getName()));
 
-  const std::size_t linkID = this->AddLink(bn, fullName);
+  const std::size_t linkID = this->AddLink(bn, fullName, _modelID);
   return this->GenerateIdentity(linkID, this->links.at(linkID));
 }
 

dartsim/src/JointFeatures_TEST.cc

@@ -433,10 +433,8 @@ TEST_F(JointFeaturesFixture, LinkCountsInJointAttachDetach)
 
   // After attaching we expect each model to have 1 link, but the current
   // behavior is that there are 2 links in model1 and 0 in model2
-  // EXPECT_EQ(1u, model1->GetLinkCount());
-  // EXPECT_EQ(1u, model2->GetLinkCount());
-  EXPECT_EQ(2u, model1->GetLinkCount());
-  EXPECT_EQ(0u, model2->GetLinkCount());
+  EXPECT_EQ(1u, model1->GetLinkCount());
+  EXPECT_EQ(1u, model2->GetLinkCount());
 
   // now detach joint and expect model2 to start moving again
   fixedJoint->Detach();
@@ -451,7 +449,7 @@ TEST_F(JointFeaturesFixture, LinkCountsInJointAttachDetach)
 
   auto model3Body = model3->GetLink(bodyName);
   auto fixedJoint2 = model3Body->AttachFixedJoint(model2Body);
-  EXPECT_EQ(2u, model2->GetLinkCount());
+  EXPECT_EQ(1u, model2->GetLinkCount());
   fixedJoint2->Detach();
   // After detaching we expect each model to have 1 link
   EXPECT_EQ(1u, model2->GetLinkCount());

dartsim/src/SDFFeatures.cc

@@ -534,7 +534,7 @@ Identity SDFFeatures::ConstructSdfLink(
     const Identity &_modelID,
     const ::sdf::Link &_sdfLink)
 {
-  const auto &modelInfo = *this->ReferenceInterface<ModelInfo>(_modelID);
+  auto &modelInfo = *this->ReferenceInterface<ModelInfo>(_modelID);
   dart::dynamics::BodyNode::Properties bodyProperties;
   bodyProperties.mName = _sdfLink.Name();
 
@@ -590,7 +590,7 @@ Identity SDFFeatures::ConstructSdfLink(
   const std::string fullName = ::sdf::JoinName(
       world->getName(),
       ::sdf::JoinName(modelInfo.model->getName(), bn->getName()));
-  const std::size_t linkID = this->AddLink(bn, fullName);
+  const std::size_t linkID = this->AddLink(bn, fullName, _modelID);
   this->AddJoint(joint);
 
   auto linkIdentity = this->GenerateIdentity(linkID, this->links.at(linkID));

dartsim/src/SDFFeatures_TEST.cc

@@ -558,16 +558,37 @@ TEST_P(SDFFeatures_TEST, WorldWithNestedModel)
   auto nestedJoint = parentModel->GetJoint("nested_joint");
   EXPECT_NE(nullptr, nestedJoint);
 
-  EXPECT_EQ(3u, parentModel->GetLinkCount());
-  EXPECT_EQ(0u, nestedModel->GetLinkCount());
-  auto nestedLink1 = parentModel->GetLink("nested_link1");
-  EXPECT_NE(nullptr, nestedLink1);
+  EXPECT_EQ(1u, parentModel->GetLinkCount());
+  EXPECT_NE(nullptr, parentModel->GetLink("link1"));
+  EXPECT_EQ(nullptr, parentModel->GetLink("nested_link1"));
+  EXPECT_EQ(nullptr, parentModel->GetLink("nested_link2"));
 
-  auto nestedLink2 = parentModel->GetLink("nested_link2");
-  EXPECT_NE(nullptr, nestedLink2);
+  ASSERT_EQ(2u, nestedModel->GetLinkCount());
+  auto nestedLink1 = nestedModel->GetLink("nested_link1");
+  ASSERT_NE(nullptr, nestedLink1);
+  EXPECT_EQ(0u, nestedLink1->GetIndex());
+  EXPECT_EQ(nestedLink1, nestedModel->GetLink(0));
 
-  auto link1 = parentModel->GetLink("link1");
-  EXPECT_NE(nullptr, link1);
+  auto nestedLink2 = nestedModel->GetLink("nested_link2");
+  ASSERT_NE(nullptr, nestedLink2);
+  EXPECT_EQ(1u, nestedLink2->GetIndex());
+  EXPECT_EQ(nestedLink2, nestedModel->GetLink(1));
+
+  auto nestedModelSkel = dartWorld->getSkeleton("parent_model::nested_model");
+  ASSERT_NE(nullptr, nestedModelSkel);
+  // nested_model::nested_link1 would have moved to the parent_model skeleton so
+  // we expect to not find it in the nested_model skeleton
+  EXPECT_EQ(nullptr, nestedModelSkel->getBodyNode("nested_link1"));
+
+  auto nestedModel2 = world->GetModel("parent_model::nested_model2");
+  ASSERT_NE(nullptr, nestedModel2);
+  EXPECT_EQ(1u, nestedModel2->GetLinkCount());
+  EXPECT_NE(nullptr, nestedModel2->GetLink("nested_link1"));
+
+  auto nestedModel3 = world->GetModel("parent_model::nested_model2");
+  ASSERT_NE(nullptr, nestedModel3);
+  EXPECT_EQ(1u, nestedModel3->GetLinkCount());
+  EXPECT_NE(nullptr, nestedModel3->GetLink("nested_link1"));
 }
 
 /////////////////////////////////////////////////