From 798f7353cdb7659c9795ef458d4b7de3ef45bfd9 Mon Sep 17 00:00:00 2001
From: Sander Mertens <sander.mertens8@gmail.com>
Date: Thu, 6 Jul 2023 17:04:51 +0200
Subject: [PATCH] #1000 Fix issue with calling .get() on table/range from stage

---
 flecs.h                                |  4 +--
 include/flecs/addons/cpp/impl/iter.hpp |  4 +--
 test/cpp_api/project.json              |  4 ++-
 test/cpp_api/src/System.cpp            | 50 ++++++++++++++++++++++++++
 test/cpp_api/src/main.cpp              | 12 ++++++-
 5 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/flecs.h b/flecs.h
index e28ea826f..939b2f6b3 100644
--- a/flecs.h
+++ b/flecs.h
@@ -29666,11 +29666,11 @@ inline flecs::type iter::type() const {
 }
 
 inline flecs::table iter::table() const {
-    return flecs::table(m_iter->world, m_iter->table);
+    return flecs::table(m_iter->real_world, m_iter->table);
 }
 
 inline flecs::table_range iter::range() const {
-    return flecs::table_range(m_iter->world, m_iter->table, 
+    return flecs::table_range(m_iter->real_world, m_iter->table, 
         m_iter->offset, m_iter->count);
 }
 
diff --git a/include/flecs/addons/cpp/impl/iter.hpp b/include/flecs/addons/cpp/impl/iter.hpp
index b6cd30ccc..2094f9057 100644
--- a/include/flecs/addons/cpp/impl/iter.hpp
+++ b/include/flecs/addons/cpp/impl/iter.hpp
@@ -56,11 +56,11 @@ inline flecs::type iter::type() const {
 }
 
 inline flecs::table iter::table() const {
-    return flecs::table(m_iter->world, m_iter->table);
+    return flecs::table(m_iter->real_world, m_iter->table);
 }
 
 inline flecs::table_range iter::range() const {
-    return flecs::table_range(m_iter->world, m_iter->table, 
+    return flecs::table_range(m_iter->real_world, m_iter->table, 
         m_iter->offset, m_iter->count);
 }
 
diff --git a/test/cpp_api/project.json b/test/cpp_api/project.json
index ec326a81c..f48b4c359 100644
--- a/test/cpp_api/project.json
+++ b/test/cpp_api/project.json
@@ -475,7 +475,9 @@
                 "startup_system",
                 "interval_tick_source",
                 "rate_tick_source",
-                "nested_rate_tick_source"
+                "nested_rate_tick_source",
+                "table_get",
+                "range_get"
             ]
         }, {
             "id": "Event",
diff --git a/test/cpp_api/src/System.cpp b/test/cpp_api/src/System.cpp
index 425d8f403..19738f367 100644
--- a/test/cpp_api/src/System.cpp
+++ b/test/cpp_api/src/System.cpp
@@ -2146,3 +2146,53 @@ void System_nested_rate_tick_source() {
     test_int(2, sys_a_invoked);
     test_int(1, sys_b_invoked);
 }
+
+void System_table_get() {
+    flecs::world ecs;
+
+    flecs::entity e1 = ecs.entity().set<Position>({10, 20});
+    flecs::entity e2 = ecs.entity().set<Position>({20, 30});
+
+    auto s = ecs.system()
+        .with<Position>()
+        .each([&](flecs::iter& iter, size_t index) {
+            flecs::entity e = iter.entity(index);
+            const Position *p = &iter.table().get<Position>()[index];
+            test_assert(p != nullptr);
+            test_assert(e == e1 || e == e2);
+            if (e == e1) {
+                test_int(p->x, 10);
+                test_int(p->y, 20);
+            } else if (e == e2) {
+                test_int(p->x, 20);
+                test_int(p->y, 30);
+            }
+        });
+
+    s.run();
+}
+
+void System_range_get() {
+    flecs::world ecs;
+
+    flecs::entity e1 = ecs.entity().set<Position>({10, 20});
+    flecs::entity e2 = ecs.entity().set<Position>({20, 30});
+
+    auto s = ecs.system()
+        .with<Position>()
+        .each([&](flecs::iter& iter, size_t index) {
+            flecs::entity e = iter.entity(index);
+            const Position *p = &iter.range().get<Position>()[index];
+            test_assert(p != nullptr);
+            test_assert(e == e1 || e == e2);
+            if (e == e1) {
+                test_int(p->x, 10);
+                test_int(p->y, 20);
+            } else if (e == e2) {
+                test_int(p->x, 20);
+                test_int(p->y, 30);
+            }
+        });
+
+    s.run();
+}
diff --git a/test/cpp_api/src/main.cpp b/test/cpp_api/src/main.cpp
index 900e20904..f2d9e9211 100644
--- a/test/cpp_api/src/main.cpp
+++ b/test/cpp_api/src/main.cpp
@@ -458,6 +458,8 @@ void System_startup_system(void);
 void System_interval_tick_source(void);
 void System_rate_tick_source(void);
 void System_nested_rate_tick_source(void);
+void System_table_get(void);
+void System_range_get(void);
 
 // Testsuite 'Event'
 void Event_evt_1_id_entity(void);
@@ -3010,6 +3012,14 @@ bake_test_case System_testcases[] = {
     {
         "nested_rate_tick_source",
         System_nested_rate_tick_source
+    },
+    {
+        "table_get",
+        System_table_get
+    },
+    {
+        "range_get",
+        System_range_get
     }
 };
 
@@ -6091,7 +6101,7 @@ static bake_test_suite suites[] = {
         "System",
         NULL,
         NULL,
-        64,
+        66,
         System_testcases
     },
     {