From 3ba680a665758c4ad70828795b633a427ff62d0b Mon Sep 17 00:00:00 2001
From: Nick Drabsch <zero@hypixelstudios.com>
Date: Sun, 9 Jul 2023 12:58:24 +0100
Subject: [PATCH] Add flecs::iter::entities() function

---
 flecs.h                           |  8 ++++++++
 include/flecs/addons/cpp/iter.hpp | 10 +++++++++-
 test/cpp_api/src/Query.cpp        | 18 ++++++++++++++++++
 test/cpp_api/src/main.cpp         |  7 ++++++-
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/flecs.h b/flecs.h
index 939b2f6b3..17b704456 100644
--- a/flecs.h
+++ b/flecs.h
@@ -20900,6 +20900,14 @@ struct iter {
         return get_unchecked_field(index);
     }
 
+    /** Get readonly access to entity ids.
+     *
+     * @return The entity ids.
+     */
+    flecs::column<const flecs::entity_t> entities() const {
+        return flecs::column<const flecs::entity_t>(m_iter->entities, static_cast<size_t>(m_iter->count), false);
+    }
+
     /** Obtain the total number of tables the iterator will iterate over. */
     int32_t table_count() const {
         return m_iter->table_count;
diff --git a/include/flecs/addons/cpp/iter.hpp b/include/flecs/addons/cpp/iter.hpp
index 8e5d7a2ba..c88e1d4cf 100644
--- a/include/flecs/addons/cpp/iter.hpp
+++ b/include/flecs/addons/cpp/iter.hpp
@@ -374,6 +374,14 @@ struct iter {
         return get_unchecked_field(index);
     }
 
+    /** Get readonly access to entity ids.
+     *
+     * @return The entity ids.
+     */
+    flecs::column<const flecs::entity_t> entities() const {
+        return flecs::column<const flecs::entity_t>(m_iter->entities, static_cast<size_t>(m_iter->count), false);
+    }
+
     /** Obtain the total number of tables the iterator will iterate over. */
     int32_t table_count() const {
         return m_iter->table_count;
@@ -471,4 +479,4 @@ struct iter {
 
 } // namespace flecs
 
-/** @} */
\ No newline at end of file
+/** @} */
diff --git a/test/cpp_api/src/Query.cpp b/test/cpp_api/src/Query.cpp
index b4a56b096..807bb5f3b 100644
--- a/test/cpp_api/src/Query.cpp
+++ b/test/cpp_api/src/Query.cpp
@@ -2229,3 +2229,21 @@ void Query_worker_iter_captured_query() {
         test_int(count, 1);
     }();
 }
+
+void Query_iter_entities() {
+    flecs::world ecs;
+
+    auto e1 = ecs.entity().set<Position>({10, 20});
+    auto e2 = ecs.entity().set<Position>({10, 20});
+    auto e3 = ecs.entity().set<Position>({10, 20});
+
+    ecs.query<Position>()
+        .iter([&](flecs::iter& it) {
+            test_int(it.count(), 3);
+
+            auto entities = it.entities();
+            test_assert(entities[0] == e1);
+            test_assert(entities[1] == e2);
+            test_assert(entities[2] == e3);
+        });
+}
diff --git a/test/cpp_api/src/main.cpp b/test/cpp_api/src/main.cpp
index f2d9e9211..2d8effd6a 100644
--- a/test/cpp_api/src/main.cpp
+++ b/test/cpp_api/src/main.cpp
@@ -563,6 +563,7 @@ void Query_instanced_nested_query_w_world(void);
 void Query_captured_query(void);
 void Query_page_iter_captured_query(void);
 void Query_worker_iter_captured_query(void);
+void Query_iter_entities(void);
 
 // Testsuite 'QueryBuilder'
 void QueryBuilder_builder_assign_same_type(void);
@@ -3417,6 +3418,10 @@ bake_test_case Query_testcases[] = {
     {
         "worker_iter_captured_query",
         Query_worker_iter_captured_query
+    },
+    {
+        "iter_entities",
+        Query_iter_entities
     }
 };
 
@@ -6122,7 +6127,7 @@ static bake_test_suite suites[] = {
         "Query",
         NULL,
         NULL,
-        78,
+        79,
         Query_testcases
     },
     {