From 05624941c5b273121949a2cc0029d3f68a5d37b5 Mon Sep 17 00:00:00 2001
From: Captain Yoshi <captain.yoshisaur@gmail.com>
Date: Mon, 29 Jan 2024 09:32:16 -0500
Subject: [PATCH] Quick access of stages by id/index

---
 core/include/moveit/task_constructor/task.h   |  6 ++++++
 core/include/moveit/task_constructor/task_p.h |  1 +
 core/src/task.cpp                             | 16 ++++++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/core/include/moveit/task_constructor/task.h b/core/include/moveit/task_constructor/task.h
index d1c123740..f04d7f563 100644
--- a/core/include/moveit/task_constructor/task.h
+++ b/core/include/moveit/task_constructor/task.h
@@ -98,6 +98,9 @@ class Task : protected WrapperBase
 	void insert(Stage::pointer&& stage, int before = -1) override;
 	void clear() final;
 
+	/// creates a flat buffer of stages in breath-first search order
+	void fillFlatBufferStages();
+
 	/// enable introspection publishing for use with rviz
 	void enableIntrospection(bool enable = true);
 	Introspection& introspection();
@@ -150,6 +153,7 @@ class Task : protected WrapperBase
 	/// access stage tree
 	ContainerBase* stages();
 	const ContainerBase* stages() const;
+	const ContainerBase* stages(uint32_t stage_id) const;
 
 	/// properties access
 	PropertyMap& properties();
@@ -165,6 +169,8 @@ class Task : protected WrapperBase
 
 private:
 	using WrapperBase::init;
+
+	std::vector<const ContainerBase*> bfs_stages_;  // breadth-first search
 };
 
 inline std::ostream& operator<<(std::ostream& os, const Task& task) {
diff --git a/core/include/moveit/task_constructor/task_p.h b/core/include/moveit/task_constructor/task_p.h
index ef42dd346..fe0ca8e68 100644
--- a/core/include/moveit/task_constructor/task_p.h
+++ b/core/include/moveit/task_constructor/task_p.h
@@ -58,6 +58,7 @@ class TaskPrivate : public WrapperBasePrivate
 
 	const std::string& ns() const { return ns_; }
 	const ContainerBase* stages() const;
+	const ContainerBase* stages(uint32_t stage_id) const;
 
 private:
 	std::string ns_;
diff --git a/core/src/task.cpp b/core/src/task.cpp
index ab375d2c7..2bd2ed019 100644
--- a/core/src/task.cpp
+++ b/core/src/task.cpp
@@ -92,6 +92,13 @@ const ContainerBase* TaskPrivate::stages() const {
 	return children().empty() ? nullptr : static_cast<ContainerBase*>(children().front().get());
 }
 
+const ContainerBase* Task::stages(uint32_t stage_id) const {
+	if (stage_id < bfs_stages_.size()) {
+		return bfs_stages_[stage_id];
+	}
+	return nullptr;
+}
+
 Task::Task(const std::string& ns, bool introspection, ContainerBase::pointer&& container)
   : WrapperBase(new TaskPrivate(this, ns), std::move(container)) {
 	setPruning(false);
@@ -143,6 +150,15 @@ void Task::loadRobotModel(const std::string& robot_description) {
 		throw Exception("Task failed to construct RobotModel");
 }
 
+void Task::fillFlatBufferStages() {
+	bfs_stages_.clear();
+	auto stage_cb = [&](const moveit::task_constructor::Stage& stage, unsigned int depth) -> bool {
+		bfs_stages_.push_back(static_cast<const ContainerBase*>(&stage));
+		return true;
+	};
+	stages()->traverseRecursively(stage_cb);
+}
+
 void Task::add(Stage::pointer&& stage) {
 	stages()->add(std::move(stage));
 }