Fix entity by index lookup, address reviewer feedback

Signed-off-by: Addisu Z. Taddese <[email protected]>

tpe/lib/src/Model.cc

@@ -16,7 +16,6 @@
 */
 
 #include <set>
-#include <algorithm>
 #include <string>
 
 #include <ignition/common/Profiler.hh>
@@ -39,7 +38,7 @@ class ignition::physics::tpelib::ModelPrivate
   /// \brief Links in the model
   public: std::vector<std::size_t> linkIds;
 
-  /// \brief Nested models links
+  /// \brief Nested models
   public: std::vector<std::size_t> nestedModelIds;
 };
 

tpe/plugin/src/Base.hh

@@ -88,26 +88,38 @@ class Base : public Implements3d<FeatureList<Feature>>
     return -1;
   }
 
-  public: inline std::size_t indexInContainerToId(
-    const std::size_t _containerId, const std::size_t _index) const
+  public: template <typename EntityType>
+  inline std::pair<std::size_t, EntityType> indexInContainerToId(
+      const std::size_t _containerId,
+      const std::size_t _index,
+      const std::map<std::size_t, EntityType> &_idMap) const
   {
     std::size_t counter = 0;
     auto it = this->childIdToParentId.begin();
 
     while (counter <= _index && it != this->childIdToParentId.end())
     {
-      if (it->second == _containerId && counter == _index)
+      if (it->second == _containerId)
       {
-        return it->first;
-      }
-      else if (it->second == _containerId)
-      {
-        ++counter;
+        auto idMapIt = _idMap.find(it->first);
+        // only count if the entity is found in the idMap. This makes sure we're
+        // counting only the entities with the correct EntityType
+        if (idMapIt != _idMap.end())
+        {
+          if (counter == _index)
+          {
+            return *idMapIt;
+          }
+          else
+          {
+            ++counter;
+          }
+        }
       }
       ++it;
     }
     // return invalid id if entity not found
-    return -1;
+    return {INVALID_ENTITY_ID, nullptr};
   }
 
 

tpe/plugin/src/Base_TEST.cc

@@ -158,24 +158,28 @@ TEST(BaseClass, AddEntities)
   // check indices
   std::size_t modelInd1 = base.idToIndexInContainer(modelId1);
   EXPECT_EQ(0u, modelInd1);
-  EXPECT_EQ(modelId1, base.indexInContainerToId(worldId, 0u));
+  EXPECT_EQ(modelId1,
+            base.indexInContainerToId(worldId, 0u, base.models).first);
   std::size_t modelInd2 = base.idToIndexInContainer(modelId2);
   EXPECT_EQ(1u, modelInd2);
-  EXPECT_EQ(modelId2, base.indexInContainerToId(worldId, 1u));
+  EXPECT_EQ(modelId2,
+            base.indexInContainerToId(worldId, 1u, base.models).first);
 
   std::size_t linkInd1 = base.idToIndexInContainer(linkId1);
   EXPECT_EQ(0u, linkInd1);
-  EXPECT_EQ(linkId1, base.indexInContainerToId(modelId1, 0u));
+  EXPECT_EQ(linkId1, base.indexInContainerToId(modelId1, 0u, base.links).first);
   std::size_t linkInd2 = base.idToIndexInContainer(linkId2);
   EXPECT_EQ(0u, linkInd2);
-  EXPECT_EQ(linkId2, base.indexInContainerToId(modelId2, 0u));
+  EXPECT_EQ(linkId2, base.indexInContainerToId(modelId2, 0u, base.links).first);
 
   std::size_t boxInd = base.idToIndexInContainer(boxId);
   EXPECT_EQ(0u, boxInd);
-  EXPECT_EQ(boxId, base.indexInContainerToId(linkId1, 0u));
+  EXPECT_EQ(boxId,
+            base.indexInContainerToId(linkId1, 0u, base.collisions).first);
   std::size_t cylinderInd = base.idToIndexInContainer(cylinderId);
   EXPECT_EQ(0u, cylinderInd);
-  EXPECT_EQ(cylinderId, base.indexInContainerToId(linkId2, 0u));
+  EXPECT_EQ(cylinderId,
+            base.indexInContainerToId(linkId2, 0u, base.collisions).first);
 }
 
 int main(int argc, char *argv[])

tpe/plugin/src/EntityManagementFeatures.cc

@@ -110,11 +110,11 @@ std::size_t EntityManagementFeatures::GetModelCount(
 Identity EntityManagementFeatures::GetModel(
   const Identity &_worldID, const std::size_t _modelIndex) const
 {
-  std::size_t modelId = this->indexInContainerToId(_worldID.id, _modelIndex);
-  auto it = this->models.find(modelId);
-  if (it != this->models.end() && it->second != nullptr)
+  const auto &[modelId, modelInfo] =
+      this->indexInContainerToId(_worldID.id, _modelIndex, this->models);
+  if (modelInfo != nullptr)
   {
-    return this->GenerateIdentity(modelId, it->second);
+    return this->GenerateIdentity(modelId, modelInfo);
   }
   return this->GenerateInvalidId();
 }
@@ -183,11 +183,11 @@ std::size_t EntityManagementFeatures::GetNestedModelCount(
 Identity EntityManagementFeatures::GetNestedModel(
   const Identity &_modelID, const std::size_t _modelIndex) const
 {
-  std::size_t modelId = this->indexInContainerToId(_modelID.id, _modelIndex);
-  auto it = this->models.find(modelId);
-  if (it != this->models.end() && it->second != nullptr)
+  const auto &[nestedModelId, nestedModelInfo] =
+      this->indexInContainerToId(_modelID.id, _modelIndex, this->models);
+  if (nestedModelInfo != nullptr)
   {
-    return this->GenerateIdentity(modelId, it->second);
+    return this->GenerateIdentity(nestedModelId, nestedModelInfo);
   }
   return this->GenerateInvalidId();
 }
@@ -196,20 +196,14 @@ Identity EntityManagementFeatures::GetNestedModel(
 Identity EntityManagementFeatures::GetNestedModel(
   const Identity &_modelID, const std::string &_modelName) const
 {
-  auto modelInfo = this->ReferenceInterface<WorldInfo>(_modelID);
+  auto modelInfo = this->ReferenceInterface<ModelInfo>(_modelID);
   if (modelInfo != nullptr)
   {
-    tpelib::Entity &modelEnt = modelInfo->world->GetChildByName(_modelName);
-    for (const auto &[parentModelId, nestedModelInfo]: this->models)
+    tpelib::Entity &modelEnt = modelInfo->model->GetChildByName(_modelName);
+    auto it = this->models.find(modelEnt.GetId());
+    if (it != this->models.end() && it->second != nullptr)
     {
-      if (nestedModelInfo != nullptr)
-      {
-        const std::string &name = nestedModelInfo->model->GetName();
-        if (parentModelId == modelEnt.GetId() && name == modelEnt.GetName())
-        {
-          return this->GenerateIdentity(parentModelId, nestedModelInfo);
-        }
-      }
+      return this->GenerateIdentity(it->first, it->second);
     }
   }
   return this->GenerateInvalidId();
@@ -226,11 +220,11 @@ std::size_t EntityManagementFeatures::GetLinkCount(
 Identity EntityManagementFeatures::GetLink(
   const Identity &_modelID, const std::size_t _linkIndex) const
 {
-  std::size_t linkId = this->indexInContainerToId(_modelID.id, _linkIndex);
-  auto it = this->links.find(linkId);
-  if (it != this->links.end() && it->second != nullptr)
+  const auto &[linkId, linkInfo] =
+      this->indexInContainerToId(_modelID.id, _linkIndex, this->links);
+  if (linkInfo != nullptr)
   {
-    return this->GenerateIdentity(it->first, it->second);
+    return this->GenerateIdentity(linkId, linkInfo);
   }
   return this->GenerateInvalidId();
 }
@@ -299,11 +293,11 @@ std::size_t EntityManagementFeatures::GetShapeCount(
 Identity EntityManagementFeatures::GetShape(
   const Identity &_linkID, const std::size_t _shapeIndex) const
 {
-  std::size_t shapeId = this->indexInContainerToId(_linkID.id, _shapeIndex);
-  auto it = this->collisions.find(shapeId);
-  if (it != this->collisions.end() && it->second != nullptr)
+  const auto &[shapeId, shapeInfo] =
+      this->indexInContainerToId(_linkID.id, _shapeIndex, this->collisions);
+  if (shapeInfo != nullptr)
   {
-    return this->GenerateIdentity(it->first, it->second);
+    return this->GenerateIdentity(shapeId, shapeInfo);
   }
   return this->GenerateInvalidId();
 }
@@ -371,8 +365,9 @@ bool EntityManagementFeatures::RemoveModelByIndex(
   auto worldInfo = this->ReferenceInterface<WorldInfo>(_worldID);
   if (worldInfo != nullptr)
   {
-    auto modelId = this->indexInContainerToId(_worldID.id, _modelIndex);
-    if (this->models.find(modelId) != this->models.end())
+    const auto [modelId, modelInfo] =
+        this->indexInContainerToId(_worldID.id, _modelIndex, this->models);
+    if (modelInfo != nullptr)
     {
       this->models.erase(modelId);
       this->childIdToParentId.erase(modelId);

tpe/plugin/src/EntityManagement_TEST.cc

@@ -85,6 +85,8 @@ TEST(EntityManagement_TEST, ConstructEmptyWorld)
   EXPECT_EQ(1u, model->GetNestedModelCount());
   ASSERT_NE(nullptr, nestedModel);
   EXPECT_EQ("empty nested model", nestedModel->GetName());
+  EXPECT_EQ(nestedModel, model->GetNestedModel("empty nested model"));
+  EXPECT_EQ(nestedModel, model->GetNestedModel(0));
   EXPECT_EQ(0u, nestedModel->GetLinkCount());
   ASSERT_NE(nullptr, nestedModel->ConstructEmptyLink("empty link"));
   EXPECT_EQ(1u, nestedModel->GetLinkCount());