From a35f41d345f83834e76af71fb7c24f13f5e6863a Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Fri, 19 Jul 2024 15:22:38 -0400
Subject: [PATCH 01/17] Import context from eosim

---
 Cargo.toml     |   1 +
 src/context.rs | 354 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs     |   1 +
 src/main.rs    |   3 -
 4 files changed, 356 insertions(+), 3 deletions(-)
 create mode 100644 src/context.rs
 create mode 100644 src/lib.rs
 delete mode 100644 src/main.rs

diff --git a/Cargo.toml b/Cargo.toml
index be2f210..6e0c8e1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,3 +8,4 @@ license = "Apache-2.0"
 homepage = "https://github.com/CDCgov/ixa"
 
 [dependencies]
+derivative = "2.2.0"
diff --git a/src/context.rs b/src/context.rs
new file mode 100644
index 0000000..4e8abe0
--- /dev/null
+++ b/src/context.rs
@@ -0,0 +1,354 @@
+use std::{
+    any::{Any, TypeId},
+    cmp::Ordering,
+    collections::{BinaryHeap, HashMap, HashSet, VecDeque},
+    rc::Rc,
+};
+
+use derivative::Derivative;
+
+pub trait Component: Any {
+    fn init(context: &mut Context);
+}
+
+pub trait Plugin: Any {
+    type DataContainer;
+
+    fn get_data_container() -> Self::DataContainer;
+}
+
+#[macro_export]
+macro_rules! define_plugin {
+    ($plugin:ident, $data_container:ty, $default: expr) => {
+        struct $plugin {}
+
+        impl $crate::context::Plugin for $plugin {
+            type DataContainer = $data_container;
+
+            fn get_data_container() -> Self::DataContainer {
+                $default
+            }
+        }
+    };
+}
+pub use define_plugin;
+
+pub struct PlanId {
+    pub id: u64,
+}
+
+#[derive(Derivative)]
+#[derivative(Eq, PartialEq, Debug)]
+pub struct TimedPlan {
+    pub time: f64,
+    plan_id: u64,
+    #[derivative(PartialEq = "ignore", Debug = "ignore")]
+    pub callback: Box<dyn FnOnce(&mut Context)>,
+}
+
+impl Ord for TimedPlan {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
+        if time_ordering == Ordering::Equal {
+            // Break time ties in order of plan id
+            self.plan_id.cmp(&other.plan_id).reverse()
+        } else {
+            time_ordering
+        }
+    }
+}
+
+impl PartialOrd for TimedPlan {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[derive(Debug)]
+struct PlanQueue {
+    queue: BinaryHeap<TimedPlan>,
+    invalid_set: HashSet<u64>,
+    plan_counter: u64,
+}
+
+impl PlanQueue {
+    pub fn new() -> PlanQueue {
+        PlanQueue {
+            queue: BinaryHeap::new(),
+            invalid_set: HashSet::new(),
+            plan_counter: 0,
+        }
+    }
+
+    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
+        // Add plan to queue and increment counter
+        let plan_id = self.plan_counter;
+        self.queue.push(TimedPlan {
+            time,
+            plan_id,
+            callback: Box::new(callback),
+        });
+        self.plan_counter += 1;
+        PlanId { id: plan_id }
+    }
+
+    pub fn cancel_plan(&mut self, id: PlanId) {
+        self.invalid_set.insert(id.id);
+    }
+
+    pub fn get_next_timed_plan(&mut self) -> Option<TimedPlan> {
+        loop {
+            let next_timed_plan = self.queue.pop();
+            match next_timed_plan {
+                Some(timed_plan) => {
+                    if self.invalid_set.contains(&timed_plan.plan_id) {
+                        self.invalid_set.remove(&timed_plan.plan_id);
+                    } else {
+                        return Some(timed_plan);
+                    }
+                }
+                None => {
+                    return None;
+                }
+            }
+        }
+    }
+}
+
+type Callback = dyn FnOnce(&mut Context);
+type EventHandler<E> = dyn Fn(&mut Context, E);
+pub struct Context {
+    plan_queue: PlanQueue,
+    callback_queue: VecDeque<Box<Callback>>,
+    event_handlers: HashMap<TypeId, Box<dyn Any>>,
+    immediate_event_handlers: HashMap<TypeId, Box<dyn Any>>,
+    plugin_data: HashMap<TypeId, Box<dyn Any>>,
+    time: f64,
+}
+
+impl Context {
+    pub fn new() -> Context {
+        Context {
+            plan_queue: PlanQueue::new(),
+            callback_queue: VecDeque::new(),
+            event_handlers: HashMap::new(),
+            immediate_event_handlers: HashMap::new(),
+            plugin_data: HashMap::new(),
+            time: 0.0,
+        }
+    }
+
+    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
+        // TODO: Handle invalid times (past, NAN, etc)
+        self.plan_queue.add_plan(time, callback)
+    }
+
+    pub fn cancel_plan(&mut self, id: PlanId) {
+        self.plan_queue.cancel_plan(id);
+    }
+
+    pub fn queue_callback(&mut self, callback: impl FnOnce(&mut Context) + 'static) {
+        self.callback_queue.push_back(Box::new(callback));
+    }
+
+    fn add_plugin<T: Plugin>(&mut self) {
+        self.plugin_data
+            .insert(TypeId::of::<T>(), Box::new(T::get_data_container()));
+    }
+
+    pub fn get_data_container_mut<T: Plugin>(&mut self) -> &mut T::DataContainer {
+        let type_id = &TypeId::of::<T>();
+        if !self.plugin_data.contains_key(type_id) {
+            self.add_plugin::<T>();
+        }
+        let data_container = self
+            .plugin_data
+            .get_mut(type_id)
+            .unwrap()
+            .downcast_mut::<T::DataContainer>();
+        match data_container {
+            Some(x) => x,
+            None => panic!("Plugin data container of incorrect type"),
+        }
+    }
+
+    pub fn get_data_container<T: Plugin>(&self) -> Option<&T::DataContainer> {
+        let type_id = &TypeId::of::<T>();
+        if !self.plugin_data.contains_key(type_id) {
+            return None;
+        }
+        let data_container = self
+            .plugin_data
+            .get(type_id)
+            .unwrap()
+            .downcast_ref::<T::DataContainer>();
+        match data_container {
+            Some(x) => Some(x),
+            None => panic!("Plugin data container of incorrect type"),
+        }
+    }
+
+    pub fn get_time(&self) -> f64 {
+        self.time
+    }
+
+    pub fn add_component<T: Component>(&mut self) {
+        T::init(self);
+    }
+
+    fn add_handlers<E: Copy + 'static>(
+        event_handlers: &mut HashMap<TypeId, Box<dyn Any>>,
+        callback: impl Fn(&mut Context, E) + 'static,
+    ) {
+        let callback_vec = event_handlers
+            .entry(TypeId::of::<E>())
+            .or_insert_with(|| Box::<Vec<Rc<EventHandler<E>>>>::default());
+        let callback_vec: &mut Vec<Rc<EventHandler<E>>> = callback_vec.downcast_mut().unwrap();
+        callback_vec.push(Rc::new(callback));
+    }
+
+    pub fn subscribe_to_event<E: Copy + 'static>(
+        &mut self,
+        callback: impl Fn(&mut Context, E) + 'static,
+    ) {
+        Self::add_handlers(&mut self.event_handlers, callback);
+    }
+
+    pub fn subscribe_immediately_to_event<E: Copy + 'static>(
+        &mut self,
+        callback: impl Fn(&mut Context, E) + 'static,
+    ) {
+        Self::add_handlers(&mut self.immediate_event_handlers, callback);
+    }
+
+    fn collect_callbacks<E: Copy + 'static>(
+        event_handlers: &HashMap<TypeId, Box<dyn Any>>,
+        event: E,
+    ) -> Vec<Box<Callback>> {
+        let mut callbacks_to_return = Vec::<Box<Callback>>::new();
+        let callback_vec = event_handlers.get(&TypeId::of::<E>());
+        if let Some(callback_vec) = callback_vec {
+            let callback_vec: &Vec<Rc<EventHandler<E>>> = callback_vec.downcast_ref().unwrap();
+            if !callback_vec.is_empty() {
+                for callback in callback_vec {
+                    let internal_callback = Rc::clone(callback);
+                    callbacks_to_return
+                        .push(Box::new(move |context| internal_callback(context, event)));
+                }
+            }
+        }
+        callbacks_to_return
+    }
+
+    pub fn release_event<E: Copy + 'static>(&mut self, event: E) {
+        // Queue standard handlers
+        for callback in Self::collect_callbacks(&self.event_handlers, event) {
+            self.queue_callback(callback);
+        }
+        // Process immediate handlers
+        for callback in Self::collect_callbacks(&self.immediate_event_handlers, event) {
+            callback(self);
+        }
+    }
+
+    pub fn execute(&mut self) {
+        // Execute callbacks if there are any in the queue
+        loop {
+            let callback = self.callback_queue.pop_front();
+            match callback {
+                Some(callback) => callback(self),
+                None => break,
+            }
+        }
+        // Start plan loop
+        loop {
+            let timed_plan = self.plan_queue.get_next_timed_plan();
+            match timed_plan {
+                Some(timed_plan) => {
+                    self.time = timed_plan.time;
+                    (timed_plan.callback)(self);
+                    loop {
+                        let callback = self.callback_queue.pop_front();
+                        match callback {
+                            Some(callback) => callback(self),
+                            None => break,
+                        }
+                    }
+                }
+                None => break,
+            }
+        }
+    }
+}
+
+impl Default for Context {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{cell::RefCell, rc::Rc};
+
+    use super::*;
+
+    define_plugin!(ComponentA, u32, 0);
+
+    impl ComponentA {
+        fn increment_counter(context: &mut Context) {
+            *(context.get_data_container_mut::<ComponentA>()) += 1;
+        }
+    }
+
+    impl Component for ComponentA {
+        fn init(context: &mut Context) {
+            context.add_plan(1.0, Self::increment_counter);
+        }
+    }
+
+    #[test]
+    fn test_component_and_planning() {
+        let mut context = Context::new();
+        context.add_component::<ComponentA>();
+        assert_eq!(context.get_time(), 0.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 0);
+        context.execute();
+        assert_eq!(context.get_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 1);
+        let plan_to_cancel = context.add_plan(3.0, ComponentA::increment_counter);
+        context.add_plan(2.0, ComponentA::increment_counter);
+        context.cancel_plan(plan_to_cancel);
+        context.execute();
+        assert_eq!(context.get_time(), 2.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 2);
+    }
+
+    #[derive(Copy, Clone)]
+    struct Event {
+        pub data: usize,
+    }
+    #[test]
+    fn test_events() {
+        let mut context = Context::new();
+
+        let obs_data = Rc::new(RefCell::new(0));
+        let immediate_obs_data = Rc::new(RefCell::new(0));
+
+        let obs_data_clone = Rc::clone(&obs_data);
+        context.subscribe_to_event::<Event>(move |_, event| {
+            *obs_data_clone.borrow_mut() = event.data;
+        });
+
+        let immediate_obs_data_clone = Rc::clone(&immediate_obs_data);
+        context.subscribe_immediately_to_event::<Event>(move |_, event| {
+            *immediate_obs_data_clone.borrow_mut() = event.data;
+        });
+
+        context.release_event(Event { data: 1 });
+        assert_eq!(*immediate_obs_data.borrow(), 1);
+
+        context.execute();
+        assert_eq!(*obs_data.borrow(), 1);
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..9efb2ab
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1 @@
+pub mod context;
diff --git a/src/main.rs b/src/main.rs
deleted file mode 100644
index e7a11a9..0000000
--- a/src/main.rs
+++ /dev/null
@@ -1,3 +0,0 @@
-fn main() {
-    println!("Hello, world!");
-}

From 19f6e6f13d3cbf6a53a085b2f51112c6e79121e3 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Sun, 28 Jul 2024 13:36:57 -0400
Subject: [PATCH 02/17] Remove event emission and handle comments from review

---
 src/context.rs | 138 +++++--------------------------------------------
 1 file changed, 14 insertions(+), 124 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 4e8abe0..f283874 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -2,15 +2,10 @@ use std::{
     any::{Any, TypeId},
     cmp::Ordering,
     collections::{BinaryHeap, HashMap, HashSet, VecDeque},
-    rc::Rc,
 };
 
 use derivative::Derivative;
 
-pub trait Component: Any {
-    fn init(context: &mut Context);
-}
-
 pub trait Plugin: Any {
     type DataContainer;
 
@@ -116,14 +111,11 @@ impl PlanQueue {
 }
 
 type Callback = dyn FnOnce(&mut Context);
-type EventHandler<E> = dyn Fn(&mut Context, E);
 pub struct Context {
     plan_queue: PlanQueue,
     callback_queue: VecDeque<Box<Callback>>,
-    event_handlers: HashMap<TypeId, Box<dyn Any>>,
-    immediate_event_handlers: HashMap<TypeId, Box<dyn Any>>,
     plugin_data: HashMap<TypeId, Box<dyn Any>>,
-    time: f64,
+    current_time: f64,
 }
 
 impl Context {
@@ -131,10 +123,8 @@ impl Context {
         Context {
             plan_queue: PlanQueue::new(),
             callback_queue: VecDeque::new(),
-            event_handlers: HashMap::new(),
-            immediate_event_handlers: HashMap::new(),
             plugin_data: HashMap::new(),
-            time: 0.0,
+            current_time: 0.0,
         }
     }
 
@@ -161,15 +151,11 @@ impl Context {
         if !self.plugin_data.contains_key(type_id) {
             self.add_plugin::<T>();
         }
-        let data_container = self
-            .plugin_data
+        self.plugin_data
             .get_mut(type_id)
             .unwrap()
-            .downcast_mut::<T::DataContainer>();
-        match data_container {
-            Some(x) => x,
-            None => panic!("Plugin data container of incorrect type"),
-        }
+            .downcast_mut::<T::DataContainer>()
+            .unwrap()
     }
 
     pub fn get_data_container<T: Plugin>(&self) -> Option<&T::DataContainer> {
@@ -177,78 +163,14 @@ impl Context {
         if !self.plugin_data.contains_key(type_id) {
             return None;
         }
-        let data_container = self
-            .plugin_data
+        self.plugin_data
             .get(type_id)
             .unwrap()
-            .downcast_ref::<T::DataContainer>();
-        match data_container {
-            Some(x) => Some(x),
-            None => panic!("Plugin data container of incorrect type"),
-        }
-    }
-
-    pub fn get_time(&self) -> f64 {
-        self.time
-    }
-
-    pub fn add_component<T: Component>(&mut self) {
-        T::init(self);
-    }
-
-    fn add_handlers<E: Copy + 'static>(
-        event_handlers: &mut HashMap<TypeId, Box<dyn Any>>,
-        callback: impl Fn(&mut Context, E) + 'static,
-    ) {
-        let callback_vec = event_handlers
-            .entry(TypeId::of::<E>())
-            .or_insert_with(|| Box::<Vec<Rc<EventHandler<E>>>>::default());
-        let callback_vec: &mut Vec<Rc<EventHandler<E>>> = callback_vec.downcast_mut().unwrap();
-        callback_vec.push(Rc::new(callback));
-    }
-
-    pub fn subscribe_to_event<E: Copy + 'static>(
-        &mut self,
-        callback: impl Fn(&mut Context, E) + 'static,
-    ) {
-        Self::add_handlers(&mut self.event_handlers, callback);
+            .downcast_ref::<T::DataContainer>()
     }
 
-    pub fn subscribe_immediately_to_event<E: Copy + 'static>(
-        &mut self,
-        callback: impl Fn(&mut Context, E) + 'static,
-    ) {
-        Self::add_handlers(&mut self.immediate_event_handlers, callback);
-    }
-
-    fn collect_callbacks<E: Copy + 'static>(
-        event_handlers: &HashMap<TypeId, Box<dyn Any>>,
-        event: E,
-    ) -> Vec<Box<Callback>> {
-        let mut callbacks_to_return = Vec::<Box<Callback>>::new();
-        let callback_vec = event_handlers.get(&TypeId::of::<E>());
-        if let Some(callback_vec) = callback_vec {
-            let callback_vec: &Vec<Rc<EventHandler<E>>> = callback_vec.downcast_ref().unwrap();
-            if !callback_vec.is_empty() {
-                for callback in callback_vec {
-                    let internal_callback = Rc::clone(callback);
-                    callbacks_to_return
-                        .push(Box::new(move |context| internal_callback(context, event)));
-                }
-            }
-        }
-        callbacks_to_return
-    }
-
-    pub fn release_event<E: Copy + 'static>(&mut self, event: E) {
-        // Queue standard handlers
-        for callback in Self::collect_callbacks(&self.event_handlers, event) {
-            self.queue_callback(callback);
-        }
-        // Process immediate handlers
-        for callback in Self::collect_callbacks(&self.immediate_event_handlers, event) {
-            callback(self);
-        }
+    pub fn get_current_time(&self) -> f64 {
+        self.current_time
     }
 
     pub fn execute(&mut self) {
@@ -265,7 +187,7 @@ impl Context {
             let timed_plan = self.plan_queue.get_next_timed_plan();
             match timed_plan {
                 Some(timed_plan) => {
-                    self.time = timed_plan.time;
+                    self.current_time = timed_plan.time;
                     (timed_plan.callback)(self);
                     loop {
                         let callback = self.callback_queue.pop_front();
@@ -289,8 +211,6 @@ impl Default for Context {
 
 #[cfg(test)]
 mod tests {
-    use std::{cell::RefCell, rc::Rc};
-
     use super::*;
 
     define_plugin!(ComponentA, u32, 0);
@@ -299,9 +219,7 @@ mod tests {
         fn increment_counter(context: &mut Context) {
             *(context.get_data_container_mut::<ComponentA>()) += 1;
         }
-    }
 
-    impl Component for ComponentA {
         fn init(context: &mut Context) {
             context.add_plan(1.0, Self::increment_counter);
         }
@@ -310,45 +228,17 @@ mod tests {
     #[test]
     fn test_component_and_planning() {
         let mut context = Context::new();
-        context.add_component::<ComponentA>();
-        assert_eq!(context.get_time(), 0.0);
+        ComponentA::init(&mut context);
+        assert_eq!(context.get_current_time(), 0.0);
         assert_eq!(*context.get_data_container_mut::<ComponentA>(), 0);
         context.execute();
-        assert_eq!(context.get_time(), 1.0);
+        assert_eq!(context.get_current_time(), 1.0);
         assert_eq!(*context.get_data_container_mut::<ComponentA>(), 1);
         let plan_to_cancel = context.add_plan(3.0, ComponentA::increment_counter);
         context.add_plan(2.0, ComponentA::increment_counter);
         context.cancel_plan(plan_to_cancel);
         context.execute();
-        assert_eq!(context.get_time(), 2.0);
+        assert_eq!(context.get_current_time(), 2.0);
         assert_eq!(*context.get_data_container_mut::<ComponentA>(), 2);
     }
-
-    #[derive(Copy, Clone)]
-    struct Event {
-        pub data: usize,
-    }
-    #[test]
-    fn test_events() {
-        let mut context = Context::new();
-
-        let obs_data = Rc::new(RefCell::new(0));
-        let immediate_obs_data = Rc::new(RefCell::new(0));
-
-        let obs_data_clone = Rc::clone(&obs_data);
-        context.subscribe_to_event::<Event>(move |_, event| {
-            *obs_data_clone.borrow_mut() = event.data;
-        });
-
-        let immediate_obs_data_clone = Rc::clone(&immediate_obs_data);
-        context.subscribe_immediately_to_event::<Event>(move |_, event| {
-            *immediate_obs_data_clone.borrow_mut() = event.data;
-        });
-
-        context.release_event(Event { data: 1 });
-        assert_eq!(*immediate_obs_data.borrow(), 1);
-
-        context.execute();
-        assert_eq!(*obs_data.borrow(), 1);
-    }
 }

From e24d5fcd0b9d8330737133ce59f33c67c33d1a91 Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Tue, 30 Jul 2024 16:38:01 -0700
Subject: [PATCH 03/17] Clean up main loop

---
 src/context.rs | 44 +++++++++++++++++++++++---------------------
 1 file changed, 23 insertions(+), 21 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index f283874..1835dd2 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -174,30 +174,32 @@ impl Context {
     }
 
     pub fn execute(&mut self) {
-        // Execute callbacks if there are any in the queue
-        loop {
-            let callback = self.callback_queue.pop_front();
-            match callback {
-                Some(callback) => callback(self),
-                None => break,
-            }
-        }
         // Start plan loop
         loop {
-            let timed_plan = self.plan_queue.get_next_timed_plan();
-            match timed_plan {
-                Some(timed_plan) => {
-                    self.current_time = timed_plan.time;
-                    (timed_plan.callback)(self);
-                    loop {
-                        let callback = self.callback_queue.pop_front();
-                        match callback {
-                            Some(callback) => callback(self),
-                            None => break,
-                        }
-                    }
+            let mut executed = false;
+
+            // First execute callbacks.
+            loop {
+                let callback = self.callback_queue.pop_front();
+                match callback {
+                    Some(callback) => {
+                        executed = true;
+                        callback(self)
+                    },
+                    None => break,
                 }
-                None => break,
+            }
+
+            // Now execute the first timed plan.
+            if let Some(timed_plan) = self.plan_queue.get_next_timed_plan() {
+                executed = true;
+                self.current_time = timed_plan.time;
+                (timed_plan.callback)(self);
+            }
+
+            // If nothing happened, then we are done.
+            if !executed {
+                break;
             }
         }
     }

From a814d8250b846d2bdc474eef92150e87de426945 Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Tue, 30 Jul 2024 16:55:43 -0700
Subject: [PATCH 04/17] Clean up even more

---
 src/context.rs | 25 +++++++------------------
 1 file changed, 7 insertions(+), 18 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 1835dd2..99b791f 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -176,29 +176,18 @@ impl Context {
     pub fn execute(&mut self) {
         // Start plan loop
         loop {
-            let mut executed = false;
-
-            // First execute callbacks.
-            loop {
-                let callback = self.callback_queue.pop_front();
-                match callback {
-                    Some(callback) => {
-                        executed = true;
-                        callback(self)
-                    },
-                    None => break,
-                }
+            // If there is a callback, run it.
+            if let Some(callback) = self.callback_queue.pop_front() {
+                callback(self);
+                continue;
             }
 
-            // Now execute the first timed plan.
+            // There aren't any callbacks, so look at the first timed plan.
             if let Some(timed_plan) = self.plan_queue.get_next_timed_plan() {
-                executed = true;
                 self.current_time = timed_plan.time;
                 (timed_plan.callback)(self);
-            }
-
-            // If nothing happened, then we are done.
-            if !executed {
+            } else {
+                // OK, there aren't any timed plans, so we're done.
                 break;
             }
         }

From ce2bb7e84f95cc053dca4b7e034c98956e121573 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Tue, 30 Jul 2024 23:28:40 -0400
Subject: [PATCH 05/17] Separate plan code from context

---
 src/context.rs | 87 ++-------------------------------------------
 src/lib.rs     |  1 +
 src/plan.rs    | 96 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 85 deletions(-)
 create mode 100644 src/plan.rs

diff --git a/src/context.rs b/src/context.rs
index 99b791f..dd10654 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,11 +1,8 @@
 use std::{
     any::{Any, TypeId},
-    cmp::Ordering,
-    collections::{BinaryHeap, HashMap, HashSet, VecDeque},
+    collections::{HashMap, VecDeque},
 };
 
-use derivative::Derivative;
-
 pub trait Plugin: Any {
     type DataContainer;
 
@@ -28,87 +25,7 @@ macro_rules! define_plugin {
 }
 pub use define_plugin;
 
-pub struct PlanId {
-    pub id: u64,
-}
-
-#[derive(Derivative)]
-#[derivative(Eq, PartialEq, Debug)]
-pub struct TimedPlan {
-    pub time: f64,
-    plan_id: u64,
-    #[derivative(PartialEq = "ignore", Debug = "ignore")]
-    pub callback: Box<dyn FnOnce(&mut Context)>,
-}
-
-impl Ord for TimedPlan {
-    fn cmp(&self, other: &Self) -> Ordering {
-        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
-        if time_ordering == Ordering::Equal {
-            // Break time ties in order of plan id
-            self.plan_id.cmp(&other.plan_id).reverse()
-        } else {
-            time_ordering
-        }
-    }
-}
-
-impl PartialOrd for TimedPlan {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-#[derive(Debug)]
-struct PlanQueue {
-    queue: BinaryHeap<TimedPlan>,
-    invalid_set: HashSet<u64>,
-    plan_counter: u64,
-}
-
-impl PlanQueue {
-    pub fn new() -> PlanQueue {
-        PlanQueue {
-            queue: BinaryHeap::new(),
-            invalid_set: HashSet::new(),
-            plan_counter: 0,
-        }
-    }
-
-    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
-        // Add plan to queue and increment counter
-        let plan_id = self.plan_counter;
-        self.queue.push(TimedPlan {
-            time,
-            plan_id,
-            callback: Box::new(callback),
-        });
-        self.plan_counter += 1;
-        PlanId { id: plan_id }
-    }
-
-    pub fn cancel_plan(&mut self, id: PlanId) {
-        self.invalid_set.insert(id.id);
-    }
-
-    pub fn get_next_timed_plan(&mut self) -> Option<TimedPlan> {
-        loop {
-            let next_timed_plan = self.queue.pop();
-            match next_timed_plan {
-                Some(timed_plan) => {
-                    if self.invalid_set.contains(&timed_plan.plan_id) {
-                        self.invalid_set.remove(&timed_plan.plan_id);
-                    } else {
-                        return Some(timed_plan);
-                    }
-                }
-                None => {
-                    return None;
-                }
-            }
-        }
-    }
-}
+use crate::plan::{PlanId, PlanQueue};
 
 type Callback = dyn FnOnce(&mut Context);
 pub struct Context {
diff --git a/src/lib.rs b/src/lib.rs
index 9efb2ab..e81eede 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1 +1,2 @@
 pub mod context;
+pub mod plan;
diff --git a/src/plan.rs b/src/plan.rs
new file mode 100644
index 0000000..15c1bd3
--- /dev/null
+++ b/src/plan.rs
@@ -0,0 +1,96 @@
+use std::{
+    cmp::Ordering,
+    collections::{BinaryHeap, HashSet},
+};
+
+use derivative::Derivative;
+
+use crate::context::Context;
+
+pub struct PlanId {
+    id: u64,
+}
+
+#[derive(Derivative)]
+#[derivative(Eq, PartialEq, Debug)]
+pub struct TimedPlan {
+    pub time: f64,
+    plan_id: u64,
+    #[derivative(PartialEq = "ignore", Debug = "ignore")]
+    pub callback: Box<dyn FnOnce(&mut Context)>,
+}
+
+impl Ord for TimedPlan {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
+        if time_ordering == Ordering::Equal {
+            // Break time ties in order of plan id
+            self.plan_id.cmp(&other.plan_id).reverse()
+        } else {
+            time_ordering
+        }
+    }
+}
+
+impl PartialOrd for TimedPlan {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+#[derive(Debug)]
+pub struct PlanQueue {
+    queue: BinaryHeap<TimedPlan>,
+    invalid_set: HashSet<u64>,
+    plan_counter: u64,
+}
+
+impl Default for PlanQueue {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl PlanQueue {
+    pub fn new() -> PlanQueue {
+        PlanQueue {
+            queue: BinaryHeap::new(),
+            invalid_set: HashSet::new(),
+            plan_counter: 0,
+        }
+    }
+
+    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
+        // Add plan to queue and increment counter
+        let plan_id = self.plan_counter;
+        self.queue.push(TimedPlan {
+            time,
+            plan_id,
+            callback: Box::new(callback),
+        });
+        self.plan_counter += 1;
+        PlanId { id: plan_id }
+    }
+
+    pub fn cancel_plan(&mut self, id: PlanId) {
+        self.invalid_set.insert(id.id);
+    }
+
+    pub fn get_next_timed_plan(&mut self) -> Option<TimedPlan> {
+        loop {
+            let next_timed_plan = self.queue.pop();
+            match next_timed_plan {
+                Some(timed_plan) => {
+                    if self.invalid_set.contains(&timed_plan.plan_id) {
+                        self.invalid_set.remove(&timed_plan.plan_id);
+                    } else {
+                        return Some(timed_plan);
+                    }
+                }
+                None => {
+                    return None;
+                }
+            }
+        }
+    }
+}

From b0b63f33a07ad4217ee05b1bcf9ea6badad40fb9 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Tue, 30 Jul 2024 23:35:04 -0400
Subject: [PATCH 06/17] Rename plugin concept

---
 src/context.rs | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index dd10654..e0dbc93 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -3,27 +3,27 @@ use std::{
     collections::{HashMap, VecDeque},
 };
 
-pub trait Plugin: Any {
+pub trait DataPlugin: Any {
     type DataContainer;
 
-    fn get_data_container() -> Self::DataContainer;
+    fn create_data_container() -> Self::DataContainer;
 }
 
 #[macro_export]
-macro_rules! define_plugin {
+macro_rules! define_data_plugin {
     ($plugin:ident, $data_container:ty, $default: expr) => {
         struct $plugin {}
 
-        impl $crate::context::Plugin for $plugin {
+        impl $crate::context::DataPlugin for $plugin {
             type DataContainer = $data_container;
 
-            fn get_data_container() -> Self::DataContainer {
+            fn create_data_container() -> Self::DataContainer {
                 $default
             }
         }
     };
 }
-pub use define_plugin;
+pub use define_data_plugin;
 
 use crate::plan::{PlanId, PlanQueue};
 
@@ -31,7 +31,7 @@ type Callback = dyn FnOnce(&mut Context);
 pub struct Context {
     plan_queue: PlanQueue,
     callback_queue: VecDeque<Box<Callback>>,
-    plugin_data: HashMap<TypeId, Box<dyn Any>>,
+    data_plugins: HashMap<TypeId, Box<dyn Any>>,
     current_time: f64,
 }
 
@@ -40,7 +40,7 @@ impl Context {
         Context {
             plan_queue: PlanQueue::new(),
             callback_queue: VecDeque::new(),
-            plugin_data: HashMap::new(),
+            data_plugins: HashMap::new(),
             current_time: 0.0,
         }
     }
@@ -58,29 +58,29 @@ impl Context {
         self.callback_queue.push_back(Box::new(callback));
     }
 
-    fn add_plugin<T: Plugin>(&mut self) {
-        self.plugin_data
-            .insert(TypeId::of::<T>(), Box::new(T::get_data_container()));
+    fn add_plugin<T: DataPlugin>(&mut self) {
+        self.data_plugins
+            .insert(TypeId::of::<T>(), Box::new(T::create_data_container()));
     }
 
-    pub fn get_data_container_mut<T: Plugin>(&mut self) -> &mut T::DataContainer {
+    pub fn get_data_container_mut<T: DataPlugin>(&mut self) -> &mut T::DataContainer {
         let type_id = &TypeId::of::<T>();
-        if !self.plugin_data.contains_key(type_id) {
+        if !self.data_plugins.contains_key(type_id) {
             self.add_plugin::<T>();
         }
-        self.plugin_data
+        self.data_plugins
             .get_mut(type_id)
             .unwrap()
             .downcast_mut::<T::DataContainer>()
             .unwrap()
     }
 
-    pub fn get_data_container<T: Plugin>(&self) -> Option<&T::DataContainer> {
+    pub fn get_data_container<T: DataPlugin>(&self) -> Option<&T::DataContainer> {
         let type_id = &TypeId::of::<T>();
-        if !self.plugin_data.contains_key(type_id) {
+        if !self.data_plugins.contains_key(type_id) {
             return None;
         }
-        self.plugin_data
+        self.data_plugins
             .get(type_id)
             .unwrap()
             .downcast_ref::<T::DataContainer>()
@@ -121,7 +121,7 @@ impl Default for Context {
 mod tests {
     use super::*;
 
-    define_plugin!(ComponentA, u32, 0);
+    define_data_plugin!(ComponentA, u32, 0);
 
     impl ComponentA {
         fn increment_counter(context: &mut Context) {

From bf06af8b2178f7275baf498be8b97046d149fd26 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Wed, 31 Jul 2024 00:58:01 -0400
Subject: [PATCH 07/17] Add tests for plan queue and handle invalid
 cancellation

---
 Cargo.toml     |   1 -
 src/context.rs |  14 +++---
 src/plan.rs    | 130 ++++++++++++++++++++++++++++++++-----------------
 3 files changed, 91 insertions(+), 54 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6e0c8e1..be2f210 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,4 +8,3 @@ license = "Apache-2.0"
 homepage = "https://github.com/CDCgov/ixa"
 
 [dependencies]
-derivative = "2.2.0"
diff --git a/src/context.rs b/src/context.rs
index e0dbc93..ffbe03a 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -29,7 +29,7 @@ use crate::plan::{PlanId, PlanQueue};
 
 type Callback = dyn FnOnce(&mut Context);
 pub struct Context {
-    plan_queue: PlanQueue,
+    plan_queue: PlanQueue<Box<Callback>>,
     callback_queue: VecDeque<Box<Callback>>,
     data_plugins: HashMap<TypeId, Box<dyn Any>>,
     current_time: f64,
@@ -47,7 +47,7 @@ impl Context {
 
     pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
         // TODO: Handle invalid times (past, NAN, etc)
-        self.plan_queue.add_plan(time, callback)
+        self.plan_queue.add_plan(time, Box::new(callback))
     }
 
     pub fn cancel_plan(&mut self, id: PlanId) {
@@ -99,12 +99,12 @@ impl Context {
                 continue;
             }
 
-            // There aren't any callbacks, so look at the first timed plan.
-            if let Some(timed_plan) = self.plan_queue.get_next_timed_plan() {
-                self.current_time = timed_plan.time;
-                (timed_plan.callback)(self);
+            // There aren't any callbacks, so look at the first plan.
+            if let Some(plan) = self.plan_queue.get_next_plan() {
+                self.current_time = plan.time;
+                (plan.data)(self);
             } else {
-                // OK, there aren't any timed plans, so we're done.
+                // OK, there aren't any plans, so we're done.
                 break;
             }
         }
diff --git a/src/plan.rs b/src/plan.rs
index 15c1bd3..532a851 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -1,90 +1,87 @@
 use std::{
     cmp::Ordering,
-    collections::{BinaryHeap, HashSet},
+    collections::{BinaryHeap, HashMap},
 };
 
-use derivative::Derivative;
-
-use crate::context::Context;
-
 pub struct PlanId {
-    id: u64,
+    id: usize,
 }
 
-#[derive(Derivative)]
-#[derivative(Eq, PartialEq, Debug)]
-pub struct TimedPlan {
+pub struct Plan<T> {
     pub time: f64,
-    plan_id: u64,
-    #[derivative(PartialEq = "ignore", Debug = "ignore")]
-    pub callback: Box<dyn FnOnce(&mut Context)>,
+    pub data: T,
 }
 
-impl Ord for TimedPlan {
+#[derive(PartialEq, Debug)]
+pub struct PlanRecord {
+    pub time: f64,
+    id: usize,
+}
+
+impl Eq for PlanRecord {}
+
+impl PartialOrd for PlanRecord {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for PlanRecord {
     fn cmp(&self, other: &Self) -> Ordering {
         let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
         if time_ordering == Ordering::Equal {
             // Break time ties in order of plan id
-            self.plan_id.cmp(&other.plan_id).reverse()
+            self.id.cmp(&other.id).reverse()
         } else {
             time_ordering
         }
     }
 }
 
-impl PartialOrd for TimedPlan {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
 #[derive(Debug)]
-pub struct PlanQueue {
-    queue: BinaryHeap<TimedPlan>,
-    invalid_set: HashSet<u64>,
-    plan_counter: u64,
+pub struct PlanQueue<T> {
+    queue: BinaryHeap<PlanRecord>,
+    data_map: HashMap<usize, T>,
+    plan_counter: usize,
 }
 
-impl Default for PlanQueue {
+impl<T> Default for PlanQueue<T> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl PlanQueue {
-    pub fn new() -> PlanQueue {
+impl<T> PlanQueue<T> {
+    pub fn new() -> PlanQueue<T> {
         PlanQueue {
             queue: BinaryHeap::new(),
-            invalid_set: HashSet::new(),
+            data_map: HashMap::new(),
             plan_counter: 0,
         }
     }
 
-    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
-        // Add plan to queue and increment counter
-        let plan_id = self.plan_counter;
-        self.queue.push(TimedPlan {
-            time,
-            plan_id,
-            callback: Box::new(callback),
-        });
+    pub fn add_plan(&mut self, time: f64, data: T) -> PlanId {
+        // Add plan to queue, store data, and increment counter
+        let id = self.plan_counter;
+        self.queue.push(PlanRecord { time, id });
+        self.data_map.insert(id, data);
         self.plan_counter += 1;
-        PlanId { id: plan_id }
+        PlanId { id }
     }
 
     pub fn cancel_plan(&mut self, id: PlanId) {
-        self.invalid_set.insert(id.id);
+        self.data_map.remove(&id.id).expect("Plan does not exist");
     }
 
-    pub fn get_next_timed_plan(&mut self) -> Option<TimedPlan> {
+    pub fn get_next_plan(&mut self) -> Option<Plan<T>> {
         loop {
-            let next_timed_plan = self.queue.pop();
-            match next_timed_plan {
-                Some(timed_plan) => {
-                    if self.invalid_set.contains(&timed_plan.plan_id) {
-                        self.invalid_set.remove(&timed_plan.plan_id);
-                    } else {
-                        return Some(timed_plan);
+            match self.queue.pop() {
+                Some(plan_record) => {
+                    if let Some(data) = self.data_map.remove(&plan_record.id) {
+                        return Some(Plan {
+                            time: plan_record.time,
+                            data,
+                        });
                     }
                 }
                 None => {
@@ -94,3 +91,44 @@ impl PlanQueue {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::PlanQueue;
+
+    #[test]
+    fn test_add_cancel() {
+        // Add some plans and cancel and make sure ordering occurs as expected
+        let mut plan_queue = PlanQueue::<usize>::new();
+        plan_queue.add_plan(1.0, 1);
+        plan_queue.add_plan(3.0, 3);
+        plan_queue.add_plan(3.0, 4);
+        let plan_to_cancel = plan_queue.add_plan(1.5, 0);
+        plan_queue.add_plan(2.0, 2);
+        plan_queue.cancel_plan(plan_to_cancel);
+
+        assert_eq!(plan_queue.get_next_plan().unwrap().time, 1.0);
+        assert_eq!(plan_queue.get_next_plan().unwrap().time, 2.0);
+
+        // Check tie handling
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 3.0);
+        assert_eq!(next_plan.data, 3);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 3.0);
+        assert_eq!(next_plan.data, 4);
+
+        assert!(plan_queue.get_next_plan().is_none());
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_invalid_cancel() {
+        // Cancel a plan that has already occured and make sure it panics
+        let mut plan_queue = PlanQueue::<()>::new();
+        let plan_to_cancel = plan_queue.add_plan(1.0, ());
+        plan_queue.get_next_plan();
+        plan_queue.cancel_plan(plan_to_cancel);
+    }
+}

From f021aeed2e464715a5c16810269f2ae93bc07e1a Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Wed, 31 Jul 2024 19:05:14 -0700
Subject: [PATCH 08/17] Refactor context tests so that they test individual
 behaviors.

---
 src/context.rs | 161 +++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 143 insertions(+), 18 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index ffbe03a..c58db4c 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -46,7 +46,9 @@ impl Context {
     }
 
     pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
-        // TODO: Handle invalid times (past, NAN, etc)
+        if time.is_nan() || time.is_infinite() || time < self.current_time {
+            panic!("Invalid time value");
+        }
         self.plan_queue.add_plan(time, Box::new(callback))
     }
 
@@ -121,32 +123,155 @@ impl Default for Context {
 mod tests {
     use super::*;
 
-    define_data_plugin!(ComponentA, u32, 0);
+    define_data_plugin!(ComponentA, Vec<u32>, vec![]);
 
-    impl ComponentA {
-        fn increment_counter(context: &mut Context) {
-            *(context.get_data_container_mut::<ComponentA>()) += 1;
-        }
+    fn add_plan(context: &mut Context, time: f64, value: u32) -> PlanId {
+        context.add_plan(time, move |context| {
+            context.get_data_container_mut::<ComponentA>().push(value);
+        })
+    }
 
-        fn init(context: &mut Context) {
-            context.add_plan(1.0, Self::increment_counter);
-        }
+    #[test]
+    #[should_panic]
+    fn negative_plan_time() {
+        let mut context = Context::new();
+        add_plan(&mut context, -1.0, 0);
     }
 
     #[test]
-    fn test_component_and_planning() {
+    #[should_panic]
+    fn infinite_plan_time() {
+        let mut context = Context::new();
+        add_plan(&mut context, f64::INFINITY, 0);
+    }
+    
+    #[test]
+    #[should_panic]
+    fn nan_plan_time() {
+        let mut context = Context::new();
+        add_plan(&mut context, f64::NAN, 0);
+    }
+    
+
+    #[test]
+    fn empty_context() {
+        let mut context = Context::new();
+        context.execute();
+        assert_eq!(context.get_current_time(), 0.0);        
+    }
+
+    #[test]    
+    fn timed_plan_only() {
         let mut context = Context::new();
-        ComponentA::init(&mut context);
-        assert_eq!(context.get_current_time(), 0.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 0);
+        add_plan(&mut context, 1.0, 1);
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 1);
-        let plan_to_cancel = context.add_plan(3.0, ComponentA::increment_counter);
-        context.add_plan(2.0, ComponentA::increment_counter);
-        context.cancel_plan(plan_to_cancel);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1]);        
+    }
+
+    #[test]
+    fn callback_only() {
+        let mut context = Context::new();
+        context.queue_callback(|context| {
+            context.get_data_container_mut::<ComponentA>().push(1);            
+        });
+        context.execute();
+        assert_eq!(context.get_current_time(), 0.0);        
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1]);
+    }
+
+    #[test]
+    fn callback_before_timed_plan() {
+        let mut context = Context::new();
+        context.queue_callback(|context| {
+            context.get_data_container_mut::<ComponentA>().push(1);            
+        });
+        add_plan(&mut context, 1.0, 2);
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);                
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);        
+    }
+
+    #[test]
+    fn callback_adds_timed_plan() {
+        let mut context = Context::new();
+        context.queue_callback(|context| {
+            context.get_data_container_mut::<ComponentA>().push(1);
+            add_plan(context, 1.0, 2);
+            context.get_data_container_mut::<ComponentA>().push(3);            
+        });
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);                
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 2]);        
+    }
+
+    #[test]
+    fn callback_adds_callback_and_timed_plan() {
+        let mut context = Context::new();
+        context.queue_callback(|context| {
+            context.get_data_container_mut::<ComponentA>().push(1);
+            add_plan(context, 1.0, 2);
+            context.queue_callback(|context| {
+                context.get_data_container_mut::<ComponentA>().push(4);
+            });
+            context.get_data_container_mut::<ComponentA>().push(3);            
+        });
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 4, 2]);        
+    }
+
+    #[test]
+    fn timed_plan_adds_callback_and_timed_plan() {
+        let mut context = Context::new();
+        context.add_plan(1.0, |context| {
+            context.get_data_container_mut::<ComponentA>().push(1);
+            // We add the plan first, but the callback will fire first.
+            add_plan(context, 2.0, 3);                
+            context.queue_callback(|context| {
+                context.get_data_container_mut::<ComponentA>().push(2);
+            });
+        });
         context.execute();
         assert_eq!(context.get_current_time(), 2.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), 2);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2, 3]);        
+    }
+
+    #[test]
+    fn cancel_plan() {
+        let mut context = Context::new();
+        let to_cancel = add_plan(&mut context, 2.0, 1);
+        context.add_plan(1.0, move |context| {
+            context.cancel_plan(to_cancel);
+        });
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![]);                
+    }
+
+    #[test]
+    fn add_plan_with_current_time() {
+        let mut context = Context::new();
+        context.add_plan(1.0, move |context| {
+            context.get_data_container_mut::<ComponentA>().push(1);
+            add_plan(context, 1.0, 2);
+            context.queue_callback(|context| {
+                context.get_data_container_mut::<ComponentA>().push(3);
+            });
+        });
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 2]);                
+    }        
+
+    #[test]
+    fn plans_at_same_time_fire_in_order() {
+        let mut context = Context::new();
+        add_plan(&mut context, 1.0, 1);
+        add_plan(&mut context, 1.0, 2);
+        context.execute();
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);                
+        
     }
 }

From f2875b3e19d7815130116bfd345af3db55f2acf0 Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Wed, 31 Jul 2024 19:44:57 -0700
Subject: [PATCH 09/17] Format

---
 src/context.rs | 54 ++++++++++++++++++++++++++++++--------------------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index c58db4c..afeff4e 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -144,39 +144,38 @@ mod tests {
         let mut context = Context::new();
         add_plan(&mut context, f64::INFINITY, 0);
     }
-    
+
     #[test]
     #[should_panic]
     fn nan_plan_time() {
         let mut context = Context::new();
         add_plan(&mut context, f64::NAN, 0);
     }
-    
 
     #[test]
     fn empty_context() {
         let mut context = Context::new();
         context.execute();
-        assert_eq!(context.get_current_time(), 0.0);        
+        assert_eq!(context.get_current_time(), 0.0);
     }
 
-    #[test]    
+    #[test]
     fn timed_plan_only() {
         let mut context = Context::new();
         add_plan(&mut context, 1.0, 1);
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1]);        
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1]);
     }
 
     #[test]
     fn callback_only() {
         let mut context = Context::new();
         context.queue_callback(|context| {
-            context.get_data_container_mut::<ComponentA>().push(1);            
+            context.get_data_container_mut::<ComponentA>().push(1);
         });
         context.execute();
-        assert_eq!(context.get_current_time(), 0.0);        
+        assert_eq!(context.get_current_time(), 0.0);
         assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1]);
     }
 
@@ -184,12 +183,12 @@ mod tests {
     fn callback_before_timed_plan() {
         let mut context = Context::new();
         context.queue_callback(|context| {
-            context.get_data_container_mut::<ComponentA>().push(1);            
+            context.get_data_container_mut::<ComponentA>().push(1);
         });
         add_plan(&mut context, 1.0, 2);
         context.execute();
-        assert_eq!(context.get_current_time(), 1.0);                
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);        
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);
     }
 
     #[test]
@@ -198,11 +197,14 @@ mod tests {
         context.queue_callback(|context| {
             context.get_data_container_mut::<ComponentA>().push(1);
             add_plan(context, 1.0, 2);
-            context.get_data_container_mut::<ComponentA>().push(3);            
+            context.get_data_container_mut::<ComponentA>().push(3);
         });
         context.execute();
-        assert_eq!(context.get_current_time(), 1.0);                
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 2]);        
+        assert_eq!(context.get_current_time(), 1.0);
+        assert_eq!(
+            *context.get_data_container_mut::<ComponentA>(),
+            vec![1, 3, 2]
+        );
     }
 
     #[test]
@@ -214,11 +216,14 @@ mod tests {
             context.queue_callback(|context| {
                 context.get_data_container_mut::<ComponentA>().push(4);
             });
-            context.get_data_container_mut::<ComponentA>().push(3);            
+            context.get_data_container_mut::<ComponentA>().push(3);
         });
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 4, 2]);        
+        assert_eq!(
+            *context.get_data_container_mut::<ComponentA>(),
+            vec![1, 3, 4, 2]
+        );
     }
 
     #[test]
@@ -227,14 +232,17 @@ mod tests {
         context.add_plan(1.0, |context| {
             context.get_data_container_mut::<ComponentA>().push(1);
             // We add the plan first, but the callback will fire first.
-            add_plan(context, 2.0, 3);                
+            add_plan(context, 2.0, 3);
             context.queue_callback(|context| {
                 context.get_data_container_mut::<ComponentA>().push(2);
             });
         });
         context.execute();
         assert_eq!(context.get_current_time(), 2.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2, 3]);        
+        assert_eq!(
+            *context.get_data_container_mut::<ComponentA>(),
+            vec![1, 2, 3]
+        );
     }
 
     #[test]
@@ -246,7 +254,7 @@ mod tests {
         });
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![]);                
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![]);
     }
 
     #[test]
@@ -261,8 +269,11 @@ mod tests {
         });
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 3, 2]);                
-    }        
+        assert_eq!(
+            *context.get_data_container_mut::<ComponentA>(),
+            vec![1, 3, 2]
+        );
+    }
 
     #[test]
     fn plans_at_same_time_fire_in_order() {
@@ -271,7 +282,6 @@ mod tests {
         add_plan(&mut context, 1.0, 2);
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
-        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);                
-        
+        assert_eq!(*context.get_data_container_mut::<ComponentA>(), vec![1, 2]);
     }
 }

From d4ca244c83f9ff33b1f335559782e4e9e0692b4e Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Wed, 31 Jul 2024 20:13:26 -0700
Subject: [PATCH 10/17] Fix the first clippy errors

---
 src/context.rs | 15 +++++++++++----
 src/plan.rs    |  5 +++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index afeff4e..80f5c8e 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -36,6 +36,7 @@ pub struct Context {
 }
 
 impl Context {
+    #[must_use]
     pub fn new() -> Context {
         Context {
             plan_queue: PlanQueue::new(),
@@ -45,10 +46,11 @@ impl Context {
         }
     }
 
+    /// # Panics
+    ///
+    /// Panics if time is in the past, infinite, or NaN.
     pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
-        if time.is_nan() || time.is_infinite() || time < self.current_time {
-            panic!("Invalid time value");
-        }
+        assert!(!time.is_nan() && !time.is_infinite() && time >= self.current_time);
         self.plan_queue.add_plan(time, Box::new(callback))
     }
 
@@ -65,6 +67,8 @@ impl Context {
             .insert(TypeId::of::<T>(), Box::new(T::create_data_container()));
     }
 
+    /// # Panics
+    #[must_use]
     pub fn get_data_container_mut<T: DataPlugin>(&mut self) -> &mut T::DataContainer {
         let type_id = &TypeId::of::<T>();
         if !self.data_plugins.contains_key(type_id) {
@@ -76,7 +80,9 @@ impl Context {
             .downcast_mut::<T::DataContainer>()
             .unwrap()
     }
-
+    
+    /// # Panics
+    #[must_use]
     pub fn get_data_container<T: DataPlugin>(&self) -> Option<&T::DataContainer> {
         let type_id = &TypeId::of::<T>();
         if !self.data_plugins.contains_key(type_id) {
@@ -88,6 +94,7 @@ impl Context {
             .downcast_ref::<T::DataContainer>()
     }
 
+    #[must_use]
     pub fn get_current_time(&self) -> f64 {
         self.current_time
     }
diff --git a/src/plan.rs b/src/plan.rs
index 532a851..3563575 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -52,6 +52,7 @@ impl<T> Default for PlanQueue<T> {
 }
 
 impl<T> PlanQueue<T> {
+    #[must_use]
     pub fn new() -> PlanQueue<T> {
         PlanQueue {
             queue: BinaryHeap::new(),
@@ -69,6 +70,10 @@ impl<T> PlanQueue<T> {
         PlanId { id }
     }
 
+    /// # Panics
+    ///
+    /// This function panics if you cancel a plan which has already
+    /// been cancelled or executed.
     pub fn cancel_plan(&mut self, id: PlanId) {
         self.data_map.remove(&id.id).expect("Plan does not exist");
     }

From 8117ed79c2475374f1d2696245a936939519f18e Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Wed, 31 Jul 2024 20:21:27 -0700
Subject: [PATCH 11/17] Clippy fixes

---
 src/context.rs | 14 +++++++-------
 src/plan.rs    | 40 ++++++++++++++++++++--------------------
 2 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 80f5c8e..173b390 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -25,11 +25,11 @@ macro_rules! define_data_plugin {
 }
 pub use define_data_plugin;
 
-use crate::plan::{PlanId, PlanQueue};
+use crate::plan::{Id, Queue};
 
 type Callback = dyn FnOnce(&mut Context);
 pub struct Context {
-    plan_queue: PlanQueue<Box<Callback>>,
+    plan_queue: Queue<Box<Callback>>,
     callback_queue: VecDeque<Box<Callback>>,
     data_plugins: HashMap<TypeId, Box<dyn Any>>,
     current_time: f64,
@@ -39,7 +39,7 @@ impl Context {
     #[must_use]
     pub fn new() -> Context {
         Context {
-            plan_queue: PlanQueue::new(),
+            plan_queue: Queue::new(),
             callback_queue: VecDeque::new(),
             data_plugins: HashMap::new(),
             current_time: 0.0,
@@ -49,12 +49,12 @@ impl Context {
     /// # Panics
     ///
     /// Panics if time is in the past, infinite, or NaN.
-    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> PlanId {
+    pub fn add_plan(&mut self, time: f64, callback: impl FnOnce(&mut Context) + 'static) -> Id {
         assert!(!time.is_nan() && !time.is_infinite() && time >= self.current_time);
         self.plan_queue.add_plan(time, Box::new(callback))
     }
 
-    pub fn cancel_plan(&mut self, id: PlanId) {
+    pub fn cancel_plan(&mut self, id: &Id) {
         self.plan_queue.cancel_plan(id);
     }
 
@@ -132,7 +132,7 @@ mod tests {
 
     define_data_plugin!(ComponentA, Vec<u32>, vec![]);
 
-    fn add_plan(context: &mut Context, time: f64, value: u32) -> PlanId {
+    fn add_plan(context: &mut Context, time: f64, value: u32) -> Id {
         context.add_plan(time, move |context| {
             context.get_data_container_mut::<ComponentA>().push(value);
         })
@@ -257,7 +257,7 @@ mod tests {
         let mut context = Context::new();
         let to_cancel = add_plan(&mut context, 2.0, 1);
         context.add_plan(1.0, move |context| {
-            context.cancel_plan(to_cancel);
+            context.cancel_plan(&to_cancel);
         });
         context.execute();
         assert_eq!(context.get_current_time(), 1.0);
diff --git a/src/plan.rs b/src/plan.rs
index 3563575..8dccd71 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -3,7 +3,7 @@ use std::{
     collections::{BinaryHeap, HashMap},
 };
 
-pub struct PlanId {
+pub struct Id {
     id: usize,
 }
 
@@ -13,20 +13,20 @@ pub struct Plan<T> {
 }
 
 #[derive(PartialEq, Debug)]
-pub struct PlanRecord {
+pub struct Record {
     pub time: f64,
     id: usize,
 }
 
-impl Eq for PlanRecord {}
+impl Eq for Record {}
 
-impl PartialOrd for PlanRecord {
+impl PartialOrd for Record {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl Ord for PlanRecord {
+impl Ord for Record {
     fn cmp(&self, other: &Self) -> Ordering {
         let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
         if time_ordering == Ordering::Equal {
@@ -39,42 +39,42 @@ impl Ord for PlanRecord {
 }
 
 #[derive(Debug)]
-pub struct PlanQueue<T> {
-    queue: BinaryHeap<PlanRecord>,
+pub struct Queue<T> {
+    queue: BinaryHeap<Record>,
     data_map: HashMap<usize, T>,
     plan_counter: usize,
 }
 
-impl<T> Default for PlanQueue<T> {
+impl<T> Default for Queue<T> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl<T> PlanQueue<T> {
+impl<T> Queue<T> {
     #[must_use]
-    pub fn new() -> PlanQueue<T> {
-        PlanQueue {
+    pub fn new() -> Queue<T> {
+        Queue {
             queue: BinaryHeap::new(),
             data_map: HashMap::new(),
             plan_counter: 0,
         }
     }
 
-    pub fn add_plan(&mut self, time: f64, data: T) -> PlanId {
+    pub fn add_plan(&mut self, time: f64, data: T) -> Id {
         // Add plan to queue, store data, and increment counter
         let id = self.plan_counter;
-        self.queue.push(PlanRecord { time, id });
+        self.queue.push(Record { time, id });
         self.data_map.insert(id, data);
         self.plan_counter += 1;
-        PlanId { id }
+        Id { id }
     }
 
     /// # Panics
     ///
     /// This function panics if you cancel a plan which has already
     /// been cancelled or executed.
-    pub fn cancel_plan(&mut self, id: PlanId) {
+    pub fn cancel_plan(&mut self, id: &Id) {
         self.data_map.remove(&id.id).expect("Plan does not exist");
     }
 
@@ -99,18 +99,18 @@ impl<T> PlanQueue<T> {
 
 #[cfg(test)]
 mod tests {
-    use super::PlanQueue;
+    use super::Queue;
 
     #[test]
     fn test_add_cancel() {
         // Add some plans and cancel and make sure ordering occurs as expected
-        let mut plan_queue = PlanQueue::<usize>::new();
+        let mut plan_queue = Queue::<usize>::new();
         plan_queue.add_plan(1.0, 1);
         plan_queue.add_plan(3.0, 3);
         plan_queue.add_plan(3.0, 4);
         let plan_to_cancel = plan_queue.add_plan(1.5, 0);
         plan_queue.add_plan(2.0, 2);
-        plan_queue.cancel_plan(plan_to_cancel);
+        plan_queue.cancel_plan(&plan_to_cancel);
 
         assert_eq!(plan_queue.get_next_plan().unwrap().time, 1.0);
         assert_eq!(plan_queue.get_next_plan().unwrap().time, 2.0);
@@ -131,9 +131,9 @@ mod tests {
     #[should_panic]
     fn test_invalid_cancel() {
         // Cancel a plan that has already occured and make sure it panics
-        let mut plan_queue = PlanQueue::<()>::new();
+        let mut plan_queue = Queue::<()>::new();
         let plan_to_cancel = plan_queue.add_plan(1.0, ());
         plan_queue.get_next_plan();
-        plan_queue.cancel_plan(plan_to_cancel);
+        plan_queue.cancel_plan(&plan_to_cancel);
     }
 }

From dd4a35506152be7b3f9e06b19038cfb84ceb6a52 Mon Sep 17 00:00:00 2001
From: EKR <ekr@rtfm.com>
Date: Wed, 31 Jul 2024 20:21:34 -0700
Subject: [PATCH 12/17] Fmt

---
 src/context.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/context.rs b/src/context.rs
index 173b390..e568f9f 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -80,7 +80,7 @@ impl Context {
             .downcast_mut::<T::DataContainer>()
             .unwrap()
     }
-    
+
     /// # Panics
     #[must_use]
     pub fn get_data_container<T: DataPlugin>(&self) -> Option<&T::DataContainer> {

From d165d0e28093c122ff8f38b741434e3b71ddd78b Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Thu, 1 Aug 2024 00:23:25 -0400
Subject: [PATCH 13/17] Update data container code and handle pedantic clippy

---
 src/context.rs | 29 ++++++++---------------------
 1 file changed, 8 insertions(+), 21 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index e568f9f..e64a2a0 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -62,36 +62,23 @@ impl Context {
         self.callback_queue.push_back(Box::new(callback));
     }
 
-    fn add_plugin<T: DataPlugin>(&mut self) {
-        self.data_plugins
-            .insert(TypeId::of::<T>(), Box::new(T::create_data_container()));
-    }
-
-    /// # Panics
     #[must_use]
+    #[allow(clippy::missing_panics_doc)]
     pub fn get_data_container_mut<T: DataPlugin>(&mut self) -> &mut T::DataContainer {
-        let type_id = &TypeId::of::<T>();
-        if !self.data_plugins.contains_key(type_id) {
-            self.add_plugin::<T>();
-        }
         self.data_plugins
-            .get_mut(type_id)
-            .unwrap()
+            .entry(TypeId::of::<T>())
+            .or_insert_with(|| Box::new(T::create_data_container()))
             .downcast_mut::<T::DataContainer>()
-            .unwrap()
+            .unwrap() // Will never panic as data container has the matching type
     }
 
-    /// # Panics
     #[must_use]
     pub fn get_data_container<T: DataPlugin>(&self) -> Option<&T::DataContainer> {
-        let type_id = &TypeId::of::<T>();
-        if !self.data_plugins.contains_key(type_id) {
-            return None;
+        if let Some(data) = self.data_plugins.get(&TypeId::of::<T>()) {
+            data.downcast_ref::<T::DataContainer>()
+        } else {
+            None
         }
-        self.data_plugins
-            .get(type_id)
-            .unwrap()
-            .downcast_ref::<T::DataContainer>()
     }
 
     #[must_use]

From 8010a4d3ec0f9ef8c8239409e8e0f94a9533ce87 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Thu, 1 Aug 2024 23:42:49 -0400
Subject: [PATCH 14/17] Document plan code and split up tests

---
 src/plan.rs | 197 ++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 136 insertions(+), 61 deletions(-)

diff --git a/src/plan.rs b/src/plan.rs
index 8dccd71..d2ddcab 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -1,57 +1,37 @@
+//! A priority queue that stores arbitrary data sorted by time
+//!
+//! Defines a `Queue<T>` that is intended to store a queue of items of type T,
+//! sorted by `f64` time, called 'plans'. This queue has methods for adding
+//! plans, cancelling plans, and retrieving the earliest plan in the queue.
+//! Adding a plan is *O*(log(*n*)) while cancellation and retrieval are *O*(1).
+//!
+//! This queue is used by `Context` to store future events where some callback
+//! closure `FnOnce(&mut Context)` will be executed at a given point in time.
+
 use std::{
     cmp::Ordering,
     collections::{BinaryHeap, HashMap},
 };
 
-pub struct Id {
-    id: usize,
-}
-
-pub struct Plan<T> {
-    pub time: f64,
-    pub data: T,
-}
-
-#[derive(PartialEq, Debug)]
-pub struct Record {
-    pub time: f64,
-    id: usize,
-}
-
-impl Eq for Record {}
-
-impl PartialOrd for Record {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for Record {
-    fn cmp(&self, other: &Self) -> Ordering {
-        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
-        if time_ordering == Ordering::Equal {
-            // Break time ties in order of plan id
-            self.id.cmp(&other.id).reverse()
-        } else {
-            time_ordering
-        }
-    }
-}
-
-#[derive(Debug)]
+/// A priority queue that stores arbitrary data sorted by time
+///
+/// Items of type T are stored in order by `f64` time and called `Plan<T>`.
+/// When plans are created they are sequentially assigned an `Id` that is a
+/// wrapped `u64`. If two plans are scheduled for the same time the plan that is
+/// scheduled first (i.e. that has the lowest id) is placed earlier.
+///
+/// The pair of time and plan id are stored in a binary heap of `Entry` objects.
+/// The data payload of the event is stored in a hash map by plan id.
+/// Plan cancellation occurs by removing the corresponding entry from the data
+/// hash map.
 pub struct Queue<T> {
-    queue: BinaryHeap<Record>,
-    data_map: HashMap<usize, T>,
-    plan_counter: usize,
-}
-
-impl<T> Default for Queue<T> {
-    fn default() -> Self {
-        Self::new()
-    }
+    queue: BinaryHeap<Entry>,
+    data_map: HashMap<u64, T>,
+    plan_counter: u64,
 }
 
 impl<T> Queue<T> {
+    /// Creates a new empty `Queue<T>`
     #[must_use]
     pub fn new() -> Queue<T> {
         Queue {
@@ -61,30 +41,43 @@ impl<T> Queue<T> {
         }
     }
 
+    /// Add a plan to the queue at the specified time
+    ///
+    /// Returns an `Id` for the newly-added plan that can be used to cancel it
+    /// if needed.
     pub fn add_plan(&mut self, time: f64, data: T) -> Id {
         // Add plan to queue, store data, and increment counter
         let id = self.plan_counter;
-        self.queue.push(Record { time, id });
+        self.queue.push(Entry { time, id });
         self.data_map.insert(id, data);
         self.plan_counter += 1;
         Id { id }
     }
 
+    /// Cancel a plan that has been added to the queue
+    ///
     /// # Panics
     ///
     /// This function panics if you cancel a plan which has already
     /// been cancelled or executed.
     pub fn cancel_plan(&mut self, id: &Id) {
+        // Delete the plan from the map, but leave in the queue
+        // It will be skipped when the plan is popped from the queue
         self.data_map.remove(&id.id).expect("Plan does not exist");
     }
 
+    /// Retrieve the earliest plan in the queue
+    ///
+    /// Returns the next plan if it exists or else `None` if the queue is empty
     pub fn get_next_plan(&mut self) -> Option<Plan<T>> {
         loop {
+            // Pop from queue until we find a plan with data or queue is empty
             match self.queue.pop() {
-                Some(plan_record) => {
-                    if let Some(data) = self.data_map.remove(&plan_record.id) {
+                Some(entry) => {
+                    // Skip plans that have been cancelled and thus have no data
+                    if let Some(data) = self.data_map.remove(&entry.id) {
                         return Some(Plan {
-                            time: plan_record.time,
+                            time: entry.time,
                             data,
                         });
                     }
@@ -97,40 +90,122 @@ impl<T> Queue<T> {
     }
 }
 
+impl<T> Default for Queue<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// A time and id pair used to order plans in the `Queue<T>`
+///
+/// `Entry` objects are sorted in increasing order of time and then plan id
+#[derive(PartialEq, Debug)]
+struct Entry {
+    time: f64,
+    id: u64,
+}
+
+impl Eq for Entry {}
+
+impl PartialOrd for Entry {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Entry {
+    fn cmp(&self, other: &Self) -> Ordering {
+        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
+        match time_ordering {
+            // Break time ties in order of plan id
+            Ordering::Equal => self.id.cmp(&other.id).reverse(),
+            _ => time_ordering,
+        }
+    }
+}
+
+/// A unique identifier for a plan added to a `Queue<T>`
+pub struct Id {
+    id: u64,
+}
+
+/// A plan that holds data of type `T` intended to be used at the specified time
+pub struct Plan<T> {
+    pub time: f64,
+    pub data: T,
+}
+
 #[cfg(test)]
 mod tests {
     use super::Queue;
 
     #[test]
-    fn test_add_cancel() {
-        // Add some plans and cancel and make sure ordering occurs as expected
-        let mut plan_queue = Queue::<usize>::new();
+    fn empty_queue() {
+        let mut plan_queue = Queue::<()>::new();
+        assert!(plan_queue.get_next_plan().is_none());
+    }
+
+    #[test]
+    fn add_plans() {
+        let mut plan_queue = Queue::new();
         plan_queue.add_plan(1.0, 1);
         plan_queue.add_plan(3.0, 3);
-        plan_queue.add_plan(3.0, 4);
-        let plan_to_cancel = plan_queue.add_plan(1.5, 0);
         plan_queue.add_plan(2.0, 2);
-        plan_queue.cancel_plan(&plan_to_cancel);
 
-        assert_eq!(plan_queue.get_next_plan().unwrap().time, 1.0);
-        assert_eq!(plan_queue.get_next_plan().unwrap().time, 2.0);
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 1.0);
+        assert_eq!(next_plan.data, 1);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 2.0);
+        assert_eq!(next_plan.data, 2);
 
-        // Check tie handling
         let next_plan = plan_queue.get_next_plan().unwrap();
         assert_eq!(next_plan.time, 3.0);
         assert_eq!(next_plan.data, 3);
 
+        assert!(plan_queue.get_next_plan().is_none());
+    }
+
+    #[test]
+    fn add_plans_at_same_time() {
+        let mut plan_queue = Queue::new();
+        plan_queue.add_plan(1.0, 1);
+        plan_queue.add_plan(1.0, 2);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 1.0);
+        assert_eq!(next_plan.data, 1);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 1.0);
+        assert_eq!(next_plan.data, 2);
+
+        assert!(plan_queue.get_next_plan().is_none());
+    }
+
+    #[test]
+    fn add_and_cancel_plans() {
+        let mut plan_queue = Queue::new();
+        plan_queue.add_plan(1.0, 1);
+        let plan_to_cancel = plan_queue.add_plan(2.0, 2);
+        plan_queue.add_plan(3.0, 3);
+        plan_queue.cancel_plan(&plan_to_cancel);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 1.0);
+        assert_eq!(next_plan.data, 1);
+
         let next_plan = plan_queue.get_next_plan().unwrap();
         assert_eq!(next_plan.time, 3.0);
-        assert_eq!(next_plan.data, 4);
+        assert_eq!(next_plan.data, 3);
 
         assert!(plan_queue.get_next_plan().is_none());
     }
 
     #[test]
     #[should_panic]
-    fn test_invalid_cancel() {
-        // Cancel a plan that has already occured and make sure it panics
+    fn cancel_invalid_plan() {
         let mut plan_queue = Queue::<()>::new();
         let plan_to_cancel = plan_queue.add_plan(1.0, ());
         plan_queue.get_next_plan();

From 724667a1c5f919bcaca7c9868de7dd9e96829c4e Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Fri, 2 Aug 2024 18:46:39 -0400
Subject: [PATCH 15/17] Add more docs

---
 src/context.rs | 101 ++++++++++++++++++++++++++++++++++++++-----------
 src/lib.rs     |  28 ++++++++++++++
 src/plan.rs    |   3 +-
 3 files changed, 108 insertions(+), 24 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index e64a2a0..8e4cc2d 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -1,33 +1,35 @@
+//! A manager for the state of a discrete-event simulation
+//!
+//! Defines a `Context` that is intended to provide the foundational mechanism
+//! for storing and manipulating the state of a given simulation.
 use std::{
     any::{Any, TypeId},
     collections::{HashMap, VecDeque},
 };
 
-pub trait DataPlugin: Any {
-    type DataContainer;
-
-    fn create_data_container() -> Self::DataContainer;
-}
-
-#[macro_export]
-macro_rules! define_data_plugin {
-    ($plugin:ident, $data_container:ty, $default: expr) => {
-        struct $plugin {}
-
-        impl $crate::context::DataPlugin for $plugin {
-            type DataContainer = $data_container;
-
-            fn create_data_container() -> Self::DataContainer {
-                $default
-            }
-        }
-    };
-}
-pub use define_data_plugin;
-
 use crate::plan::{Id, Queue};
 
-type Callback = dyn FnOnce(&mut Context);
+/// A manager for the state of a discrete-event simulation
+///
+/// Provides core simulation services including
+/// * Maintaining a notion of time
+/// * Scheduling events to occur at some point in the future and executing them
+///   at that time
+/// * Holding data that can be accessed by simulation modules
+///
+/// Simulations are constructed out of a series of interacting modules that
+/// take turns manipulating the Context through a mutable reference. Modules
+/// store data in the simulation using the `DataPlugin` trait that allows them
+/// to retrieve data by type.
+///
+/// The future event list of the simulation is a queue of `Callback` objects -
+/// called `plans` - that will assume control of the Context at a future point
+/// in time and execute the logic in the associated `FnOnce(&mut Context)`
+/// closure. Modules can add plans to this queue through the `Context`.
+///
+/// The simulation also has an immediate callback queue to allow for the
+/// accumulation of side effects of mutations by the current controlling module.
+///
 pub struct Context {
     plan_queue: Queue<Box<Callback>>,
     callback_queue: VecDeque<Box<Callback>>,
@@ -36,6 +38,7 @@ pub struct Context {
 }
 
 impl Context {
+    /// Create a new empty `Context`
     #[must_use]
     pub fn new() -> Context {
         Context {
@@ -46,6 +49,10 @@ impl Context {
         }
     }
 
+    /// Add a plan to the future event list at the specified time
+    ///
+    /// Returns an `Id` for the newly-added plan that can be used to cancel it
+    /// if needed.
     /// # Panics
     ///
     /// Panics if time is in the past, infinite, or NaN.
@@ -54,14 +61,29 @@ impl Context {
         self.plan_queue.add_plan(time, Box::new(callback))
     }
 
+    /// Cancel a plan that has been added to the queue
+    ///
+    /// # Panics
+    ///
+    /// This function panics if you cancel a plan which has already been
+    /// cancelled or executed.
     pub fn cancel_plan(&mut self, id: &Id) {
         self.plan_queue.cancel_plan(id);
     }
 
+    /// Add a `Callback` to the queue to be executed before the next plan
     pub fn queue_callback(&mut self, callback: impl FnOnce(&mut Context) + 'static) {
         self.callback_queue.push_back(Box::new(callback));
     }
 
+    /// Retrieve a mutable reference to the data container associated with a
+    /// `DataPlugin`
+    ///
+    /// If the data container has not been already added to the `Context` then
+    /// this function will use the `DataPlugin::create_data_container` method
+    /// to construct a new data container and store it in the `Context`.
+    ///
+    /// Returns a mutable reference to the data container
     #[must_use]
     #[allow(clippy::missing_panics_doc)]
     pub fn get_data_container_mut<T: DataPlugin>(&mut self) -> &mut T::DataContainer {
@@ -72,6 +94,10 @@ impl Context {
             .unwrap() // Will never panic as data container has the matching type
     }
 
+    /// Retrieve a reference to the data container associated with a
+    /// `DataPlugin`
+    ///
+    /// Returns a reference to the data container if it exists or else `None`
     #[must_use]
     pub fn get_data_container<T: DataPlugin>(&self) -> Option<&T::DataContainer> {
         if let Some(data) = self.data_plugins.get(&TypeId::of::<T>()) {
@@ -81,11 +107,15 @@ impl Context {
         }
     }
 
+    /// Get the current time in the simulation
+    ///
+    /// Returns the current time
     #[must_use]
     pub fn get_current_time(&self) -> f64 {
         self.current_time
     }
 
+    /// Execute the simulation until the plan and callback queues are empty
     pub fn execute(&mut self) {
         // Start plan loop
         loop {
@@ -113,6 +143,31 @@ impl Default for Context {
     }
 }
 
+type Callback = dyn FnOnce(&mut Context);
+
+/// A trait for objects that can provide data containers to be held by `Context`
+pub trait DataPlugin: Any {
+    type DataContainer;
+
+    fn create_data_container() -> Self::DataContainer;
+}
+
+#[macro_export]
+macro_rules! define_data_plugin {
+    ($plugin:ident, $data_container:ty, $default: expr) => {
+        struct $plugin {}
+
+        impl $crate::context::DataPlugin for $plugin {
+            type DataContainer = $data_container;
+
+            fn create_data_container() -> Self::DataContainer {
+                $default
+            }
+        }
+    };
+}
+pub use define_data_plugin;
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/lib.rs b/src/lib.rs
index e81eede..dfe09e7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,2 +1,30 @@
+//! A framework for building discrete-event simulations
+//!
+//! Ixa is a framework designed to support the creation of large-scale
+//! discrete event simulations. The primary use case is the construction of
+//! agent-based models for disease transmission, but the approach is applicable
+//! in a wide array of circumstances.
+//!
+//! The central object of an Eosim simulation is the `Context` that is
+//! responsible for managing all the behavior of the simulation. All of the
+//! simulation-specific logic is embedded in modules that rely on the `Context`
+//! for core services such as:
+//! * Maintaining a notion of time for the simulation
+//! * Scheduling events to occur at some point in the future and executing them
+//!   at that time
+//! * Holding module-specific data so that the module and other modules can
+//!   access it
+//!
+//! In practice, a simulation usually consists of a set of modules that work
+//! together to provide all of the functions of the simulation. For instance,
+//! For instance, a simple disease transmission model might consist of the
+//! following modules:
+//! * A population loader that initializes the set of people represented
+//!   by the simulation.
+//! * An infection seeder that introduces the pathogen into the population.
+//! * A disease progression manager that transitions infected people through
+//!   stages of disease until recovery.
+//! * A transmission manager that models the process of an infected
+//!   person trying to infect susceptible people in the population.
 pub mod context;
 pub mod plan;
diff --git a/src/plan.rs b/src/plan.rs
index d2ddcab..678f5dd 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -31,7 +31,7 @@ pub struct Queue<T> {
 }
 
 impl<T> Queue<T> {
-    /// Creates a new empty `Queue<T>`
+    /// Create a new empty `Queue<T>`
     #[must_use]
     pub fn new() -> Queue<T> {
         Queue {
@@ -113,6 +113,7 @@ impl PartialOrd for Entry {
     }
 }
 
+/// Entry objects are ordered in increasing order by time and then plan id
 impl Ord for Entry {
     fn cmp(&self, other: &Self) -> Ordering {
         let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();

From 70a6bfb29bd88f66b1c7a1350f6b14fca7f80d10 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Fri, 2 Aug 2024 23:40:12 -0400
Subject: [PATCH 16/17] Address comments from review and add tests

---
 src/context.rs | 41 ++++++++++++++++++++++++++++++-----------
 src/plan.rs    | 27 +++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/src/context.rs b/src/context.rs
index 8e4cc2d..268ec5f 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -9,6 +9,9 @@ use std::{
 
 use crate::plan::{Id, Queue};
 
+/// The common callback used by multiple `Context` methods for future events
+type Callback = dyn FnOnce(&mut Context);
+
 /// A manager for the state of a discrete-event simulation
 ///
 /// Provides core simulation services including
@@ -27,8 +30,10 @@ use crate::plan::{Id, Queue};
 /// in time and execute the logic in the associated `FnOnce(&mut Context)`
 /// closure. Modules can add plans to this queue through the `Context`.
 ///
-/// The simulation also has an immediate callback queue to allow for the
-/// accumulation of side effects of mutations by the current controlling module.
+/// The simulation also has a separate callback mechanism. Callbacks
+/// fire before the next timed event (even if it is scheduled for the
+/// current time. This allows modules to schedule actions for immediate
+/// execution but outside of the current iteration of the event loop.
 ///
 pub struct Context {
     plan_queue: Queue<Box<Callback>>,
@@ -143,8 +148,6 @@ impl Default for Context {
     }
 }
 
-type Callback = dyn FnOnce(&mut Context);
-
 /// A trait for objects that can provide data containers to be held by `Context`
 pub trait DataPlugin: Any {
     type DataContainer;
@@ -174,6 +177,29 @@ mod tests {
 
     define_data_plugin!(ComponentA, Vec<u32>, vec![]);
 
+    #[test]
+    fn empty_context() {
+        let mut context = Context::new();
+        context.execute();
+        assert_eq!(context.get_current_time(), 0.0);
+    }
+
+    #[test]
+    fn get_data_container() {
+        let mut context = Context::new();
+        context.get_data_container_mut::<ComponentA>().push(1);
+        assert_eq!(
+            *context.get_data_container::<ComponentA>().unwrap(),
+            vec![1],
+        );
+    }
+
+    #[test]
+    fn get_uninitialized_data_container() {
+        let context = Context::new();
+        assert!(context.get_data_container::<ComponentA>().is_none());
+    }
+
     fn add_plan(context: &mut Context, time: f64, value: u32) -> Id {
         context.add_plan(time, move |context| {
             context.get_data_container_mut::<ComponentA>().push(value);
@@ -201,13 +227,6 @@ mod tests {
         add_plan(&mut context, f64::NAN, 0);
     }
 
-    #[test]
-    fn empty_context() {
-        let mut context = Context::new();
-        context.execute();
-        assert_eq!(context.get_current_time(), 0.0);
-    }
-
     #[test]
     fn timed_plan_only() {
         let mut context = Context::new();
diff --git a/src/plan.rs b/src/plan.rs
index 678f5dd..e29ec58 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -15,10 +15,10 @@ use std::{
 
 /// A priority queue that stores arbitrary data sorted by time
 ///
-/// Items of type T are stored in order by `f64` time and called `Plan<T>`.
+/// Items of type `T` are stored in order by `f64` time and called `Plan<T>`.
 /// When plans are created they are sequentially assigned an `Id` that is a
 /// wrapped `u64`. If two plans are scheduled for the same time the plan that is
-/// scheduled first (i.e. that has the lowest id) is placed earlier.
+/// scheduled first (i.e., that has the lowest id) is placed earlier.
 ///
 /// The pair of time and plan id are stored in a binary heap of `Entry` objects.
 /// The data payload of the event is stored in a hash map by plan id.
@@ -204,6 +204,29 @@ mod tests {
         assert!(plan_queue.get_next_plan().is_none());
     }
 
+    #[test]
+    fn add_and_get_plans() {
+        let mut plan_queue = Queue::new();
+        plan_queue.add_plan(1.0, 1);
+        plan_queue.add_plan(2.0, 2);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 1.0);
+        assert_eq!(next_plan.data, 1);
+
+        plan_queue.add_plan(3.0, 3);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 2.0);
+        assert_eq!(next_plan.data, 2);
+
+        let next_plan = plan_queue.get_next_plan().unwrap();
+        assert_eq!(next_plan.time, 3.0);
+        assert_eq!(next_plan.data, 3);
+
+        assert!(plan_queue.get_next_plan().is_none());
+    }
+
     #[test]
     #[should_panic]
     fn cancel_invalid_plan() {

From 327c410280b9ab3d2cb6eb4444bc07bf574cef31 Mon Sep 17 00:00:00 2001
From: Jason Asher <swo9@cdc.gov>
Date: Sat, 3 Aug 2024 23:32:10 -0400
Subject: [PATCH 17/17] Modify test based on review

---
 src/plan.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/plan.rs b/src/plan.rs
index e29ec58..724b994 100644
--- a/src/plan.rs
+++ b/src/plan.rs
@@ -214,15 +214,15 @@ mod tests {
         assert_eq!(next_plan.time, 1.0);
         assert_eq!(next_plan.data, 1);
 
-        plan_queue.add_plan(3.0, 3);
+        plan_queue.add_plan(1.5, 3);
 
         let next_plan = plan_queue.get_next_plan().unwrap();
-        assert_eq!(next_plan.time, 2.0);
-        assert_eq!(next_plan.data, 2);
+        assert_eq!(next_plan.time, 1.5);
+        assert_eq!(next_plan.data, 3);
 
         let next_plan = plan_queue.get_next_plan().unwrap();
-        assert_eq!(next_plan.time, 3.0);
-        assert_eq!(next_plan.data, 3);
+        assert_eq!(next_plan.time, 2.0);
+        assert_eq!(next_plan.data, 2);
 
         assert!(plan_queue.get_next_plan().is_none());
     }