From 43aacd23fa43dac176446dd3d91115a05eb5b701 Mon Sep 17 00:00:00 2001
From: Mingwei Samuel <mingwei.samuel@gmail.com>
Date: Thu, 9 Jan 2025 16:23:50 -0800
Subject: [PATCH] subgraph_loops wip

---
 dfir_rs/src/scheduled/graph.rs |  3 +-
 dfir_rs/src/util/slot_vec.rs   | 66 +++++++++++++++++++++++++++++++---
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/dfir_rs/src/scheduled/graph.rs b/dfir_rs/src/scheduled/graph.rs
index 48427754313..12d8a789be2 100644
--- a/dfir_rs/src/scheduled/graph.rs
+++ b/dfir_rs/src/scheduled/graph.rs
@@ -23,13 +23,14 @@ use super::state::StateHandle;
 use super::subgraph::Subgraph;
 use super::{HandoffId, HandoffTag, SubgraphId, SubgraphTag};
 use crate::scheduled::ticks::{TickDuration, TickInstant};
-use crate::util::slot_vec::SlotVec;
+use crate::util::slot_vec::{SecondarySlotVec, SlotVec};
 use crate::Never;
 
 /// A DFIR graph. Owns, schedules, and runs the compiled subgraphs.
 #[derive(Default)]
 pub struct Dfir<'a> {
     pub(super) subgraphs: SlotVec<SubgraphTag, SubgraphData<'a>>,
+    pub(super) subgraph_loop: SecondarySlotVec<SubgraphTag, LoopId>,
     pub(super) context: Context,
 
     handoffs: SlotVec<HandoffTag, HandoffData>,
diff --git a/dfir_rs/src/util/slot_vec.rs b/dfir_rs/src/util/slot_vec.rs
index cca950e182b..b8d5ea1f526 100644
--- a/dfir_rs/src/util/slot_vec.rs
+++ b/dfir_rs/src/util/slot_vec.rs
@@ -55,7 +55,7 @@ pub struct SlotVec<Tag: ?Sized, Val> {
     _phantom: PhantomData<Tag>,
 }
 impl<Tag: ?Sized, Val> SlotVec<Tag, Val> {
-    /// Creates a new SlotVec.
+    /// Creates a new `SlotVec`.
     pub fn new() -> Self {
         Self {
             slots: Vec::default(),
@@ -63,14 +63,14 @@ impl<Tag: ?Sized, Val> SlotVec<Tag, Val> {
         }
     }
 
-    /// Inserts a value into the SlotVec and returns the key.
+    /// Inserts a value into the `SlotVec` and returns the key.
     pub fn insert(&mut self, value: Val) -> Key<Tag> {
         let key = Key::from_raw(self.slots.len());
         self.slots.push(value);
         key
     }
 
-    /// Use the provided function to generate a value given the key and insert it into the SlotVec.
+    /// Use the provided function to generate a value given the key and insert it into the `SlotVec`.
     pub fn insert_with_key<F>(&mut self, func: F) -> Key<Tag>
     where
         F: FnOnce(Key<Tag>) -> Val,
@@ -90,12 +90,12 @@ impl<Tag: ?Sized, Val> SlotVec<Tag, Val> {
         self.slots.get_mut(key.index)
     }
 
-    /// Returns the number of elements in the SlotVec.
+    /// Returns the number of elements in the `SlotVec`.
     pub fn len(&self) -> usize {
         self.slots.len()
     }
 
-    /// Returns true if the SlotVec is empty.
+    /// Returns true if the `SlotVec` is empty.
     pub fn is_empty(&self) -> bool {
         self.slots.is_empty()
     }
@@ -117,3 +117,59 @@ impl<Key: ?Sized, Val> Default for SlotVec<Key, Val> {
         Self::new()
     }
 }
+
+/// A secondary map used to associated data with keys from elements in an existing [`SlotVec`].
+pub struct SecondarySlotVec<Tag: ?Sized, Val> {
+    slots: Vec<Option<Val>>,
+    _phantom: PhantomData<Tag>,
+}
+impl<Tag: ?Sized, Val> SecondarySlotVec<Tag, Val> {
+    /// Creates a new `SecondarySlotVec`.
+    pub fn new() -> Self {
+        Self {
+            slots: Vec::default(),
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Inserts a value into the `SecondarySlotVec` and returns the previous value associated with the key.
+    pub fn insert(&mut self, key: Key<Tag>, value: Val) -> Option<Val> {
+        if key.index >= self.slots.len() {
+            self.slots.resize_with(key.index + 1, || None);
+        }
+        self.slots[key.index].replace(value)
+    }
+
+    /// Removes a value associated with the key from the `SecondarySlotVec` and returns it.
+    pub fn remove(&mut self, key: Key<Tag>) -> Option<Val> {
+        // TODO(mingwei): Shrink the vector?
+        self.slots[key.index].take()
+    }
+
+    /// Returns a reference to the value associated with the key.
+    pub fn get(&self, key: Key<Tag>) -> Option<&Val> {
+        self.slots.get(key.index).and_then(|v| v.as_ref())
+    }
+
+    /// Returns a mutable reference to the value associated with the key.
+    pub fn get_mut(&mut self, key: Key<Tag>) -> Option<&mut Val> {
+        self.slots.get_mut(key.index).and_then(|v| v.as_mut())
+    }
+}
+impl<Tag: ?Sized, Val> Default for SecondarySlotVec<Tag, Val> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+impl<Tag: ?Sized, Val> Index<Key<Tag>> for SecondarySlotVec<Tag, Val> {
+    type Output = Val;
+
+    fn index(&self, key: Key<Tag>) -> &Self::Output {
+        self.get(key).unwrap()
+    }
+}
+impl<Tag: ?Sized, Val> IndexMut<Key<Tag>> for SecondarySlotVec<Tag, Val> {
+    fn index_mut(&mut self, key: Key<Tag>) -> &mut Self::Output {
+        self.get_mut(key).unwrap()
+    }
+}