From 2bacc55748580e9543a3768a6841c648ef050d28 Mon Sep 17 00:00:00 2001
From: Martin Pecka <peckama2@fel.cvut.cz>
Date: Tue, 23 Feb 2021 19:46:06 +0100
Subject: [PATCH] Add a convenience function for getting possibly non-existing
 components. (#629)

Signed-off-by: Martin Pecka <peckama2@fel.cvut.cz>
---
 .../ignition/gazebo/EntityComponentManager.hh | 13 ++++++++++
 .../gazebo/detail/EntityComponentManager.hh   | 14 +++++++++++
 src/EntityComponentManager_TEST.cc            | 25 +++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/include/ignition/gazebo/EntityComponentManager.hh b/include/ignition/gazebo/EntityComponentManager.hh
index aafcdb9d75..7e9c77b00c 100644
--- a/include/ignition/gazebo/EntityComponentManager.hh
+++ b/include/ignition/gazebo/EntityComponentManager.hh
@@ -214,6 +214,19 @@ namespace ignition
       public: template<typename ComponentTypeT>
               ComponentTypeT *Component(const ComponentKey &_key);
 
+      /// \brief Get a mutable component assigned to an entity based on a
+      /// component type. If the component doesn't exist, create it and
+      /// initialize with the given default value.
+      /// \param[in] _entity The entity.
+      /// \param[in] _default The value that should be used to construct
+      /// the component in case the component doesn't exist.
+      /// \return The component of the specified type assigned to the specified
+      /// entity.
+      public: template<typename ComponentTypeT>
+              ComponentTypeT *ComponentDefault(Entity _entity,
+              const typename ComponentTypeT::Type &_default =
+                  typename ComponentTypeT::Type());
+
       /// \brief Get the data from a component.
       /// * If the component type doesn't hold any data, this won't compile.
       /// * If the entity doesn't have that component, it will return nullopt.
diff --git a/include/ignition/gazebo/detail/EntityComponentManager.hh b/include/ignition/gazebo/detail/EntityComponentManager.hh
index 880dba745c..dcdb37324b 100644
--- a/include/ignition/gazebo/detail/EntityComponentManager.hh
+++ b/include/ignition/gazebo/detail/EntityComponentManager.hh
@@ -127,6 +127,20 @@ ComponentTypeT *EntityComponentManager::Component(const ComponentKey &_key)
       this->ComponentImplementation(_key));
 }
 
+//////////////////////////////////////////////////
+template<typename ComponentTypeT>
+ComponentTypeT *EntityComponentManager::ComponentDefault(Entity _entity,
+    const typename ComponentTypeT::Type &_default)
+{
+  auto comp = this->Component<ComponentTypeT>(_entity);
+  if (!comp)
+  {
+    this->CreateComponent(_entity, ComponentTypeT(_default));
+    comp = this->Component<ComponentTypeT>(_entity);
+  }
+  return comp;
+}
+
 //////////////////////////////////////////////////
 template<typename ComponentTypeT>
 std::optional<typename ComponentTypeT::Type>
diff --git a/src/EntityComponentManager_TEST.cc b/src/EntityComponentManager_TEST.cc
index 59db43c671..22c61b4e7c 100644
--- a/src/EntityComponentManager_TEST.cc
+++ b/src/EntityComponentManager_TEST.cc
@@ -443,6 +443,31 @@ TEST_P(EntityComponentManagerFixture, EntitiesAndComponents)
   EXPECT_FALSE(manager.EntityHasComponentType(kNullEntity,
         IntComponent::typeId));
 
+  // Query non-existing component, the default value is default-constructed
+  BoolComponent *boolComp = manager.ComponentDefault<BoolComponent>(entity);
+  ASSERT_NE(nullptr, boolComp);
+  EXPECT_TRUE(manager.HasComponentType(BoolComponent::typeId));
+  EXPECT_TRUE(manager.EntityHasComponentType(entity, BoolComponent::typeId));
+  EXPECT_EQ(false, boolComp->Data());
+
+  // Query non-existing component, the default value is used
+  DoubleComponent *doubleComp =
+    manager.ComponentDefault<DoubleComponent>(entity, 1.0);
+  ASSERT_NE(nullptr, doubleComp);
+  EXPECT_TRUE(manager.HasComponentType(DoubleComponent::typeId));
+  EXPECT_TRUE(manager.EntityHasComponentType(entity, IntComponent::typeId));
+  EXPECT_TRUE(manager.EntityHasComponentType(entity, DoubleComponent::typeId));
+  EXPECT_FALSE(
+    manager.EntityHasComponentType(entity2, DoubleComponent::typeId));
+  EXPECT_FLOAT_EQ(1.0, doubleComp->Data());
+
+  // Query existing component, the default value is not used
+  IntComponent *intComp = manager.ComponentDefault<IntComponent>(entity, 124);
+  ASSERT_NE(nullptr, intComp);
+  EXPECT_TRUE(manager.HasComponentType(IntComponent::typeId));
+  EXPECT_TRUE(manager.EntityHasComponentType(entity, IntComponent::typeId));
+  EXPECT_EQ(123, intComp->Data());
+
   // Remove all entities
   manager.RequestRemoveEntities();
   EXPECT_EQ(3u, manager.EntityCount());