From 28752ea64738aeea1034039e9eb747bfbb5e1fe6 Mon Sep 17 00:00:00 2001
From: Levi Armstrong <levi.armstrong@gmail.com>
Date: Thu, 25 Apr 2024 13:32:37 -0500
Subject: [PATCH] Add utility methods to task composer node info container

---
 .../core/task_composer_node_info.h            | 18 ++++++++++++++-
 .../core/src/task_composer_node_info.cpp      | 23 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/tesseract_task_composer/core/include/tesseract_task_composer/core/task_composer_node_info.h b/tesseract_task_composer/core/include/tesseract_task_composer/core/task_composer_node_info.h
index ffe99f4127..c52e3968e8 100644
--- a/tesseract_task_composer/core/include/tesseract_task_composer/core/task_composer_node_info.h
+++ b/tesseract_task_composer/core/include/tesseract_task_composer/core/task_composer_node_info.h
@@ -32,6 +32,7 @@ TESSERACT_COMMON_IGNORE_WARNINGS_PUSH
 #include <shared_mutex>
 #include <map>
 #include <chrono>
+#include <functional>
 #include <boost/serialization/access.hpp>
 #include <boost/serialization/export.hpp>
 #include <boost/uuid/uuid.hpp>
@@ -60,9 +61,12 @@ class TaskComposerNodeInfo
   TaskComposerNodeInfo(TaskComposerNodeInfo&&) = default;
   TaskComposerNodeInfo& operator=(TaskComposerNodeInfo&&) = default;
 
-  /** @brief The name */
+  /** @brief The name of the task */
   std::string name;
 
+  /** @brief The namespace of the task */
+  std::string ns;
+
   /** @brief The task uuid */
   boost::uuids::uuid uuid{};
 
@@ -165,6 +169,15 @@ class TaskComposerNodeInfoContainer
    */
   TaskComposerNodeInfo::UPtr getInfo(const boost::uuids::uuid& key) const;
 
+  /**
+   * @brief Get task info by name
+   * @param name The name of task info to search
+   * @param ns The namespace to search under. If empty all namespace's are include
+   * @return A list of task infos with the following name.
+   */
+  std::vector<TaskComposerNodeInfo::UPtr>
+  find(const std::function<bool(const TaskComposerNodeInfo& node_info)>& search_fn) const;
+
   /** @brief Get a copy of the task_info_map_ in case it gets resized*/
   std::map<boost::uuids::uuid, TaskComposerNodeInfo::UPtr> getInfoMap() const;
 
@@ -189,6 +202,9 @@ class TaskComposerNodeInfoContainer
   /** @brief Clear the contents */
   void clear();
 
+  /** @brief Prune data from node infos to save space */
+  void prune(const std::function<void(TaskComposerNodeInfo&)>& prune_fn);
+
   bool operator==(const TaskComposerNodeInfoContainer& rhs) const;
   bool operator!=(const TaskComposerNodeInfoContainer& rhs) const;
 
diff --git a/tesseract_task_composer/core/src/task_composer_node_info.cpp b/tesseract_task_composer/core/src/task_composer_node_info.cpp
index ab895a3486..5b8be0d5a5 100644
--- a/tesseract_task_composer/core/src/task_composer_node_info.cpp
+++ b/tesseract_task_composer/core/src/task_composer_node_info.cpp
@@ -45,6 +45,7 @@ namespace tesseract_planning
 {
 TaskComposerNodeInfo::TaskComposerNodeInfo(const TaskComposerNode& node)
   : name(node.name_)
+  , ns(node.ns_)
   , uuid(node.uuid_)
   , parent_uuid(node.parent_uuid_)
   , inbound_edges(node.inbound_edges_)
@@ -58,6 +59,7 @@ bool TaskComposerNodeInfo::operator==(const TaskComposerNodeInfo& rhs) const
 
   bool equal = true;
   equal &= name == rhs.name;
+  equal &= ns == rhs.ns;
   equal &= uuid == rhs.uuid;
   equal &= parent_uuid == rhs.parent_uuid;
   equal &= return_value == rhs.return_value;
@@ -84,6 +86,7 @@ template <class Archive>
 void TaskComposerNodeInfo::serialize(Archive& ar, const unsigned int /*version*/)
 {
   ar& boost::serialization::make_nvp("name", name);
+  ar& boost::serialization::make_nvp("ns", ns);
   ar& boost::serialization::make_nvp("uuid", uuid);
   ar& boost::serialization::make_nvp("parent_uuid", parent_uuid);
   ar& boost::serialization::make_nvp("return_value", return_value);
@@ -160,6 +163,19 @@ TaskComposerNodeInfo::UPtr TaskComposerNodeInfoContainer::getInfo(const boost::u
   return it->second->clone();
 }
 
+std::vector<TaskComposerNodeInfo::UPtr>
+TaskComposerNodeInfoContainer::find(const std::function<bool(const TaskComposerNodeInfo&)>& search_fn) const
+{
+  std::unique_lock<std::shared_mutex> lock(mutex_);
+  std::vector<TaskComposerNodeInfo::UPtr> results;
+  for (const auto& info : info_map_)
+  {
+    if (search_fn(*info.second))
+      results.push_back(info.second->clone());
+  }
+  return results;
+}
+
 void TaskComposerNodeInfoContainer::setAborted(const boost::uuids::uuid& node_uuid)
 {
   assert(!node_uuid.is_nil());
@@ -180,6 +196,13 @@ void TaskComposerNodeInfoContainer::clear()
   info_map_.clear();
 }
 
+void TaskComposerNodeInfoContainer::prune(const std::function<void(TaskComposerNodeInfo& node_info)>& prune_fn)
+{
+  std::unique_lock<std::shared_mutex> lock(mutex_);
+  for (auto& info : info_map_)
+    prune_fn(*info.second);
+}
+
 std::map<boost::uuids::uuid, TaskComposerNodeInfo::UPtr> TaskComposerNodeInfoContainer::getInfoMap() const
 {
   std::shared_lock<std::shared_mutex> lock(mutex_);