Add iterator methods to container-like classes (#1644)

dartsim · Jan 6, 2022 · f07186f · f07186f
1 parent 5ef8c1f
commit f07186f
Show file tree

Hide file tree

Showing 41 changed files with 1,111 additions and 305 deletions.
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
@@ -13,3 +13,4 @@
 - [ ] Set version target by selecting a milestone on the right side
 - [ ] Summarize this change in `CHANGELOG.md`
 - [ ] Add unit test(s) for this change
+- [ ] Add Python bindings for new methods and classes
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@
 * Dynamics
 
   * Added deep copy for shapes: [#1612](https://github.com/dartsim/dart/pull/1612)
+  * Added iterator methods to container-like classes: [#1644](https://github.com/dartsim/dart/pull/1644)
   * Fixed grouping of constraints: [#1624](https://github.com/dartsim/dart/pull/1624), [#1628](https://github.com/dartsim/dart/pull/1628)
   * Fixed issue with removing skeletons without shapes: [#1625](https://github.com/dartsim/dart/pull/1625)
 

diff --git a/dart/collision/CollisionGroup.cpp b/dart/collision/CollisionGroup.cpp
@@ -480,15 +480,12 @@ bool CollisionGroup::updateSkeletonSource(SkeletonSources::value_type& entry)
       // CollisionGroup.
       updateNeeded = true;
 
-      const auto& collisionShapeNodes
-          = bn->getShapeNodesWith<dynamics::CollisionAspect>();
-
-      for (const auto& shapeNode : collisionShapeNodes)
-      {
-        source.mObjects.insert(
-            {shapeNode, addShapeFrameImpl(shapeNode, meta.get())});
-        child->second.mFrames.insert(shapeNode);
-      }
+      bn->eachShapeNodeWith<dynamics::CollisionAspect>(
+          [&](const dynamics::ShapeNode* shapeNode) {
+            source.mObjects.insert(
+                {shapeNode, addShapeFrameImpl(shapeNode, meta.get())});
+            child->second.mFrames.insert(shapeNode);
+          });
 
       // Skip the rest of the for-loop, because it's only for updating old
       // children
@@ -505,33 +502,30 @@ bool CollisionGroup::updateSkeletonSource(SkeletonSources::value_type& entry)
     std::unordered_set<const dynamics::ShapeFrame*> unusedFrames
         = child->second.mFrames;
 
-    const std::vector<const dynamics::ShapeNode*> nodes
-        = child->first->getShapeNodesWith<dynamics::CollisionAspect>();
-
-    for (const dynamics::ShapeNode* node : nodes)
-    {
-      unusedFrames.erase(node);
-
-      auto frameInsertion
-          = source.mObjects.insert(std::make_pair(node, nullptr));
-
-      const auto& it = frameInsertion.first;
-      if (frameInsertion.second)
-      {
-        // If the insertion occurred, then this is a new ShapeFrame, and we
-        // need to create a collision object for it.
-        updateNeeded = true;
-
-        it->second = addShapeFrameImpl(node, meta.get());
-        child->second.mFrames.insert(node);
-      }
-      else
-      {
-        // If the insertion did not occur, then this is an old ShapeFrame, and
-        // we should check if it needs an update.
-        updateNeeded |= updateShapeFrame(it->second);
-      }
-    }
+    child->first->eachShapeNodeWith<dynamics::CollisionAspect>(
+        [&](const dynamics::ShapeNode* shapeNode) {
+          unusedFrames.erase(shapeNode);
+
+          auto frameInsertion
+              = source.mObjects.insert(std::make_pair(shapeNode, nullptr));
+
+          const auto& it = frameInsertion.first;
+          if (frameInsertion.second)
+          {
+            // If the insertion occurred, then this is a new ShapeFrame, and we
+            // need to create a collision object for it.
+            updateNeeded = true;
+
+            it->second = addShapeFrameImpl(shapeNode, meta.get());
+            child->second.mFrames.insert(shapeNode);
+          }
+          else
+          {
+            // If the insertion did not occur, then this is an old ShapeFrame,
+            // and we should check if it needs an update.
+            updateNeeded |= updateShapeFrame(it->second);
+          }
+        });
 
     for (const dynamics::ShapeFrame* unused : unusedFrames)
     {
@@ -585,30 +579,28 @@ bool CollisionGroup::updateBodyNodeSource(BodyNodeSources::value_type& entry)
   std::unordered_map<const dynamics::ShapeFrame*, ObjectInfo*> unusedFrames
       = source.mObjects;
 
-  const std::vector<const dynamics::ShapeNode*> nodes
-      = bn->getShapeNodesWith<dynamics::CollisionAspect>();
-
-  for (const dynamics::ShapeNode* node : nodes)
-  {
-    unusedFrames.erase(node);
-
-    auto frameInsertion = source.mObjects.insert(std::make_pair(node, nullptr));
-
-    const auto& it = frameInsertion.first;
-    if (frameInsertion.second)
-    {
-      updateNeeded = true;
-      // If the insertion occurred, then this is a new ShapeFrame, and we need
-      // to create a collision object for it.
-      it->second = addShapeFrameImpl(node, bn.get());
-    }
-    else
-    {
-      // If the insertion did not occur, then this is an old ShapeFrame, and
-      // we should check if it needs an update.
-      updateNeeded |= updateShapeFrame(it->second);
-    }
-  }
+  bn->eachShapeNodeWith<dynamics::CollisionAspect>(
+      [&](const dynamics::ShapeNode* shapeNode) {
+        unusedFrames.erase(shapeNode);
+
+        auto frameInsertion
+            = source.mObjects.insert(std::make_pair(shapeNode, nullptr));
+
+        const auto& it = frameInsertion.first;
+        if (frameInsertion.second)
+        {
+          updateNeeded = true;
+          // If the insertion occurred, then this is a new ShapeFrame, and we
+          // need to create a collision object for it.
+          it->second = addShapeFrameImpl(shapeNode, bn.get());
+        }
+        else
+        {
+          // If the insertion did not occur, then this is an old ShapeFrame, and
+          // we should check if it needs an update.
+          updateNeeded |= updateShapeFrame(it->second);
+        }
+      });
 
   // Remove from this group and ShapeFrames that no longer belong to the
   // BodyNode

diff --git a/dart/collision/detail/CollisionGroup.hpp b/dart/collision/detail/CollisionGroup.hpp
@@ -84,11 +84,10 @@ void CollisionGroup::addShapeFramesOf(
 {
   assert(bodyNode);
 
-  auto collisionShapeNodes
-      = bodyNode->getShapeNodesWith<dynamics::CollisionAspect>();
-
-  for (auto& shapeNode : collisionShapeNodes)
-    addShapeFrame(shapeNode);
+  bodyNode->eachShapeNodeWith<dynamics::CollisionAspect>(
+      [this](const dynamics::ShapeNode* shapeNode) {
+        addShapeFrame(shapeNode);
+      });
 
   addShapeFramesOf(others...);
 }
@@ -118,15 +117,11 @@ void CollisionGroup::subscribeTo(
   if (inserted.second)
   {
     const BodyNodeSources::iterator& entry = inserted.first;
-
-    const auto collisionShapeNodes
-        = bodyNode->getShapeNodesWith<dynamics::CollisionAspect>();
-
-    for (const auto& shapeNode : collisionShapeNodes)
-    {
-      entry->second.mObjects.insert(
-          {shapeNode, addShapeFrameImpl(shapeNode, bodyNode.get())});
-    }
+    bodyNode->eachShapeNodeWith<dynamics::CollisionAspect>(
+        [&](const dynamics::ShapeNode* shapeNode) {
+          entry->second.mObjects.insert(
+              {shapeNode, addShapeFrameImpl(shapeNode, bodyNode.get())});
+        });
   }
 
   subscribeTo(others...);
@@ -149,21 +144,18 @@ void CollisionGroup::subscribeTo(
     {
       const dynamics::BodyNode* bn = skeleton->getBodyNode(i);
 
-      const auto& collisionShapeNodes
-          = bn->getShapeNodesWith<dynamics::CollisionAspect>();
-
       auto& childInfo
           = entry.mChildren
                 .insert(std::make_pair(
                     bn, SkeletonSource::ChildInfo(bn->getVersion())))
                 .first->second;
 
-      for (const auto& shapeNode : collisionShapeNodes)
-      {
-        entry.mObjects.insert(
-            {shapeNode, addShapeFrameImpl(shapeNode, skeleton.get())});
-        childInfo.mFrames.insert(shapeNode);
-      }
+      bn->eachShapeNodeWith<dynamics::CollisionAspect>(
+          [&](const dynamics::ShapeNode* shapeNode) {
+            entry.mObjects.insert(
+                {shapeNode, addShapeFrameImpl(shapeNode, skeleton.get())});
+            childInfo.mFrames.insert(shapeNode);
+          });
     }
   }
 
@@ -220,11 +212,10 @@ void CollisionGroup::removeShapeFramesOf(
 {
   assert(bodyNode);
 
-  auto collisionShapeNodes
-      = bodyNode->getShapeNodesWith<dynamics::CollisionAspect>();
-
-  for (auto& shapeNode : collisionShapeNodes)
-    removeShapeFrame(shapeNode);
+  bodyNode->eachShapeNodeWith<dynamics::CollisionAspect>(
+      [&](const dynamics::ShapeNode* shapeNode) {
+        removeShapeFrame(shapeNode);
+      });
 
   removeShapeFramesOf(others...);
 }

diff --git a/dart/dynamics/BodyNode.cpp b/dart/dynamics/BodyNode.cpp
@@ -664,24 +664,22 @@ Eigen::Vector6d BodyNode::getCOMSpatialAcceleration(
 }
 
 //==============================================================================
-void BodyNode::setFrictionCoeff(double _coeff)
+void BodyNode::setFrictionCoeff(double coeff)
 {
   // Below code block is to set the friction coefficients of all the current
   // dynamics aspects of this BodyNode as a stopgap solution. However, this
   // won't help for new ShapeNodes that get added later.
-  auto shapeNodes = getShapeNodesWith<DynamicsAspect>();
-  for (auto& shapeNode : shapeNodes)
-  {
+  eachShapeNodeWith<DynamicsAspect>([&](ShapeNode* shapeNode) {
     auto* dynamicsAspect = shapeNode->getDynamicsAspect();
-    dynamicsAspect->setFrictionCoeff(_coeff);
-  }
+    dynamicsAspect->setFrictionCoeff(coeff);
+  });
 
-  if (mAspectProperties.mFrictionCoeff == _coeff)
+  if (mAspectProperties.mFrictionCoeff == coeff)
     return;
 
   assert(
-      0.0 <= _coeff && "Coefficient of friction should be non-negative value.");
-  mAspectProperties.mFrictionCoeff = _coeff;
+      0.0 <= coeff && "Coefficient of friction should be non-negative value.");
+  mAspectProperties.mFrictionCoeff = coeff;
 
   incrementVersion();
 }
@@ -693,25 +691,23 @@ double BodyNode::getFrictionCoeff() const
 }
 
 //==============================================================================
-void BodyNode::setRestitutionCoeff(double _coeff)
+void BodyNode::setRestitutionCoeff(double coeff)
 {
   // Below code block is to set the restitution coefficients of all the current
   // dynamics aspects of this BodyNode as a stopgap solution. However, this
   // won't help for new ShapeNodes that get added later.
-  auto shapeNodes = getShapeNodesWith<DynamicsAspect>();
-  for (auto& shapeNode : shapeNodes)
-  {
+  eachShapeNodeWith<DynamicsAspect>([&](ShapeNode* shapeNode) {
     auto* dynamicsAspect = shapeNode->getDynamicsAspect();
-    dynamicsAspect->setFrictionCoeff(_coeff);
-  }
+    dynamicsAspect->setFrictionCoeff(coeff);
+  });
 
-  if (_coeff == mAspectProperties.mRestitutionCoeff)
+  if (coeff == mAspectProperties.mRestitutionCoeff)
     return;
 
   assert(
-      0.0 <= _coeff && _coeff <= 1.0
+      0.0 <= coeff && coeff <= 1.0
       && "Coefficient of restitution should be in range of [0, 1].");
-  mAspectProperties.mRestitutionCoeff = _coeff;
+  mAspectProperties.mRestitutionCoeff = coeff;
 
   incrementVersion();
 }
@@ -982,9 +978,7 @@ const std::vector<const ShapeNode*> BodyNode::getShapeNodes() const
 //==============================================================================
 void BodyNode::removeAllShapeNodes()
 {
-  auto shapeNodes = getShapeNodes();
-  for (auto shapeNode : shapeNodes)
-    shapeNode->remove();
+  eachShapeNode([](ShapeNode* shapeNode) { shapeNode->remove(); });
 }
 
 //==============================================================================
@@ -1616,6 +1610,30 @@ void BodyNode::dirtyCoriolisForces()
   SKEL_SET_FLAGS(mCoriolisAndGravityForces);
 }
 
+//==============================================================================
+void BodyNode::setColor(const Eigen::Vector3d& color)
+{
+  eachShapeNodeWith<VisualAspect>([&](ShapeNode* shapeNode) {
+    shapeNode->getVisualAspect()->setColor(color);
+  });
+}
+
+//==============================================================================
+void BodyNode::setColor(const Eigen::Vector4d& color)
+{
+  eachShapeNodeWith<VisualAspect>([&](ShapeNode* shapeNode) {
+    shapeNode->getVisualAspect()->setColor(color);
+  });
+}
+
+//==============================================================================
+void BodyNode::setAlpha(double alpha)
+{
+  eachShapeNodeWith<VisualAspect>([&](ShapeNode* shapeNode) {
+    shapeNode->getVisualAspect()->setAlpha(alpha);
+  });
+}
+
 //==============================================================================
 void BodyNode::updateTransform()
 {