From 1050c12f36666bc1c7f162f0bd81eccc91099e1a Mon Sep 17 00:00:00 2001
From: Sander Mertens <sander.mertens8@gmail.com>
Date: Mon, 16 Dec 2024 20:59:35 -0800
Subject: [PATCH] Add ecs_query_get

---
 distr/flecs.c           | 13 +++++++++++++
 distr/flecs.h           | 13 +++++++++++++
 include/flecs.h         | 13 +++++++++++++
 src/query/api.c         | 13 +++++++++++++
 test/query/project.json |  4 +++-
 test/query/src/Basic.c  | 30 ++++++++++++++++++++++++++++++
 test/query/src/main.c   | 12 +++++++++++-
 7 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/distr/flecs.c b/distr/flecs.c
index 64e0a038a..7012edc92 100644
--- a/distr/flecs.c
+++ b/distr/flecs.c
@@ -33672,6 +33672,19 @@ const ecs_query_t* ecs_query_get_cache_query(
     }
 }
 
+const ecs_query_t* ecs_query_get(
+    const ecs_world_t *world,
+    ecs_entity_t query)
+{
+    const EcsPoly *poly_comp = ecs_get_pair(world, query, EcsPoly, EcsQuery);
+    if (!poly_comp) {
+        return NULL;
+    } else {
+        flecs_poly_assert(poly_comp->poly, ecs_query_t);
+        return poly_comp->poly;
+    }
+}
+
  /**
  * @file query/util.c
  * @brief Query utilities.
diff --git a/distr/flecs.h b/distr/flecs.h
index a3b51e15f..9e3b7341d 100644
--- a/distr/flecs.h
+++ b/distr/flecs.h
@@ -8136,6 +8136,19 @@ FLECS_API
 bool ecs_query_changed(
     ecs_query_t *query);
 
+/** Get query object.
+ * Returns the query object. Can be used to access various information about
+ * the query.
+ *
+ * @param world The world.
+ * @param query The query.
+ * @return The query object.
+ */
+FLECS_API
+const ecs_query_t* ecs_query_get(
+    const ecs_world_t *world,
+    ecs_entity_t query);
+
 /** Skip a table while iterating.
  * This operation lets the query iterator know that a table was skipped while
  * iterating. A skipped table will not reset its changed state, and the query
diff --git a/include/flecs.h b/include/flecs.h
index 8219c9925..a7c31f83c 100644
--- a/include/flecs.h
+++ b/include/flecs.h
@@ -4896,6 +4896,19 @@ FLECS_API
 bool ecs_query_changed(
     ecs_query_t *query);
 
+/** Get query object.
+ * Returns the query object. Can be used to access various information about
+ * the query.
+ *
+ * @param world The world.
+ * @param query The query.
+ * @return The query object.
+ */
+FLECS_API
+const ecs_query_t* ecs_query_get(
+    const ecs_world_t *world,
+    ecs_entity_t query);
+
 /** Skip a table while iterating.
  * This operation lets the query iterator know that a table was skipped while
  * iterating. A skipped table will not reset its changed state, and the query
diff --git a/src/query/api.c b/src/query/api.c
index 3af9ff1cc..6cfa24aac 100644
--- a/src/query/api.c
+++ b/src/query/api.c
@@ -541,3 +541,16 @@ const ecs_query_t* ecs_query_get_cache_query(
         return impl->cache->query;
     }
 }
+
+const ecs_query_t* ecs_query_get(
+    const ecs_world_t *world,
+    ecs_entity_t query)
+{
+    const EcsPoly *poly_comp = ecs_get_pair(world, query, EcsPoly, EcsQuery);
+    if (!poly_comp) {
+        return NULL;
+    } else {
+        flecs_poly_assert(poly_comp->poly, ecs_query_t);
+        return poly_comp->poly;
+    }
+}
diff --git a/test/query/project.json b/test/query/project.json
index 319b68525..7289c8460 100644
--- a/test/query/project.json
+++ b/test/query/project.json
@@ -697,7 +697,9 @@
                 "entity_iteration_w_match_empty_tables",
                 "get_cache_query_uncached",
                 "get_cache_query_cached",
-                "get_cache_query_partially_cached"
+                "get_cache_query_partially_cached",
+                "get_query",
+                "get_query_not_a_query"
             ]
         }, {
             "id": "Combinations",
diff --git a/test/query/src/Basic.c b/test/query/src/Basic.c
index 15859eea4..5b78703f8 100644
--- a/test/query/src/Basic.c
+++ b/test/query/src/Basic.c
@@ -11428,3 +11428,33 @@ void Basic_get_cache_query_partially_cached(void) {
 
     ecs_fini(world);
 }
+
+void Basic_get_query(void) {
+    ecs_world_t *world = ecs_mini();
+
+    ECS_COMPONENT(world, Position);
+
+    ecs_entity_t qe = ecs_new(world);
+
+    ecs_query_t *q = ecs_query(world, {
+        .entity = qe,
+        .expr = "Position",
+        .cache_kind = EcsQueryCacheAuto
+    });
+
+    test_assert(q != NULL);
+    test_assert(q->entity == qe);
+    test_assert(ecs_query_get(world, qe) == q);
+
+    ecs_fini(world);
+}
+
+void Basic_get_query_not_a_query(void) {
+    ecs_world_t *world = ecs_mini();
+
+    ecs_entity_t qe = ecs_new(world);
+
+    test_assert(ecs_query_get(world, qe) == NULL);
+
+    ecs_fini(world);
+}
diff --git a/test/query/src/main.c b/test/query/src/main.c
index f64723648..810942b5c 100644
--- a/test/query/src/main.c
+++ b/test/query/src/main.c
@@ -686,6 +686,8 @@ void Basic_entity_iteration_w_match_empty_tables(void);
 void Basic_get_cache_query_uncached(void);
 void Basic_get_cache_query_cached(void);
 void Basic_get_cache_query_partially_cached(void);
+void Basic_get_query(void);
+void Basic_get_query_not_a_query(void);
 
 // Testsuite 'Combinations'
 void Combinations_setup(void);
@@ -4815,6 +4817,14 @@ bake_test_case Basic_testcases[] = {
     {
         "get_cache_query_partially_cached",
         Basic_get_cache_query_partially_cached
+    },
+    {
+        "get_query",
+        Basic_get_query
+    },
+    {
+        "get_query_not_a_query",
+        Basic_get_query_not_a_query
     }
 };
 
@@ -10477,7 +10487,7 @@ static bake_test_suite suites[] = {
         "Basic",
         Basic_setup,
         NULL,
-        230,
+        232,
         Basic_testcases,
         1,
         Basic_params