refactor: switch SimpleMapStore from legacy to current store implementations

SimpleMapStore now wraps SimpleMapStoreGlobal/SimpleMapStorePerKey (the
epoch-based implementations) instead of the legacy stores. Adds
diagnostic_info() to both current stores. Legacy stores remain available
for benchmark comparisons only.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: zzylol

asap-query-engine/src/stores/simple_map_store/global.rs

@@ -154,6 +154,49 @@ impl SimpleMapStoreGlobal {
             cleanup_policy,
         }
     }
+
+    /// Collect diagnostic info about store contents.
+    pub fn diagnostic_info(&self) -> super::StoreDiagnostics {
+        use super::{AggregationDiagnostic, StoreDiagnostics};
+
+        let data = self.lock.lock().unwrap();
+        let mut per_aggregation = Vec::new();
+        let mut total_time_map_entries: usize = 0;
+        let total_sketch_bytes: usize = 0;
+
+        for (&agg_id, per_key) in &data.stores {
+            let time_map_len = per_key.current_epoch.window_count()
+                + per_key
+                    .sealed_epochs
+                    .values()
+                    .map(|e| e.distinct_window_count())
+                    .sum::<usize>();
+            let read_counts_len = data
+                .read_counts
+                .get(&agg_id)
+                .map(|rc| rc.len())
+                .unwrap_or(0);
+            total_time_map_entries += time_map_len;
+
+            let num_aggregate_objects = per_key.current_epoch.len()
+                + per_key.sealed_epochs.values().map(|e| e.entries.len()).sum::<usize>();
+
+            per_aggregation.push(AggregationDiagnostic {
+                aggregation_id: agg_id,
+                time_map_len,
+                read_counts_len,
+                num_aggregate_objects,
+                sketch_bytes: 0, // skip serialization for diagnostics
+            });
+        }
+
+        StoreDiagnostics {
+            num_aggregations: data.stores.len(),
+            total_time_map_entries,
+            total_sketch_bytes,
+            per_aggregation,
+        }
+    }
 }
 
 /// Extracted config fields needed inside the locked batch loop.

asap-query-engine/src/stores/simple_map_store/mod.rs

@@ -7,17 +7,32 @@ use crate::data_model::{
     AggregateCore, CleanupPolicy, LockStrategy, PrecomputedOutput, StreamingConfig,
 };
 use crate::stores::{Store, StoreResult, TimestampedBucketsMap};
+use global::SimpleMapStoreGlobal;
+use per_key::SimpleMapStorePerKey;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-pub use legacy::LegacySimpleMapStoreGlobal;
-pub use legacy::LegacySimpleMapStorePerKey;
-pub use legacy::{AggregationDiagnostic, StoreDiagnostics};
+/// Diagnostic snapshot from a single aggregation ID in the store.
+pub struct AggregationDiagnostic {
+    pub aggregation_id: u64,
+    pub time_map_len: usize,
+    pub read_counts_len: usize,
+    pub num_aggregate_objects: usize,
+    pub sketch_bytes: usize,
+}
+
+/// Diagnostic snapshot of the entire store.
+pub struct StoreDiagnostics {
+    pub num_aggregations: usize,
+    pub total_time_map_entries: usize,
+    pub total_sketch_bytes: usize,
+    pub per_aggregation: Vec<AggregationDiagnostic>,
+}
 
 /// Enum wrapper that dispatches to either global or per-key lock implementation
 pub enum SimpleMapStore {
-    Global(LegacySimpleMapStoreGlobal),
-    PerKey(LegacySimpleMapStorePerKey),
+    Global(SimpleMapStoreGlobal),
+    PerKey(SimpleMapStorePerKey),
 }
 
 impl SimpleMapStore {
@@ -26,7 +41,7 @@ impl SimpleMapStore {
         Self::new_with_strategy(streaming_config, cleanup_policy, LockStrategy::PerKey)
     }
 
-    /// Collect diagnostic info for memory leak investigation.
+    /// Collect diagnostic info for memory investigation.
     pub fn diagnostic_info(&self) -> StoreDiagnostics {
         match self {
             SimpleMapStore::Global(store) => store.diagnostic_info(),
@@ -41,11 +56,11 @@ impl SimpleMapStore {
         lock_strategy: LockStrategy,
     ) -> Self {
         match lock_strategy {
-            LockStrategy::Global => SimpleMapStore::Global(LegacySimpleMapStoreGlobal::new(
+            LockStrategy::Global => SimpleMapStore::Global(SimpleMapStoreGlobal::new(
                 streaming_config,
                 cleanup_policy,
             )),
-            LockStrategy::PerKey => SimpleMapStore::PerKey(LegacySimpleMapStorePerKey::new(
+            LockStrategy::PerKey => SimpleMapStore::PerKey(SimpleMapStorePerKey::new(
                 streaming_config,
                 cleanup_policy,
             )),

asap-query-engine/src/stores/simple_map_store/per_key.rs

@@ -183,6 +183,49 @@ impl SimpleMapStorePerKey {
         }
     }
 
+    /// Collect diagnostic info about store contents.
+    pub fn diagnostic_info(&self) -> super::StoreDiagnostics {
+        use super::{AggregationDiagnostic, StoreDiagnostics};
+
+        let mut per_aggregation = Vec::new();
+        let mut total_time_map_entries: usize = 0;
+        let total_sketch_bytes: usize = 0;
+
+        for entry in self.store.iter() {
+            let agg_id = *entry.key();
+            let data = match entry.value().read() {
+                Ok(d) => d,
+                Err(_) => continue,
+            };
+            let time_map_len = data.current_epoch.window_count()
+                + data
+                    .sealed_epochs
+                    .values()
+                    .map(|e| e.distinct_window_count())
+                    .sum::<usize>();
+            let read_counts_len = data.read_counts.lock().map(|rc| rc.len()).unwrap_or(0);
+            total_time_map_entries += time_map_len;
+
+            let num_aggregate_objects = data.current_epoch.len()
+                + data.sealed_epochs.values().map(|e| e.entries.len()).sum::<usize>();
+
+            per_aggregation.push(AggregationDiagnostic {
+                aggregation_id: agg_id,
+                time_map_len,
+                read_counts_len,
+                num_aggregate_objects,
+                sketch_bytes: 0, // skip serialization for diagnostics
+            });
+        }
+
+        StoreDiagnostics {
+            num_aggregations: self.store.len(),
+            total_time_map_entries,
+            total_sketch_bytes,
+            per_aggregation,
+        }
+    }
+
     fn cleanup_old_aggregates(
         &self,
         data: &mut StoreKeyData,