From 3f9c4bf8137fef7ef09dea5d8e58ad802a029193 Mon Sep 17 00:00:00 2001
From: Sander Mertens <sander.mertens8@gmail.com>
Date: Wed, 11 Dec 2024 17:57:53 +0000
Subject: [PATCH] Fix leaks when using FLECS_USE_OS_ALLOC

---
 distr/flecs.c                | 17 ++++++++++++-----
 distr/flecs.h                |  2 +-
 include/flecs.h              |  2 +-
 src/datastructures/hashmap.c |  2 +-
 src/datastructures/map.c     |  2 +-
 src/query/engine/eval_iter.c |  4 +++-
 src/world.c                  |  9 +++++++--
 7 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/distr/flecs.c b/distr/flecs.c
index e8b5ac455f..1a408490c9 100644
--- a/distr/flecs.c
+++ b/distr/flecs.c
@@ -20662,8 +20662,13 @@ ecs_entities_t ecs_get_entities(
 ecs_flags32_t ecs_world_get_flags(
     const ecs_world_t *world)
 {
-    flecs_poly_assert(world, ecs_world_t);
-    return world->flags;
+    if (flecs_poly_is(world, ecs_world_t)) {
+        return world->flags;
+    } else {
+        flecs_poly_assert(world, ecs_stage_t);
+        const ecs_stage_t *stage = (const ecs_stage_t*)world;
+        return stage->world->flags;
+    }
 }
 
 /**
@@ -29984,7 +29989,7 @@ void flecs_hashmap_fini(
         ecs_hm_bucket_t *bucket = ecs_map_ptr(&it);
         ecs_vec_fini(a, &bucket->keys, map->key_size);
         ecs_vec_fini(a, &bucket->values, map->value_size);
-#ifdef FLECS_SANITIZE        
+#if defined(FLECS_SANITIZE) || defined(FLECS_USE_OS_ALLOC)
         flecs_bfree(&map->bucket_allocator, bucket);
 #endif
     }
@@ -30439,7 +30444,7 @@ void ecs_map_fini(
     }
 
     bool sanitize = false;
-#ifdef FLECS_SANITIZE
+#if defined(FLECS_SANITIZE) || defined(FLECS_USE_OS_ALLOC)
     sanitize = true;
 #endif
 
@@ -71024,7 +71029,9 @@ ecs_iter_t ecs_query_iter(
     if (cache) {
         /* If monitors changed, do query rematching */
         ecs_flags32_t flags = q->flags;
-        if (!(world->flags & EcsWorldReadonly) && flags & EcsQueryHasRefs) {
+        if (!(ecs_world_get_flags(world) & EcsWorldReadonly) && 
+             (flags & EcsQueryHasRefs)) 
+        {
             flecs_eval_component_monitors(q->world);
         }
     }
diff --git a/distr/flecs.h b/distr/flecs.h
index 7c4b9951e9..f7947fd656 100644
--- a/distr/flecs.h
+++ b/distr/flecs.h
@@ -265,7 +265,7 @@
  * When enabled, Flecs will use the OS allocator provided in the OS API directly
  * instead of the builtin block allocator. This can decrease memory utilization
  * as memory will be freed more often, at the cost of decreased performance. */
-#define FLECS_USE_OS_ALLOC
+// #define FLECS_USE_OS_ALLOC
 
 /** @def FLECS_ID_DESC_MAX
  * Maximum number of ids to add ecs_entity_desc_t / ecs_bulk_desc_t */
diff --git a/include/flecs.h b/include/flecs.h
index ec96077c27..027abc2392 100644
--- a/include/flecs.h
+++ b/include/flecs.h
@@ -263,7 +263,7 @@
  * When enabled, Flecs will use the OS allocator provided in the OS API directly
  * instead of the builtin block allocator. This can decrease memory utilization
  * as memory will be freed more often, at the cost of decreased performance. */
-#define FLECS_USE_OS_ALLOC
+// #define FLECS_USE_OS_ALLOC
 
 /** @def FLECS_ID_DESC_MAX
  * Maximum number of ids to add ecs_entity_desc_t / ecs_bulk_desc_t */
diff --git a/src/datastructures/hashmap.c b/src/datastructures/hashmap.c
index 987e9e6a6d..e60c39ce27 100644
--- a/src/datastructures/hashmap.c
+++ b/src/datastructures/hashmap.c
@@ -53,7 +53,7 @@ void flecs_hashmap_fini(
         ecs_hm_bucket_t *bucket = ecs_map_ptr(&it);
         ecs_vec_fini(a, &bucket->keys, map->key_size);
         ecs_vec_fini(a, &bucket->values, map->value_size);
-#ifdef FLECS_SANITIZE        
+#if defined(FLECS_SANITIZE) || defined(FLECS_USE_OS_ALLOC)
         flecs_bfree(&map->bucket_allocator, bucket);
 #endif
     }
diff --git a/src/datastructures/map.c b/src/datastructures/map.c
index 271dcfd252..cda7b19219 100644
--- a/src/datastructures/map.c
+++ b/src/datastructures/map.c
@@ -243,7 +243,7 @@ void ecs_map_fini(
     }
 
     bool sanitize = false;
-#ifdef FLECS_SANITIZE
+#if defined(FLECS_SANITIZE) || defined(FLECS_USE_OS_ALLOC)
     sanitize = true;
 #endif
 
diff --git a/src/query/engine/eval_iter.c b/src/query/engine/eval_iter.c
index 052955dda1..7e2be7789c 100644
--- a/src/query/engine/eval_iter.c
+++ b/src/query/engine/eval_iter.c
@@ -379,7 +379,9 @@ ecs_iter_t ecs_query_iter(
     if (cache) {
         /* If monitors changed, do query rematching */
         ecs_flags32_t flags = q->flags;
-        if (!(world->flags & EcsWorldReadonly) && flags & EcsQueryHasRefs) {
+        if (!(ecs_world_get_flags(world) & EcsWorldReadonly) && 
+             (flags & EcsQueryHasRefs)) 
+        {
             flecs_eval_component_monitors(q->world);
         }
     }
diff --git a/src/world.c b/src/world.c
index cbe66941cc..ab2008ea5d 100644
--- a/src/world.c
+++ b/src/world.c
@@ -2400,6 +2400,11 @@ ecs_entities_t ecs_get_entities(
 ecs_flags32_t ecs_world_get_flags(
     const ecs_world_t *world)
 {
-    flecs_poly_assert(world, ecs_world_t);
-    return world->flags;
+    if (flecs_poly_is(world, ecs_world_t)) {
+        return world->flags;
+    } else {
+        flecs_poly_assert(world, ecs_stage_t);
+        const ecs_stage_t *stage = (const ecs_stage_t*)world;
+        return stage->world->flags;
+    }
 }