#861 Add support for setting this variable in queries

SanderMertens · Aug 24, 2023 · 41fb856 · 41fb856
1 parent d1541f1
commit 41fb856
Show file tree

Hide file tree

Showing 9 changed files with 274 additions and 20 deletions.
diff --git a/flecs.c b/flecs.c
@@ -13329,6 +13329,10 @@ void ecs_iter_set_var(
 
     it->constrained_vars |= flecs_ito(uint64_t, 1 << var_id);
 
+    if (it->set_var) {
+        it->set_var(it);
+    }
+
 error:
     return;
 }
@@ -19234,6 +19238,38 @@ const ecs_filter_t* ecs_query_get_filter(
     return &query->filter;
 }
 
+static
+void flecs_query_set_var(
+    ecs_iter_t *it)
+{
+    ecs_check(it->constrained_vars == 1, ECS_INVALID_OPERATION, 
+        "can only set $this variable for queries");
+
+    ecs_var_t *var = &it->variables[0];
+    ecs_table_t *table = var->range.table;
+    if (!table) {
+        goto nodata;
+    }
+
+    ecs_query_iter_t *qit = &it->priv.iter.query;
+    ecs_query_t *query = qit->query;
+    ecs_query_table_t *qt = ecs_table_cache_get(&query->cache, table);
+    if (!qt) {
+        goto nodata;
+    }
+
+    qit->node = qt->first;
+    qit->last = qt->last->next_match;
+    it->offset = var->range.offset;
+    it->count = var->range.count;
+    return;
+error:
+nodata:
+    it->priv.iter.query.node = NULL;
+    it->priv.iter.query.last = NULL;
+    return;
+}
+
 ecs_iter_t ecs_query_iter(
     const ecs_world_t *stage,
     ecs_query_t *query)
@@ -19282,8 +19318,10 @@ ecs_iter_t ecs_query_iter(
         .terms = query->filter.terms,
         .field_count = query->filter.field_count,
         .table_count = table_count,
+        .variable_count = 1,
         .priv.iter.query = it,
         .next = ecs_query_next,
+        .set_var = flecs_query_set_var
     };
 
     flecs_filter_apply_iter_flags(&result, &query->filter);
@@ -19318,7 +19356,8 @@ ecs_iter_t ecs_query_iter(
         }
     } else {
         /* Trivial iteration, use arrays from query cache */
-        flecs_iter_init(stage, &result, flecs_iter_cache_ptrs);
+        flecs_iter_init(stage, &result, 
+            flecs_iter_cache_ptrs|flecs_iter_cache_variables);
     }
 
     result.sizes = query->filter.sizes;
@@ -19476,15 +19515,20 @@ void flecs_query_populate_trivial(
     ecs_query_table_match_t *match)
 {;
     ecs_table_t *table = match->table;
-    int32_t count = ecs_table_count(table);
+    int32_t offset, count;
+    if (!it->constrained_vars) {
+        it->offset = offset = 0;
+        it->count = count = ecs_table_count(table);
+    } else {
+        offset = it->offset;
+        count = it->count;
+    }
 
     it->ids = match->ids;
     it->sources = match->sources;
     it->columns = match->columns;
     it->group_id = match->group_id;
     it->instance_count = 0;
-    it->offset = 0;
-    it->count = count;
     it->references = ecs_vec_first(&match->refs);
 
     if (!it->references) {
@@ -19505,15 +19549,16 @@ void flecs_query_populate_trivial(
                 }
 
                 it->ptrs[i] = ecs_vec_get(&data->columns[column].data,
-                    it->sizes[i], 0);
+                    it->sizes[i], offset);
             }
         }
 
         it->frame_offset += it->table ? ecs_table_count(it->table) : 0;
         it->table = table;
-        it->entities = ecs_vec_first(&data->entities);
+        it->entities = ecs_vec_get_t(&data->entities, ecs_entity_t, offset);
     } else {
-        flecs_iter_populate_data(it->real_world, it, table, 0, count, it->ptrs);
+        flecs_iter_populate_data(
+            it->real_world, it, table, offset, count, it->ptrs);
     }
 }
 
@@ -51765,7 +51810,7 @@ void flecs_json_serialize_iter_variables(ecs_iter_t *it, ecs_strbuf_t *buf) {
     int32_t var_count = it->variable_count;
     int32_t actual_count = 0;
 
-    for (int i = 0; i < var_count; i ++) {
+    for (int i = 1; i < var_count; i ++) {
         const char *var_name = variable_names[i];
         if (flecs_json_skip_variable(var_name)) continue;
 
@@ -51897,7 +51942,7 @@ void flecs_json_serialize_iter_result_variables(
     int32_t var_count = it->variable_count;
     int32_t actual_count = 0;
 
-    for (int i = 0; i < var_count; i ++) {
+    for (int i = 1; i < var_count; i ++) {
         const char *var_name = variable_names[i];
         if (flecs_json_skip_variable(var_name)) continue;
 

diff --git a/flecs.h b/flecs.h
@@ -3274,6 +3274,7 @@ struct ecs_iter_t {
     /* Chained iterators */
     ecs_iter_next_action_t next;  /* Function to progress iterator */
     ecs_iter_action_t callback;   /* Callback of system or observer */
+    ecs_iter_action_t set_var;    /* Invoked after setting variable (optionally set) */
     ecs_iter_fini_action_t fini;  /* Function to cleanup iterator resources */
     ecs_iter_t *chain_it;         /* Optional, allows for creating iterator chains */
 };

diff --git a/include/flecs/private/api_types.h b/include/flecs/private/api_types.h
@@ -302,6 +302,7 @@ struct ecs_iter_t {
     /* Chained iterators */
     ecs_iter_next_action_t next;  /* Function to progress iterator */
     ecs_iter_action_t callback;   /* Callback of system or observer */
+    ecs_iter_action_t set_var;    /* Invoked after setting variable (optionally set) */
     ecs_iter_fini_action_t fini;  /* Function to cleanup iterator resources */
     ecs_iter_t *chain_it;         /* Optional, allows for creating iterator chains */
 };

diff --git a/src/addons/json/serialize.c b/src/addons/json/serialize.c
@@ -1420,7 +1420,7 @@ void flecs_json_serialize_iter_variables(ecs_iter_t *it, ecs_strbuf_t *buf) {
     int32_t var_count = it->variable_count;
     int32_t actual_count = 0;
 
-    for (int i = 0; i < var_count; i ++) {
+    for (int i = 1; i < var_count; i ++) {
         const char *var_name = variable_names[i];
         if (flecs_json_skip_variable(var_name)) continue;
 
@@ -1552,7 +1552,7 @@ void flecs_json_serialize_iter_result_variables(
     int32_t var_count = it->variable_count;
     int32_t actual_count = 0;
 
-    for (int i = 0; i < var_count; i ++) {
+    for (int i = 1; i < var_count; i ++) {
         const char *var_name = variable_names[i];
         if (flecs_json_skip_variable(var_name)) continue;
 

diff --git a/src/iter.c b/src/iter.c
@@ -763,6 +763,10 @@ void ecs_iter_set_var(
 
     it->constrained_vars |= flecs_ito(uint64_t, 1 << var_id);
 
+    if (it->set_var) {
+        it->set_var(it);
+    }
+
 error:
     return;
 }

diff --git a/src/query.c b/src/query.c
@@ -2182,6 +2182,38 @@ const ecs_filter_t* ecs_query_get_filter(
     return &query->filter;
 }
 
+static
+void flecs_query_set_var(
+    ecs_iter_t *it)
+{
+    ecs_check(it->constrained_vars == 1, ECS_INVALID_OPERATION, 
+        "can only set $this variable for queries");
+
+    ecs_var_t *var = &it->variables[0];
+    ecs_table_t *table = var->range.table;
+    if (!table) {
+        goto nodata;
+    }
+
+    ecs_query_iter_t *qit = &it->priv.iter.query;
+    ecs_query_t *query = qit->query;
+    ecs_query_table_t *qt = ecs_table_cache_get(&query->cache, table);
+    if (!qt) {
+        goto nodata;
+    }
+
+    qit->node = qt->first;
+    qit->last = qt->last->next_match;
+    it->offset = var->range.offset;
+    it->count = var->range.count;
+    return;
+error:
+nodata:
+    it->priv.iter.query.node = NULL;
+    it->priv.iter.query.last = NULL;
+    return;
+}
+
 ecs_iter_t ecs_query_iter(
     const ecs_world_t *stage,
     ecs_query_t *query)
@@ -2230,8 +2262,10 @@ ecs_iter_t ecs_query_iter(
         .terms = query->filter.terms,
         .field_count = query->filter.field_count,
         .table_count = table_count,
+        .variable_count = 1,
         .priv.iter.query = it,
         .next = ecs_query_next,
+        .set_var = flecs_query_set_var
     };
 
     flecs_filter_apply_iter_flags(&result, &query->filter);
@@ -2266,7 +2300,8 @@ ecs_iter_t ecs_query_iter(
         }
     } else {
         /* Trivial iteration, use arrays from query cache */
-        flecs_iter_init(stage, &result, flecs_iter_cache_ptrs);
+        flecs_iter_init(stage, &result, 
+            flecs_iter_cache_ptrs|flecs_iter_cache_variables);
     }
 
     result.sizes = query->filter.sizes;
@@ -2424,15 +2459,20 @@ void flecs_query_populate_trivial(
     ecs_query_table_match_t *match)
 {;
     ecs_table_t *table = match->table;
-    int32_t count = ecs_table_count(table);
+    int32_t offset, count;
+    if (!it->constrained_vars) {
+        it->offset = offset = 0;
+        it->count = count = ecs_table_count(table);
+    } else {
+        offset = it->offset;
+        count = it->count;
+    }
 
     it->ids = match->ids;
     it->sources = match->sources;
     it->columns = match->columns;
     it->group_id = match->group_id;
     it->instance_count = 0;
-    it->offset = 0;
-    it->count = count;
     it->references = ecs_vec_first(&match->refs);
 
     if (!it->references) {
@@ -2453,15 +2493,16 @@ void flecs_query_populate_trivial(
                 }
 
                 it->ptrs[i] = ecs_vec_get(&data->columns[column].data,
-                    it->sizes[i], 0);
+                    it->sizes[i], offset);
             }
         }
 
         it->frame_offset += it->table ? ecs_table_count(it->table) : 0;
         it->table = table;
-        it->entities = ecs_vec_first(&data->entities);
+        it->entities = ecs_vec_get_t(&data->entities, ecs_entity_t, offset);
     } else {
-        flecs_iter_populate_data(it->real_world, it, table, 0, count, it->ptrs);
+        flecs_iter_populate_data(
+            it->real_world, it, table, offset, count, it->ptrs);
     }
 }
 

diff --git a/test/api/project.json b/test/api/project.json
@@ -1657,7 +1657,11 @@
                 "set_get_context",
                 "set_get_binding_context",
                 "set_get_context_w_free",
-                "set_get_binding_context_w_free"
+                "set_get_binding_context_w_free",
+                "set_this",
+                "set_this_no_match",
+                "set_this_is_true",
+                "set_this_w_wildcard"
             ]
         }, {
             "id": "Iter",