From 1e2bcdd18c96cfc6fb14e12c7656b2b6945fa410 Mon Sep 17 00:00:00 2001
From: Ian Chen <ichen@openrobotics.org>
Date: Tue, 14 Nov 2023 21:50:57 +0000
Subject: [PATCH] update remove logic, cleanup

Signed-off-by: Ian Chen <ichen@openrobotics.org>
---
 bullet-featherstone/src/Base.hh               | 25 +++++++++++--------
 .../src/EntityManagementFeatures.cc           |  4 ---
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/bullet-featherstone/src/Base.hh b/bullet-featherstone/src/Base.hh
index ae0f8fbec..53197d7ee 100644
--- a/bullet-featherstone/src/Base.hh
+++ b/bullet-featherstone/src/Base.hh
@@ -405,17 +405,22 @@ class Base : public Implements3d<FeatureList<Feature>>
     model->nestedModelEntityIds.clear();
 
     // remove references in parent model or world model
-    auto *parentModel = this->ReferenceInterface<ModelInfo>(_parentID);
-    auto nestedModelIt =
-        parentModel->nestedModelNameToEntityId.find(model->name);
-    if (nestedModelIt != parentModel->nestedModelNameToEntityId.end())
+    auto parentModelIt = this->models.find(_parentID);
+    if (parentModelIt != this->models.end())
     {
-      std::size_t nestedModelID = nestedModelIt->second;
-      parentModel->nestedModelNameToEntityId.erase(nestedModelIt);
-      parentModel->nestedModelEntityIds.erase(std::remove(
-          parentModel->nestedModelEntityIds.begin(),
-          parentModel->nestedModelEntityIds.end(), nestedModelID),
-          parentModel->nestedModelEntityIds.end());
+      auto parentModel = parentModelIt->second;
+      auto nestedModelIt =
+          parentModel->nestedModelNameToEntityId.find(model->name);
+      if (nestedModelIt !=
+          parentModel->nestedModelNameToEntityId.end())
+      {
+        std::size_t nestedModelID = nestedModelIt->second;
+        parentModel->nestedModelNameToEntityId.erase(nestedModelIt);
+        parentModel->nestedModelEntityIds.erase(std::remove(
+            parentModel->nestedModelEntityIds.begin(),
+            parentModel->nestedModelEntityIds.end(), nestedModelID),
+            parentModel->nestedModelEntityIds.end());
+      }
     }
 
     // If nested, we are done here. No need to remove multibody
diff --git a/bullet-featherstone/src/EntityManagementFeatures.cc b/bullet-featherstone/src/EntityManagementFeatures.cc
index c86f5440e..99e40b827 100644
--- a/bullet-featherstone/src/EntityManagementFeatures.cc
+++ b/bullet-featherstone/src/EntityManagementFeatures.cc
@@ -215,10 +215,6 @@ bool EntityManagementFeatures::RemoveModel(const Identity &_modelID)
 bool EntityManagementFeatures::ModelRemoved(
     const Identity &_modelID) const
 {
-//  auto *model = this->ReferenceInterface<ModelInfo>(_modelID);
-//  auto *world = this->ReferenceInterface<WorldInfo>(model->world);
-//  return world->modelIndexToEntityId.count(model->indexInWorld) == 0;
-
   return this->models.find(_modelID) == this->models.end();
 }