From 1d98af52fb243673e7915d608b37a2740fe3b8a4 Mon Sep 17 00:00:00 2001
From: Robert Jacobson <RJacobson@cdc.gov>
Date: Mon, 13 Jul 2026 16:47:55 -0500
Subject: [PATCH] fix: Eagerly compute property value count indexes

---
 src/entity/context_extension.rs       |  7 ++++
 src/entity/index/value_count_index.rs | 55 +++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/src/entity/context_extension.rs b/src/entity/context_extension.rs
index 50d2d58b..495187fd 100644
--- a/src/entity/context_extension.rs
+++ b/src/entity/context_extension.rs
@@ -376,10 +376,17 @@ impl ContextEntitiesExt for Context {
     }
 
     fn index_property_counts<E: Entity, P: IndexableProperty<E>>(&mut self) {
+        let property_id = P::id();
+        let context_ptr: *const Context = self;
         let property_store = self.entity_store.get_property_store_mut::<E>();
         let current_index_type = property_store.get::<P>().index_type();
         if current_index_type != PropertyIndexType::FullIndex {
             property_store.set_property_indexed::<P>(PropertyIndexType::ValueCountIndex);
+            // SAFETY: This only creates a shared reference to `Context` while mutably borrowing
+            // the property store to update index internals.
+            unsafe {
+                property_store.index_unindexed_entities_for_property_id(&*context_ptr, property_id);
+            }
         }
     }
 
diff --git a/src/entity/index/value_count_index.rs b/src/entity/index/value_count_index.rs
index 0cc80318..44a1e6d0 100644
--- a/src/entity/index/value_count_index.rs
+++ b/src/entity/index/value_count_index.rs
@@ -142,6 +142,61 @@ mod tests {
         );
     }
 
+    #[test]
+    fn enabling_multi_property_value_count_index_indexes_existing_entities() {
+        let mut context = Context::new();
+
+        context
+            .add_entity(with!(Person, Age(1u8), Weight(2u8), Height(3u8)))
+            .unwrap();
+        context
+            .add_entity(with!(Person, Age(1u8), Weight(2u8), Height(3u8)))
+            .unwrap();
+        context
+            .add_entity(with!(Person, Age(3u8), Weight(2u8), Height(1u8)))
+            .unwrap();
+
+        context.index_property_counts::<Person, AWH>();
+
+        assert_eq!(
+            context.query_entity_count(with!(Person, Age(1u8), Weight(2u8), Height(3u8))),
+            2
+        );
+        assert_eq!(
+            context.query_entity_count(with!(Person, Age(3u8), Weight(2u8), Height(1u8))),
+            1
+        );
+    }
+
+    #[test]
+    fn enabling_value_count_index_indexes_existing_entities() {
+        let mut context = Context::new();
+
+        context.add_entity(with!(Person, Age(10))).unwrap();
+        context.add_entity(with!(Person, Age(10))).unwrap();
+
+        context.index_property_counts::<Person, Age>();
+
+        assert_eq!(context.query_entity_count(with!(Person, Age(10))), 2);
+    }
+
+    #[test]
+    fn changing_existing_entity_before_adding_another_does_not_double_count() {
+        let mut context = Context::new();
+
+        let first = context.add_entity(with!(Person, Age(10))).unwrap();
+        context.add_entity(with!(Person, Age(10))).unwrap();
+
+        context.index_property_counts::<Person, Age>();
+        context.set_property(first, Age(20));
+
+        // Adding another entity must not count the changed entity a second time.
+        context.add_entity(with!(Person, Age(10))).unwrap();
+
+        assert_eq!(context.query_entity_count(with!(Person, Age(20))), 1);
+        assert_eq!(context.query_entity_count(with!(Person, Age(10))), 2);
+    }
+
     #[test]
     fn test_index_value_compute_same_values() {
         let value = one_shot_128(&"test value");