dartsim: Add support for category bitmask (gazebosim#885)

Signed-off-by: Ian Chen <ichen@openrobotics.org>
Signed-off-by: Ian Chen <iche@intrinsic.ai>
Co-authored-by: Ian Chen <iche@intrinsic.ai>
Signed-off-by: Ian Chen <iche@intrinsic.ai>
(cherry picked from commit e608ccf)

Author: iche033

dartsim/src/EntityManagementFeatures.cc

@@ -44,6 +44,7 @@ class BitmaskContactFilter : public dart::collision::BodyNodeCollisionFilter
   public: using DartShapeConstPtr = const dart::dynamics::ShapeNode*;
 
   private: std::unordered_map<DartShapeConstPtr, uint16_t> bitmaskMap;
+  private: std::unordered_map<DartShapeConstPtr, uint16_t> categoryBitmaskMap;
 
   public: bool ignoresCollision(
       DartCollisionConstPtr _object1,
@@ -58,9 +59,20 @@ class BitmaskContactFilter : public dart::collision::BodyNodeCollisionFilter
 
     auto shape1Iter = bitmaskMap.find(shapeNode1);
     auto shape2Iter = bitmaskMap.find(shapeNode2);
-    if (shape1Iter != bitmaskMap.end() && shape2Iter != bitmaskMap.end() &&
-        ((shape1Iter->second & shape2Iter->second) == 0))
-      return true;
+    if (shape1Iter != bitmaskMap.end() && shape2Iter != bitmaskMap.end())
+    {
+      // For backward compatibility, if category bitmask is not set, it
+      // defaults to the same value as collide bitmask.
+      auto category1Iter = categoryBitmaskMap.find(shapeNode1);
+      auto category2Iter = categoryBitmaskMap.find(shapeNode2);
+      uint16_t categoryMask1 = (category1Iter != categoryBitmaskMap.end()) ?
+          category1Iter->second : shape1Iter->second;
+      uint16_t categoryMask2 = (category2Iter != categoryBitmaskMap.end()) ?
+          category2Iter->second : shape2Iter->second;
+
+      return !((categoryMask1 & shape2Iter->second) |
+             (categoryMask2 & shape1Iter->second));
+    }
 
     return false;
   }
@@ -85,12 +97,35 @@ class BitmaskContactFilter : public dart::collision::BodyNodeCollisionFilter
       bitmaskMap.erase(shapeIter);
   }
 
+  public: void SetIgnoredCategory(DartShapeConstPtr _shapePtr, uint16_t _mask)
+  {
+    categoryBitmaskMap[_shapePtr] = _mask;
+  }
+
+  public: uint16_t GetIgnoredCategory(DartShapeConstPtr _shapePtr) const
+  {
+    auto shapeIter = categoryBitmaskMap.find(_shapePtr);
+    if (shapeIter != categoryBitmaskMap.end())
+      return shapeIter->second;
+    // For backward compatibility, if category bitmask is not set, it
+    // defaults to the same value as collide bitmask.
+    return GetIgnoredCollision(_shapePtr);
+  }
+
+  public: void RemoveIgnoredCategory(DartShapeConstPtr _shapePtr)
+  {
+    auto shapeIter = categoryBitmaskMap.find(_shapePtr);
+    if (shapeIter != categoryBitmaskMap.end())
+      categoryBitmaskMap.erase(shapeIter);
+  }
+
   public: void RemoveSkeletonCollisions(dart::dynamics::SkeletonPtr _skelPtr)
   {
     for (std::size_t i = 0; i < _skelPtr->getNumShapeNodes(); ++i)
     {
       auto shapePtr = _skelPtr->getShapeNode(i);
       this->RemoveIgnoredCollision(shapePtr);
+      this->RemoveIgnoredCategory(shapePtr);
     }
   }
 
@@ -850,6 +885,33 @@ void EntityManagementFeatures::RemoveCollisionFilterMask(
   filterPtr->RemoveIgnoredCollision(shapeNode);
 }
 
+void EntityManagementFeatures::SetCategoryFilterMask(
+    const Identity &_shapeID, uint16_t _mask)
+{
+  const auto shapeNode = this->ReferenceInterface<ShapeInfo>(_shapeID)->node;
+  const std::size_t worldID = GetWorldOfShapeNode(this, shapeNode);
+  const auto filterPtr = GetFilterPtr(this, worldID);
+  filterPtr->SetIgnoredCategory(shapeNode, _mask);
+}
+
+uint16_t EntityManagementFeatures::GetCategoryFilterMask(
+    const Identity &_shapeID) const
+{
+  const auto shapeNode = this->ReferenceInterface<ShapeInfo>(_shapeID)->node;
+  const std::size_t worldID = GetWorldOfShapeNode(this, shapeNode);
+  const auto filterPtr = GetFilterPtr(this, worldID);
+  return filterPtr->GetIgnoredCategory(shapeNode);
+}
+
+void EntityManagementFeatures::RemoveCategoryFilterMask(
+    const Identity &_shapeID)
+{
+  const auto shapeNode = this->ReferenceInterface<ShapeInfo>(_shapeID)->node;
+  const std::size_t worldID = GetWorldOfShapeNode(this, shapeNode);
+  const auto filterPtr = GetFilterPtr(this, worldID);
+  filterPtr->RemoveIgnoredCategory(shapeNode);
+}
+
 Identity EntityManagementFeatures::GetWorldModel(const Identity &_worldID) const
 {
   auto modelID = this->modelProxiesToWorld.MaybeAt(_worldID);

dartsim/src/EntityManagementFeatures.hh

@@ -37,6 +37,7 @@ namespace dartsim {
 struct EntityManagementFeatureList : FeatureList<
   GetEntities,
   RemoveEntities,
+  CategoryFilterMaskFeature,
   ConstructEmptyWorldFeature,
   ConstructEmptyModelFeature,
   ConstructEmptyNestedModelFeature,
@@ -178,6 +179,15 @@ class GZ_PHYSICS_DARTSIM_PLUGIN_VISIBLE EntityManagementFeatures :
 
   public: void RemoveCollisionFilterMask(const Identity &_shapeID) override;
 
+  // ----- Manage category filter masks -----
+  public: void SetCategoryFilterMask(
+      const Identity &_shapeID, uint16_t _mask) override;
+
+  public: uint16_t GetCategoryFilterMask(
+      const Identity &_shapeID) const override;
+
+  public: void RemoveCategoryFilterMask(const Identity &_shapeID) override;
+
   // ----- World model feature -----
   public: Identity GetWorldModel(const Identity &_worldID) const override;
 };

dartsim/src/SDFFeatures.cc

@@ -20,6 +20,7 @@
 #include <cmath>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 
@@ -885,6 +886,7 @@ Identity SDFFeatures::ConstructSdfCollision(
   // TODO(addisu) We are using the coefficient specified in the <ode> tag.
   // Either add parameters specific to DART or generic to all physics engines.
   uint16_t collideBitmask = 0xFF;
+  std::optional<uint16_t> categoryBitmask;
   if (_collision.Element())
   {
     const auto &odeFriction = _collision.Element()
@@ -958,12 +960,13 @@ Identity SDFFeatures::ConstructSdfCollision(
           surfaceBounce->Get<double>("restitution_coefficient"));
     }
 #endif
-    // TODO(anyone) add category_bitmask as well
-    const auto bitmaskElement = _collision.Element()
+    const auto contactElement = _collision.Element()
                                      ->GetElement("surface")
                                      ->GetElement("contact");
-    if (bitmaskElement->HasElement("collide_bitmask"))
-      collideBitmask = bitmaskElement->Get<int>("collide_bitmask");
+    if (contactElement->HasElement("collide_bitmask"))
+      collideBitmask = contactElement->Get<int>("collide_bitmask");
+    if (contactElement->HasElement("category_bitmask"))
+      categoryBitmask = contactElement->Get<int>("category_bitmask");
   }
 
   node->setRelativeTransform(ResolveSdfPose(_collision.SemanticPose()) *
@@ -974,6 +977,8 @@ Identity SDFFeatures::ConstructSdfCollision(
   auto identity = this->GenerateIdentity(shapeID, this->shapes.at(shapeID));
 
   this->SetCollisionFilterMask(identity, collideBitmask);
+  if (categoryBitmask.has_value())
+    this->SetCategoryFilterMask(identity, categoryBitmask.value());
   return identity;
 }