Add entity_column

flecs.h

@@ -20642,6 +20642,56 @@ struct column {
     bool m_is_shared;
 };
 
+struct entity_column {
+    /** Create column from component array.
+     *
+     * @param array Pointer to the component array.
+     * @param count Number of elements in component array.
+     * @param is_shared Is the component shared or not.
+     */
+    entity_column(flecs::world_t ecs, flecs::entity_t* array, size_t count, bool is_shared = false)
+        : m_ecs(ecs)
+        , m_array(array)
+        , m_count(count) 
+        , m_is_shared(is_shared) {}
+
+    /** Create column from iterator.
+     *
+     * @param iter Iterator object.
+     * @param column Index of the signature of the query being iterated over.
+     */
+    entity_column(iter &iter, int column);
+
+    /** Return element in component array.
+     * This operator may only be used if the column is not shared.
+     * 
+     * @param index Index of element.
+     * @return Reference to element.
+     */
+    flecs::entity operator[](size_t index) const {
+        ecs_assert(index < m_count, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+        ecs_assert(!index || !m_is_shared, ECS_INVALID_PARAMETER, NULL);
+        ecs_assert(m_array != nullptr, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+        return flecs::entity{m_ecs, m_array[index]};
+    }
+
+    /** Return first element of component array.
+     * This operator is typically used when the column is shared.
+     *
+     * @return Reference to the first element.
+     */
+    flecs::entity operator*() const {
+      ecs_assert(m_array != nullptr, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+      return flecs::entity{m_ecs, m_array};
+    }
+
+protected:
+    flecs::world_t m_ecs;
+    flecs::entity_t* m_array;
+    size_t m_count;
+    bool m_is_shared;
+};
+
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -20890,6 +20940,19 @@ struct iter {
         return get_field<A>(index);
     }
 
+    /** Get readonly access to field data.
+     * If the specified field index does not match with the provided type, the
+     * function will assert.
+     *
+     * @tparam T Type of the field.
+     * @param index The field index.
+     * @return The field data.
+     */
+    template<>
+    flecs::entity_column field<flecs::entity>(int32_t index) const {
+        return get_entity_field<A>(index);
+    }
+
     /** Get unchecked access to field data.
      * Unchecked access is required when a system does not know the type of a
      * field at compile time.
@@ -20970,6 +21033,35 @@ struct iter {
             count, is_shared);
     }
 
+    /* Get field, check if correct type is used */
+    flecs::entity_column get_entity_field(int32_t index) const {
+
+#ifndef FLECS_NDEBUG
+        ecs_entity_t term_id = ecs_field_id(m_iter, index);
+        ecs_assert(ECS_HAS_ID_FLAG(term_id, PAIR) ||
+            term_id == _::cpp_type<flecs::entity_t>::id(m_iter->world), 
+            ECS_COLUMN_TYPE_MISMATCH, NULL);
+#endif
+
+        size_t count;
+        bool is_shared = !ecs_field_is_self(m_iter, index);
+
+        /* If a shared column is retrieved with 'column', there will only be a
+         * single value. Ensure that the application does not accidentally read
+         * out of bounds. */
+        if (is_shared) {
+            count = 1;
+        } else {
+            /* If column is owned, there will be as many values as there are
+             * entities. */
+            count = static_cast<size_t>(m_iter->count);
+        }
+
+        return flecs::entity_column(m_iter->world, 
+            static_cast<flecs::entity_t*>(ecs_field_w_size(m_iter, sizeof(flecs::entity_t), index)), 
+            count, is_shared);
+    }
+
     flecs::untyped_column get_unchecked_field(int32_t index) const {
         size_t count;
         size_t size = ecs_field_size(m_iter, index);
@@ -29645,6 +29737,10 @@ inline column<T>::column(iter &iter, int32_t index) {
     *this = iter.field<T>(index);
 }
 
+inline entity_column::entity_column(iter &iter, int32_t index) {
+    *this = iter.field<flecs::entity>(index);
+}
+
 inline flecs::entity iter::src(int32_t index) const {
     return flecs::entity(m_iter->world, ecs_field_src(m_iter, index));
 }

include/flecs/addons/cpp/impl/iter.hpp

@@ -35,6 +35,10 @@ inline column<T>::column(iter &iter, int32_t index) {
     *this = iter.field<T>(index);
 }
 
+inline entity_column::entity_column(iter &iter, int32_t index) {
+    *this = iter.field<flecs::entity>(index);
+}
+
 inline flecs::entity iter::src(int32_t index) const {
     return flecs::entity(m_iter->world, ecs_field_src(m_iter, index));
 }

include/flecs/addons/cpp/iter.hpp

@@ -116,6 +116,55 @@ struct column {
     bool m_is_shared;
 };
 
+struct entity_column {
+    /** Create column from component array.
+     *
+     * @param array Pointer to the component array.
+     * @param count Number of elements in component array.
+     * @param is_shared Is the component shared or not.
+     */
+    entity_column(flecs::world_t ecs, flecs::entity_t* array, size_t count, bool is_shared = false)
+        : m_ecs(ecs)
+        , m_array(array)
+        , m_count(count) 
+        , m_is_shared(is_shared) {}
+
+    /** Create column from iterator.
+     *
+     * @param iter Iterator object.
+     * @param column Index of the signature of the query being iterated over.
+     */
+    entity_column(iter &iter, int column);
+
+    /** Return element in component array.
+     * This operator may only be used if the column is not shared.
+     * 
+     * @param index Index of element.
+     * @return Reference to element.
+     */
+    flecs::entity& operator[](size_t index) const {
+        ecs_assert(index < m_count, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+        ecs_assert(!index || !m_is_shared, ECS_INVALID_PARAMETER, NULL);
+        ecs_assert(m_array != nullptr, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+        return flecs::entity{m_ecs, m_array[index]};
+    }
+
+    /** Return first element of component array.
+     * This operator is typically used when the column is shared.
+     *
+     * @return Reference to the first element.
+     */
+    flecs::entity& operator*() const {
+      ecs_assert(m_array != nullptr, ECS_COLUMN_INDEX_OUT_OF_RANGE, NULL);
+      return flecs::entity{m_ecs, m_array};
+    }
+
+protected:
+    flecs::world_t m_ecs;
+    flecs::entity_t* m_array;
+    size_t m_count;
+    bool m_is_shared;
+};
 
 ////////////////////////////////////////////////////////////////////////////////
 
@@ -364,6 +413,19 @@ struct iter {
         return get_field<A>(index);
     }
 
+    /** Get readonly access to field data.
+     * If the specified field index does not match with the provided type, the
+     * function will assert.
+     *
+     * @tparam T Type of the field.
+     * @param index The field index.
+     * @return The field data.
+     */
+    template<>
+    flecs::entity_column field<flecs::entity>(int32_t index) const {
+        return get_entity_field<A>(index);
+    }
+
     /** Get unchecked access to field data.
      * Unchecked access is required when a system does not know the type of a
      * field at compile time.
@@ -444,6 +506,35 @@ struct iter {
             count, is_shared);
     }
 
+    /* Get field, check if correct type is used */
+    flecs::entity_column get_entity_field(int32_t index) const {
+
+#ifndef FLECS_NDEBUG
+        ecs_entity_t term_id = ecs_field_id(m_iter, index);
+        ecs_assert(ECS_HAS_ID_FLAG(term_id, PAIR) ||
+            term_id == _::cpp_type<flecs::entity_t>::id(m_iter->world), 
+            ECS_COLUMN_TYPE_MISMATCH, NULL);
+#endif
+
+        size_t count;
+        bool is_shared = !ecs_field_is_self(m_iter, index);
+
+        /* If a shared column is retrieved with 'column', there will only be a
+         * single value. Ensure that the application does not accidentally read
+         * out of bounds. */
+        if (is_shared) {
+            count = 1;
+        } else {
+            /* If column is owned, there will be as many values as there are
+             * entities. */
+            count = static_cast<size_t>(m_iter->count);
+        }
+
+        return flecs::entity_column(m_iter->world, 
+            static_cast<flecs::entity_t*>(ecs_field_w_size(m_iter, sizeof(flecs::entity_t), index)), 
+            count, is_shared);
+    }
+
     flecs::untyped_column get_unchecked_field(int32_t index) const {
         size_t count;
         size_t size = ecs_field_size(m_iter, index);
@@ -471,4 +562,4 @@ struct iter {
 
 } // namespace flecs
 
-/** @} */
+/** @} */

test/cpp_api/src/Iterable.cpp

@@ -118,3 +118,38 @@ void Iterable_worker_iter() {
 
     test_int(count, 2);
 }
+
+
+struct UnionTag {};
+
+void Iterable_entity_column_iter() {
+    flecs::world ecs;
+
+    ecs.component<UnionTag>()
+        .add(flecs::Union);
+
+    auto e1 = ecs.entity();
+    auto e2 = ecs.entity().add<UnionTag>(e1);
+    auto e3 = ecs.entity().add<UnionTag>(e2);
+    auto e4 = ecs.entity().add<UnionTag>(e3);
+
+    ecs.query<>()
+        .with<UnionTag>()
+        .iter([&](flecs::iter it) {
+            test_int(it.count(), 3);
+
+            test_assert(it.entity(0) == e2);
+            test_assert(it.entity(1) == e3);
+            test_assert(it.entity(2) == e4);
+
+            auto column = it.field<flecs::entity_t>();
+            test_assert(column[0] == e1);
+            test_assert(column[1] == e2);
+            test_assert(column[2] == e3);
+
+            auto column = it.field<flecs::entity>();
+            test_assert(column[0].id() == e1);
+            test_assert(column[1].id() == e2);
+            test_assert(column[2].id() == e3);
+        });
+}