Fix collision shape caching

tesseract_collision/bullet/include/tesseract_collision/bullet/bullet_collision_shape_cache.h

@@ -29,17 +29,28 @@
 
 #include <map>
 #include <memory>
-#include <shared_mutex>
+#include <vector>
+#include <mutex>
 
 #include <tesseract_common/macros.h>
 TESSERACT_COMMON_IGNORE_WARNINGS_PUSH
 #include <BulletCollision/CollisionShapes/btCollisionShape.h>
+#include <boost/uuid/uuid.hpp>
 TESSERACT_COMMON_IGNORE_WARNINGS_POP
 
 #include <tesseract_geometry/fwd.h>
 
 namespace tesseract_collision::tesseract_collision_bullet
 {
+struct BulletCollisionShape
+{
+  BulletCollisionShape() = default;
+  BulletCollisionShape(std::shared_ptr<btCollisionShape> top_level_);
+
+  std::shared_ptr<btCollisionShape> top_level;
+  std::vector<std::shared_ptr<btCollisionShape>> children;
+};
+
 class BulletCollisionShapeCache
 {
 public:
@@ -49,20 +60,23 @@ class BulletCollisionShapeCache
    * @param value The value to store
    */
   static void insert(const std::shared_ptr<const tesseract_geometry::Geometry>& key,
-                     const std::shared_ptr<btCollisionShape>& value);
+                     const std::shared_ptr<BulletCollisionShape>& value);
 
   /**
    * @brief Retrieve the cache entry by key
    * @param key The cache key
    * @return If key exists the entry is returned, otherwise a nullptr is returned
    */
-  static std::shared_ptr<btCollisionShape> get(const std::shared_ptr<const tesseract_geometry::Geometry>& key);
+  static std::shared_ptr<BulletCollisionShape> get(const std::shared_ptr<const tesseract_geometry::Geometry>& key);
+
+  /** @brief Remove any entries which are no longer valid */
+  static void prune();
 
 private:
   /** @brief The static cache */
-  static std::map<std::shared_ptr<const tesseract_geometry::Geometry>, std::shared_ptr<btCollisionShape>> cache_;
+  static std::map<boost::uuids::uuid, std::weak_ptr<BulletCollisionShape>> cache_;
   /** @brief The shared mutex for thread safety */
-  static std::shared_mutex mutex_;
+  static std::mutex mutex_;
 };
 }  // namespace tesseract_collision::tesseract_collision_bullet
 

tesseract_collision/bullet/include/tesseract_collision/bullet/bullet_utils.h

@@ -53,6 +53,7 @@ TESSERACT_COMMON_IGNORE_WARNINGS_POP
 
 #include <tesseract_collision/core/types.h>
 #include <tesseract_collision/core/common.h>
+#include <tesseract_collision/bullet/bullet_collision_shape_cache.h>
 
 namespace tesseract_collision::tesseract_collision_bullet
 {
@@ -129,7 +130,7 @@ class CollisionObjectWrapper : public btCollisionObject
    */
   std::shared_ptr<CollisionObjectWrapper> clone();
 
-  void manage(const std::shared_ptr<btCollisionShape>& t);
+  void manage(const std::shared_ptr<BulletCollisionShape>& t);
 
   void manageReserve(std::size_t s);
 
@@ -143,7 +144,7 @@ class CollisionObjectWrapper : public btCollisionObject
   /**< @brief The shapes poses information */
   tesseract_common::VectorIsometry3d m_shape_poses{};
   /** @brief This manages the collision shape pointer so they get destroyed */
-  std::vector<std::shared_ptr<btCollisionShape>> m_data{};
+  std::vector<std::shared_ptr<BulletCollisionShape>> m_data{};
 };
 
 using COW = CollisionObjectWrapper;
@@ -362,10 +363,9 @@ class TesseractOverlapFilterCallback : public btOverlapFilterCallback
 /**
  * @brief Create a bullet collision shape from tesseract collision shape
  * @param geom Tesseract collision shape
- * @param cow The collision object wrapper the collision shape is associated with
  * @return Bullet collision shape.
  */
-std::shared_ptr<btCollisionShape> createShapePrimitive(const CollisionShapeConstPtr& geom, CollisionObjectWrapper* cow);
+std::shared_ptr<BulletCollisionShape> createShapePrimitive(const CollisionShapeConstPtr& geom);
 
 /**
  * @brief Update a collision objects filters

tesseract_collision/bullet/src/bullet_collision_shape_cache.cpp

@@ -32,23 +32,49 @@
 namespace tesseract_collision::tesseract_collision_bullet
 {
 // Static member definitions
-std::map<std::shared_ptr<const tesseract_geometry::Geometry>, std::shared_ptr<btCollisionShape>>
-    BulletCollisionShapeCache::cache_;
-std::shared_mutex BulletCollisionShapeCache::mutex_;
+std::map<boost::uuids::uuid, std::weak_ptr<BulletCollisionShape>> BulletCollisionShapeCache::cache_;
+std::mutex BulletCollisionShapeCache::mutex_;
+
+BulletCollisionShape::BulletCollisionShape(std::shared_ptr<btCollisionShape> top_level_)
+  : top_level(std::move(top_level_))
+{
+}
 
 void BulletCollisionShapeCache::insert(const std::shared_ptr<const tesseract_geometry::Geometry>& key,
-                                       const std::shared_ptr<btCollisionShape>& value)
+                                       const std::shared_ptr<BulletCollisionShape>& value)
 {
-  std::unique_lock lock(mutex_);
-  cache_[key] = value;
+  assert(!key->getUUID().is_nil());
+  std::scoped_lock lock(mutex_);
+  cache_[key->getUUID()] = value;
 }
 
-std::shared_ptr<btCollisionShape>
+std::shared_ptr<BulletCollisionShape>
 BulletCollisionShapeCache::get(const std::shared_ptr<const tesseract_geometry::Geometry>& key)
 {
-  std::shared_lock lock(mutex_);
-  auto it = cache_.find(key);
-  return ((it != cache_.end()) ? it->second : nullptr);
+  assert(!key->getUUID().is_nil());
+  std::scoped_lock lock(mutex_);
+  auto it = cache_.find(key->getUUID());
+  if (it != cache_.end())
+  {
+    std::shared_ptr<BulletCollisionShape> collision_shape = it->second.lock();
+    if (collision_shape != nullptr)
+      return collision_shape;
+
+    cache_.erase(key->getUUID());
+  }
+  return nullptr;
+}
+
+void BulletCollisionShapeCache::prune()
+{
+  std::scoped_lock lock(mutex_);
+  for (auto it = cache_.begin(); it != cache_.end();)
+  {
+    if (it->second.expired())
+      it = cache_.erase(it);
+    else
+      ++it;
+  }
 }
 
 }  // namespace tesseract_collision::tesseract_collision_bullet