From 67bf92b0cbdfb2937b110f44903e504adc96a1a4 Mon Sep 17 00:00:00 2001
From: Steve Peters <scpeters@openrobotics.org>
Date: Sat, 28 Oct 2023 01:39:17 -0700
Subject: [PATCH] Add Link auto_inertia_params to API

Pass //link/inertial/auto_inertia_params ElementPtr to
Collision::CalculateInertial and use it if the
Collision does not have auto_inertia_params of its own.

Signed-off-by: Steve Peters <scpeters@openrobotics.org>
---
 include/sdf/Collision.hh |  6 +++++-
 include/sdf/Link.hh      | 12 ++++++++++++
 src/Collision.cc         | 15 ++++++++++++---
 src/Link.cc              | 25 ++++++++++++++++++++++++-
 4 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/include/sdf/Collision.hh b/include/sdf/Collision.hh
index 6667896d6..2ab4114ce 100644
--- a/include/sdf/Collision.hh
+++ b/include/sdf/Collision.hh
@@ -167,11 +167,15 @@ namespace sdf
     /// \param[in] _density An optional density value to override the default
     /// collision density. This value is used instead of DefaultDensity()
     // if this collision's density has not been explicitly set.
+    /// \param[in] _autoInertiaParams An ElementPtr to the auto_inertia_params
+    /// element to be used if the auto_inertia_params have not been set in this
+    /// collision.
     public: void CalculateInertial(
                     sdf::Errors &_errors,
                     gz::math::Inertiald &_inertial,
                     const ParserConfig &_config,
-                    const std::optional<double> &_density);
+                    const std::optional<double> &_density,
+                    sdf::ElementPtr _autoInertiaParams);
 
     /// \brief Get a pointer to the SDF element that was used during
     /// load.
diff --git a/include/sdf/Link.hh b/include/sdf/Link.hh
index 00a65b080..3e5260e73 100644
--- a/include/sdf/Link.hh
+++ b/include/sdf/Link.hh
@@ -89,6 +89,18 @@ namespace sdf
     /// \param[in] _density Density of the inertial.
     public: void SetDensity(double _density);
 
+    /// \brief Get the ElementPtr to the <auto_inertia_params> element
+    /// This element can be used as a parent element to hold user-defined
+    /// params for the custom moment of inertia calculator.
+    /// \return ElementPtr object for the <auto_inertia_params> element.
+    public: sdf::ElementPtr AutoInertiaParams() const;
+
+    /// \brief Function to set the auto inertia params using a
+    /// sdf ElementPtr object
+    /// \param[in] _autoInertiaParams ElementPtr to <auto_inertia_params>
+    /// element
+    public: void SetAutoInertiaParams(const sdf::ElementPtr _autoInertiaParams);
+
     /// \brief Get the number of visuals.
     /// \return Number of visuals contained in this Link object.
     public: uint64_t VisualCount() const;
diff --git a/src/Collision.cc b/src/Collision.cc
index 7c29b7657..05ca351bf 100644
--- a/src/Collision.cc
+++ b/src/Collision.cc
@@ -262,7 +262,8 @@ void Collision::CalculateInertial(
   gz::math::Inertiald &_inertial,
   const ParserConfig &_config)
 {
-  this->CalculateInertial(_errors, _inertial, _config, std::nullopt);
+  this->CalculateInertial(
+    _errors, _inertial, _config, std::nullopt, ElementPtr());
 }
 
 /////////////////////////////////////////////////
@@ -270,7 +271,8 @@ void Collision::CalculateInertial(
   sdf::Errors &_errors,
   gz::math::Inertiald &_inertial,
   const ParserConfig &_config,
-  const std::optional<double> &_density)
+  const std::optional<double> &_density,
+  sdf::ElementPtr _autoInertiaParams)
 {
   // Order of precedence for density:
   double density;
@@ -303,9 +305,16 @@ void Collision::CalculateInertial(
     );
   }
 
+  // If this Collision's auto inertia params have not been set, then use the
+  // params passed into this function.
+  sdf::ElementPtr autoInertiaParams = this->dataPtr->autoInertiaParams;
+  if (!autoInertiaParams)
+  {
+    autoInertiaParams = _autoInertiaParams;
+  }
   auto geomInertial =
     this->dataPtr->geom.CalculateInertial(_errors, _config,
-      density, this->dataPtr->autoInertiaParams);
+      density, autoInertiaParams);
 
   if (!geomInertial)
   {
diff --git a/src/Link.cc b/src/Link.cc
index c3d772dd5..659b25cbe 100644
--- a/src/Link.cc
+++ b/src/Link.cc
@@ -73,6 +73,9 @@ class sdf::Link::Implementation
   /// calculations if the collision density has not been set.
   public: std::optional<double> density;
 
+  /// \brief SDF element pointer to <auto_inertia_params> tag
+  public: sdf::ElementPtr autoInertiaParams{nullptr};
+
   /// \brief The inertial information for this link.
   public: gz::math::Inertiald inertial {{1.0,
             gz::math::Vector3d::One, gz::math::Vector3d::Zero},
@@ -189,6 +192,13 @@ Errors Link::Load(ElementPtr _sdf, const ParserConfig &_config)
       this->dataPtr->density = inertialElem->Get<double>("density");
     }
 
+    // Load the auto_inertia_params element
+    if (inertialElem->HasElement("auto_inertia_params"))
+    {
+      this->dataPtr->autoInertiaParams =
+        inertialElem->GetElement("auto_inertia_params");
+    }
+
     if (inertialElem->Get<bool>("auto"))
     {
       this->dataPtr->autoInertia = true;
@@ -330,6 +340,18 @@ void Link::SetDensity(double _density)
   this->dataPtr->density = _density;
 }
 
+/////////////////////////////////////////////////
+sdf::ElementPtr Link::AutoInertiaParams() const
+{
+  return this->dataPtr->autoInertiaParams;
+}
+
+/////////////////////////////////////////////////
+void Link::SetAutoInertiaParams(const sdf::ElementPtr _autoInertiaParams)
+{
+  this->dataPtr->autoInertiaParams = _autoInertiaParams;
+}
+
 /////////////////////////////////////////////////
 uint64_t Link::VisualCount() const
 {
@@ -642,7 +664,8 @@ void Link::ResolveAutoInertials(sdf::Errors &_errors,
     {
       gz::math::Inertiald collisionInertia;
       collision.CalculateInertial(_errors, collisionInertia, _config,
-                                  this->dataPtr->density);
+                                  this->dataPtr->density,
+                                  this->dataPtr->autoInertiaParams);
       totalInertia = totalInertia + collisionInertia;
     }