From b65a6f8e046bac21ed600be28903658fd5f4d257 Mon Sep 17 00:00:00 2001 From: k88hudson-cfa Date: Thu, 28 Nov 2024 15:50:01 -0800 Subject: [PATCH] Periodic reports --- src/people.rs | 140 +++++++++++++++++++++++++++++++++++++++++--------- src/report.rs | 130 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 246 insertions(+), 24 deletions(-) diff --git a/src/people.rs b/src/people.rs index 70fecd4d..0664de65 100644 --- a/src/people.rs +++ b/src/people.rs @@ -101,10 +101,25 @@ seq!(Z in 1..20 { }); pub trait Tabulator { + fn setup(&self, context: &mut Context); fn get_typelist(&self) -> Vec; fn get_columns(&self) -> Vec; } +impl Tabulator for (T,) { + fn setup(&self, context: &mut Context) { + context.index_property_noargs::(); + } + fn get_typelist(&self) -> Vec { + vec![std::any::TypeId::of::()] + } + fn get_columns(&self) -> Vec { + vec![String::from( + std::any::type_name::().split("::").last().unwrap(), + )] + } +} + macro_rules! impl_tabulator { ($ct:expr) => { seq!(N in 0..$ct { @@ -118,6 +133,11 @@ macro_rules! impl_tabulator { )* ) { + fn setup(&self, context: &mut Context) { + #( + context.index_property_noargs::(); + )* + } fn get_typelist(&self) -> Vec { vec![ #( @@ -138,7 +158,7 @@ macro_rules! impl_tabulator { } } -seq!(Z in 1..20 { +seq!(Z in 2..20 { impl_tabulator!(Z); }); @@ -194,6 +214,7 @@ impl Index { max_indexed: 0, } } + fn add_person(&mut self, context: &Context, person_id: PersonId) { let (hash, display) = (self.indexer)(context, person_id); self.lookup @@ -617,14 +638,20 @@ impl PeopleData { } } - fn get_index_ref_mut_by_prop( + fn get_index_ref_mut_by_prop_noargs( &self, - _property: T, ) -> Option> { let type_id = TypeId::of::(); self.get_index_ref_mut(type_id) } + fn get_index_ref_mut_by_prop( + &self, + _property: T, + ) -> Option> { + self.get_index_ref_mut_by_prop_noargs::() + } + // Convenience function to iterate over the current population. // Note that this doesn't hold a reference to PeopleData, so if // you change the population while using it, it won't notice. @@ -730,28 +757,37 @@ pub trait ContextPeopleExt { // Returns a PersonId for a usize fn get_person_id(&self, person_id: usize) -> PersonId; + fn index_property_noargs(&mut self); fn index_property(&mut self, property: T); fn query_people(&self, q: T) -> Vec; fn match_person(&self, person_id: PersonId, q: T) -> bool; - fn get_counts(&self, tabulator: T, print_fn: &dyn Fn(&[String], usize)); + fn get_counts(&self, tabulator: &T, print_fn: F) + where + F: Fn(&Context, &[String], usize); } fn process_indices( + context: &Context, remaining_indices: &[&Index], accumulated_values: &mut [IndexValue], display_values: &mut Vec, current_matches: &mut HashSet, intersect: bool, - print_fn: &dyn Fn(&[String], usize), + print_fn: &dyn Fn(&Context, &[String], usize), ) { if remaining_indices.is_empty() { - print_fn(display_values, current_matches.len()); + print_fn(context, display_values, current_matches.len()); return; } if let Some((next_index, rest_indices)) = remaining_indices.split_first() { - // TODO this might be empty? let lookup = next_index.lookup.as_ref().unwrap(); + + // If there is nothing in the index, we don't need to process it + if lookup.is_empty() { + return; + } + for (value, (people, display)) in lookup { let mut updated_values = accumulated_values.to_owned(); updated_values.push(value.clone()); @@ -768,6 +804,7 @@ fn process_indices( }; process_indices( + context, rest_indices, &mut updated_values, display_values, @@ -775,6 +812,8 @@ fn process_indices( true, print_fn, ); + display_values.pop(); + updated_values.pop(); } } } @@ -923,7 +962,11 @@ impl ContextPeopleExt for Context { PersonId { id: person_id } } - fn index_property(&mut self, property: T) { + fn index_property(&mut self, _property: T) { + self.index_property_noargs::(); + } + + fn index_property_noargs(&mut self) { // Ensure that the data container exists { let _ = self.get_data_container_mut(PeoplePlugin); @@ -933,7 +976,9 @@ impl ContextPeopleExt for Context { self.register_indexer::(); let data_container = self.get_data_container(PeoplePlugin).unwrap(); - let mut index = data_container.get_index_ref_mut_by_prop(property).unwrap(); + let mut index = data_container + .get_index_ref_mut_by_prop_noargs::() + .unwrap(); if index.lookup.is_none() { index.lookup = Some(HashMap::new()); } @@ -986,7 +1031,10 @@ impl ContextPeopleExt for Context { } } - fn get_counts(&self, tabulator: T, print_fn: &dyn Fn(&[String], usize)) { + fn get_counts(&self, tabulator: &T, print_fn: F) + where + F: Fn(&Context, &[String], usize), + { let type_ids = tabulator.get_typelist(); // First, update indexes @@ -1016,6 +1064,7 @@ impl ContextPeopleExt for Context { .collect::>(); process_indices( + self, indices.as_slice(), &mut Vec::new(), &mut Vec::new(), @@ -1171,7 +1220,7 @@ mod test { error::IxaError, people::{Index, IndexValue, PeoplePlugin, PersonPropertyHolder}, }; - use std::{any::TypeId, cell::RefCell, rc::Rc, vec}; + use std::{any::TypeId, cell::RefCell, collections::HashSet, hash::Hash, rc::Rc, vec}; define_person_property!(Age, u8); #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] @@ -1855,21 +1904,66 @@ mod test { assert_eq!(cols.get_columns(), vec!["Age", "RiskCategoryType"]); } - #[test] - fn test_get_counts() { + fn get_counts_test_setup( + tabulator: &T, + setup: impl FnOnce(&mut Context), + expected_values: &HashSet<(Vec, usize)>, + ) { let mut context = Context::new(); - context.index_property(IsRunner); - context.index_property(IsSwimmer); + setup(&mut context); + tabulator.setup(&mut context); + + let results: RefCell, usize)>> = RefCell::new(HashSet::new()); + context.get_counts(tabulator, |_context, values, count| { + results.borrow_mut().insert((values.to_vec(), count)); + }); - let anne = context.add_person(()).unwrap(); - let charlie = context.add_person(()).unwrap(); + let results = &*results.borrow(); + assert_eq!(results, expected_values); + } - context.set_person_property(anne, IsRunner, true); - context.set_person_property(charlie, IsRunner, true); - context.set_person_property(anne, IsSwimmer, true); + #[test] + fn test_get_counts() { + let tabulator = (IsRunner,); + let mut expected = HashSet::new(); + expected.insert((vec!["true".to_string()], 1)); + expected.insert((vec!["false".to_string()], 1)); + get_counts_test_setup( + &tabulator, + |context| { + let bob = context.add_person(()).unwrap(); + context.add_person(()).unwrap(); + context.set_person_property(bob, IsRunner, true); + }, + &expected, + ); + } - context.get_counts((IsRunner, IsSwimmer), &|values, count| { - println!("{values:?} {count}"); - }); + #[test] + fn test_get_counts_multi() { + let tabulator = (IsRunner, IsSwimmer); + let mut expected = HashSet::new(); + expected.insert((vec!["false".to_string(), "false".to_string()], 3)); + expected.insert((vec!["false".to_string(), "true".to_string()], 1)); + expected.insert((vec!["true".to_string(), "false".to_string()], 1)); + expected.insert((vec!["true".to_string(), "true".to_string()], 1)); + + get_counts_test_setup( + &tabulator, + |context| { + context.add_person(()).unwrap(); + context.add_person(()).unwrap(); + context.add_person(()).unwrap(); + let bob = context.add_person(()).unwrap(); + let anne = context.add_person(()).unwrap(); + let charlie = context.add_person(()).unwrap(); + + context.set_person_property(bob, IsRunner, true); + context.set_person_property(charlie, IsRunner, true); + context.set_person_property(anne, IsSwimmer, true); + context.set_person_property(charlie, IsSwimmer, true); + }, + &expected, + ); } } diff --git a/src/report.rs b/src/report.rs index 350f15cc..759843ed 100644 --- a/src/report.rs +++ b/src/report.rs @@ -1,5 +1,6 @@ use crate::context::Context; use crate::error::IxaError; +use crate::people::{ContextPeopleExt, Tabulator}; use csv::Writer; use std::any::TypeId; use std::cell::RefCell; @@ -60,6 +61,27 @@ pub trait Report: 'static { fn serialize(&self, writer: &mut Writer); } +#[allow(clippy::module_name_repetitions)] +pub struct PeriodicReport { + t: f64, + values: Vec, + count: usize, +} + +impl Report for PeriodicReport { + fn type_id(&self) -> TypeId { + TypeId::of::() + } + + fn serialize(&self, writer: &mut Writer) { + let mut row = vec![self.t.to_string()]; + row.extend(self.values.clone()); + row.push(self.count.to_string()); + + writer.write_record(&row).expect("Failed to write row"); + } +} + /// Use this macro to define a unique report type #[macro_export] macro_rules! create_report_trait { @@ -115,6 +137,17 @@ pub trait ContextReportExt { /// If the file already exists and `overwrite` is set to false, raises an error and info message. /// If the file cannot be created, raises an error. fn add_report(&mut self, short_name: &str) -> Result<(), IxaError>; + /// Adds a periodic report ad the end of period `period` which summarizes the + /// number of people in each combination of properties in `tabulator`. + /// # Errors + /// If the file already exists and `overwrite` is set to false, raises an error and info message. + /// If the file cannot be created, raises an error. + fn add_periodic_report( + &mut self, + short_name: &str, + period: f64, + tabulator: T, + ) -> Result<(), IxaError>; fn send_report(&self, report: T); fn report_options(&mut self) -> &mut ConfigReportOptions; } @@ -151,6 +184,56 @@ impl ContextReportExt for Context { Ok(()) } + fn add_periodic_report( + &mut self, + short_name: &str, + period: f64, + tabulator: T, + ) -> Result<(), IxaError> { + self.add_report::(short_name)?; + + { + // Write the header + let columns = tabulator.get_columns(); + let mut header = vec!["t".to_string()]; + header.extend(columns); + header.push("count".to_string()); + + let data_container = self + .get_data_container(ReportPlugin) + .expect("No writer found for the report type"); + let mut writer_cell = data_container.file_writers.try_borrow_mut().unwrap(); + let writer = writer_cell + .get_mut(&TypeId::of::()) + .expect("No writer found for the report type"); + writer + .write_record(&header) + .expect("Failed to write header"); + writer.flush().expect("Failed to flush writer"); + } + + // TODO add indexes + tabulator.setup(self); + + self.add_periodic_plan_with_phase( + period, + move |context: &mut Context| { + context.get_counts(&tabulator, move |context, values, count| { + println!("{} {}", values.join(","), count); + context.send_report(PeriodicReport { + t: context.get_current_time(), + values: values.to_vec(), + count, + }); + }); + // Hello world + }, + crate::context::ExecutionPhase::Last, + ); + + Ok(()) + } + /// Write a new row to the appropriate report file fn send_report(&self, report: T) { // No data container will exist if no reports have been added @@ -174,12 +257,16 @@ impl ContextReportExt for Context { #[cfg(test)] mod test { + use crate::{define_person_property, define_person_property_with_default}; + use super::*; use core::convert::TryInto; use serde_derive::{Deserialize, Serialize}; use std::thread; use tempfile::tempdir; + define_person_property_with_default!(IsRunner, bool, false); + #[derive(Serialize, Deserialize)] struct SampleReport { id: u32, @@ -451,10 +538,51 @@ mod test { .overwrite(true); let result = context2.add_report::("sample_report"); assert!(result.is_ok()); - // file should also be empty let file = File::open(file_path).unwrap(); let reader = csv::Reader::from_reader(file); let records = reader.into_records(); assert_eq!(records.count(), 0); } + + #[test] + fn add_periodic_report() { + let mut context = Context::new(); + let temp_dir = tempdir().unwrap(); + let path = PathBuf::from(&temp_dir.path()); + let config = context.report_options(); + config + .file_prefix("test_".to_string()) + .directory(path.clone()); + let _ = context.add_periodic_report("periodic", 1.0, (IsRunner,)); + let person = context.add_person(()).unwrap(); + context.add_person(()).unwrap(); + + context.add_plan(1.0, move |context: &mut Context| { + context.set_person_property(person, IsRunner, true); + }); + + context.execute(); + + let file_path = path.join("test_periodic.csv"); + assert!(file_path.exists(), "CSV file should exist"); + + let mut reader = csv::Reader::from_path(file_path).unwrap(); + + assert_eq!(reader.headers().unwrap(), vec!["t", "IsRunner", "count"]); + + let mut actual: Vec> = reader + .records() + .map(|result| result.unwrap().iter().map(String::from).collect()) + .collect(); + let mut expected = vec![ + vec!["0", "false", "2"], + vec!["1", "false", "1"], + vec!["1", "true", "1"], + ]; + + actual.sort(); + expected.sort(); + + assert_eq!(actual, expected, "CSV file should contain the correct data"); + } }