perf(analyzer): track sparse SSA branch deltas

Author: tignear

crates/analyzer/src/comb_loop_detect/procedure.rs

@@ -1,7 +1,7 @@
 //! Analyzer-IR procedure evaluation for combinational dependency extraction.
 
 use super::region::{BitPartition, NodeKey, PackedSpan, dst_writes, var_reads};
-use super::ssa::{SsaStore, VersionId};
+use super::ssa::{BranchState, SsaStore, VersionId};
 use crate::HashSet;
 use crate::conv::Context;
 use crate::ir::VarId;
@@ -100,8 +100,9 @@ impl<'a> ProcedureAnalysis<'a> {
         if let Some((_, expression)) = &destination.select.1 {
             dependencies.extend(self.eval_expr(expression));
         }
-        for key in self.write_keys(destination) {
-            let version = self.ssa.definition(dependencies.clone());
+        let keys = self.write_keys(destination);
+        let version = self.ssa.definition(dependencies);
+        for key in keys {
             self.ssa.bind(key, version);
             self.written.insert(key);
         }
@@ -240,22 +241,15 @@ impl<'a> ProcedureAnalysis<'a> {
         let condition = self.eval_expr(&statement.cond);
         let mut nested_controls = controls.to_vec();
         nested_controls.extend_from_slice(&condition);
-        let saved_state = self.ssa.snapshot();
-        let saved_written = self.written.clone();
-
+        let checkpoint = self.ssa.checkpoint();
         self.eval_block(&statement.true_side, &nested_controls);
-        let true_state = self.ssa.snapshot();
-        let true_written = self.written.clone();
+        let true_state = self.ssa.capture_and_rollback(checkpoint);
 
-        self.ssa.restore(&saved_state);
-        self.written = saved_written;
+        let checkpoint = self.ssa.checkpoint();
         self.eval_block(&statement.false_side, &nested_controls);
-        let false_state = self.ssa.snapshot();
-        let false_written = self.written.clone();
+        let false_state = self.ssa.capture_and_rollback(checkpoint);
 
-        self.ssa.merge(&saved_state, &[true_state, false_state]);
-        self.written = true_written;
-        self.written.extend(false_written);
+        self.ssa.merge(&[true_state, false_state]);
     }
 
     fn eval_case(&mut self, statement: &CaseStatement, controls: &[VersionId]) {
@@ -275,24 +269,16 @@ impl<'a> ProcedureAnalysis<'a> {
         }
         let mut nested_controls = controls.to_vec();
         nested_controls.extend(condition);
-        let saved_state = self.ssa.snapshot();
-        let saved_written = self.written.clone();
         let mut states = Vec::with_capacity(statement.arms.len() + 1);
-        let mut written = saved_written.clone();
         for arm in &statement.arms {
-            self.ssa.restore(&saved_state);
-            self.written = saved_written.clone();
+            let checkpoint = self.ssa.checkpoint();
             self.eval_block(&arm.body, &nested_controls);
-            states.push(self.ssa.snapshot());
-            written.extend(self.written.iter().copied());
+            states.push(self.ssa.capture_and_rollback(checkpoint));
         }
-        self.ssa.restore(&saved_state);
-        self.written = saved_written;
+        let checkpoint = self.ssa.checkpoint();
         self.eval_block(&statement.default, &nested_controls);
-        states.push(self.ssa.snapshot());
-        written.extend(self.written.iter().copied());
-        self.ssa.merge(&saved_state, &states);
-        self.written = written;
+        states.push(self.ssa.capture_and_rollback(checkpoint));
+        self.ssa.merge(&states);
     }
 
     fn eval_for(&mut self, statement: &ForStatement, controls: &[VersionId]) {
@@ -329,15 +315,10 @@ impl<'a> ProcedureAnalysis<'a> {
         // Runtime loops have a zero-trip path. One symbolic body traversal is
         // enough to expose all explicit reads; the exit phi keeps LiveOnEntry
         // separate so retained state does not become a loop edge.
-        let saved_state = self.ssa.snapshot();
-        let saved_written = self.written.clone();
+        let checkpoint = self.ssa.checkpoint();
         self.eval_block(&statement.body, &range_controls);
-        let body_state = self.ssa.snapshot();
-        let body_written = self.written.clone();
-        self.ssa
-            .merge(&saved_state, &[saved_state.clone(), body_state]);
-        self.written = saved_written;
-        self.written.extend(body_written);
+        let body_state = self.ssa.capture_and_rollback(checkpoint);
+        self.ssa.merge(&[BranchState::unchanged(), body_state]);
     }
 
     fn eval_expr_requested(

crates/analyzer/src/comb_loop_detect/ssa.rs

@@ -12,15 +12,34 @@ enum Version<K> {
     Phi(Vec<VersionId>),
 }
 
-#[derive(Clone)]
-pub(super) struct Snapshot<K> {
-    current: HashMap<K, VersionId>,
+pub(super) struct Checkpoint {
+    undo_start: usize,
+    depth: usize,
+}
+
+pub(super) struct BranchState<K> {
+    bindings: HashMap<K, VersionId>,
+}
+
+impl<K> BranchState<K> {
+    pub(super) fn unchanged() -> Self {
+        Self {
+            bindings: HashMap::default(),
+        }
+    }
+}
+
+struct Undo<K> {
+    key: K,
+    previous: Option<VersionId>,
 }
 
 pub(super) struct SsaStore<K> {
     versions: Vec<Version<K>>,
     entries: HashMap<K, VersionId>,
     current: HashMap<K, VersionId>,
+    undo: Vec<Undo<K>>,
+    checkpoints: Vec<usize>,
 }
 
 impl<K> Default for SsaStore<K> {
@@ -29,6 +48,8 @@ impl<K> Default for SsaStore<K> {
             versions: Vec::new(),
             entries: HashMap::default(),
             current: HashMap::default(),
+            undo: Vec::new(),
+            checkpoints: Vec::new(),
         }
     }
 }
@@ -51,9 +72,7 @@ where
         if let Some(version) = self.current.get(&key) {
             *version
         } else {
-            let version = self.entry(key);
-            self.current.insert(key, version);
-            version
+            self.entry(key)
         }
     }
 
@@ -66,38 +85,65 @@ where
     }
 
     pub(super) fn bind(&mut self, key: K, version: VersionId) {
-        self.current.insert(key, version);
+        let previous = self.current.insert(key, version);
+        if !self.checkpoints.is_empty() {
+            self.undo.push(Undo { key, previous });
+        }
     }
 
-    pub(super) fn snapshot(&self) -> Snapshot<K> {
-        Snapshot {
-            current: self.current.clone(),
-        }
+    pub(super) fn checkpoint(&mut self) -> Checkpoint {
+        let checkpoint = Checkpoint {
+            undo_start: self.undo.len(),
+            depth: self.checkpoints.len(),
+        };
+        self.checkpoints.push(checkpoint.undo_start);
+        checkpoint
     }
 
-    pub(super) fn restore(&mut self, snapshot: &Snapshot<K>) {
-        self.current.clone_from(&snapshot.current);
+    pub(super) fn capture_and_rollback(&mut self, checkpoint: Checkpoint) -> BranchState<K> {
+        assert_eq!(checkpoint.depth + 1, self.checkpoints.len());
+        assert_eq!(self.checkpoints.pop(), Some(checkpoint.undo_start));
+
+        let mut bindings = HashMap::default();
+        for undo in &self.undo[checkpoint.undo_start..] {
+            let version = self
+                .current
+                .get(&undo.key)
+                .copied()
+                .expect("a branch binding must exist until rollback");
+            bindings.insert(undo.key, version);
+        }
+
+        while self.undo.len() > checkpoint.undo_start {
+            let undo = self.undo.pop().expect("undo length checked above");
+            if let Some(previous) = undo.previous {
+                self.current.insert(undo.key, previous);
+            } else {
+                self.current.remove(&undo.key);
+            }
+        }
+        bindings.retain(|key, version| self.current.get(key).copied() != Some(*version));
+        BranchState { bindings }
     }
 
-    pub(super) fn merge(&mut self, base: &Snapshot<K>, states: &[Snapshot<K>]) {
-        let mut keys: HashSet<K> = base.current.keys().copied().collect();
+    pub(super) fn merge(&mut self, states: &[BranchState<K>]) {
+        let mut keys = HashSet::default();
         for state in states {
-            keys.extend(state.current.keys().copied());
+            keys.extend(state.bindings.keys().copied());
         }
-        let mut merged = HashMap::default();
         for key in keys {
-            let fallback = base
+            let fallback = self
                 .current
                 .get(&key)
                 .copied()
                 .unwrap_or_else(|| self.entry(key));
             let inputs = states
                 .iter()
-                .map(|state| state.current.get(&key).copied().unwrap_or(fallback))
+                .map(|state| state.bindings.get(&key).copied().unwrap_or(fallback))
                 .collect();
-            merged.insert(key, self.phi(inputs));
+            let version = self.phi(inputs);
+            self.bind(key, version);
         }
-        self.current = merged;
     }
 
     pub(super) fn root_sources(&self, version: VersionId) -> HashSet<K> {
@@ -179,32 +225,77 @@ mod tests {
     #[test]
     fn retained_live_on_entry_is_not_a_combinational_read() {
         let mut ssa = SsaStore::default();
-        let base = ssa.snapshot();
         let retained = ssa.read("destination");
+        let checkpoint = ssa.checkpoint();
         let assigned = ssa.definition(Vec::new());
         ssa.bind("destination", assigned);
-        let branch = ssa.snapshot();
+        let branch = ssa.capture_and_rollback(checkpoint);
 
-        ssa.merge(&base, &[base.clone(), branch]);
+        ssa.merge(&[BranchState::unchanged(), branch]);
 
         let merged = ssa.read("destination");
         assert_ne!(merged, retained);
         assert!(ssa.root_sources(merged).is_empty());
     }
 
     #[test]
-    fn restore_discards_current_bindings_without_discarding_versions() {
+    fn rollback_discards_current_bindings_without_discarding_versions() {
         let mut ssa = SsaStore::default();
-        let base = ssa.snapshot();
+        let checkpoint = ssa.checkpoint();
         let source = ssa.read("source");
         let definition = ssa.definition(vec![source]);
         ssa.bind("destination", definition);
 
-        ssa.restore(&base);
+        let _ = ssa.capture_and_rollback(checkpoint);
         let restored = ssa.read("destination");
 
         let expected = ["source"].into_iter().collect::<HashSet<_>>();
         assert_eq!(ssa.root_sources(definition), expected);
         assert!(ssa.root_sources(restored).is_empty());
     }
+
+    #[test]
+    fn branch_state_contains_only_keys_changed_since_checkpoint() {
+        let mut ssa = SsaStore::default();
+        for key in 0..1_000 {
+            let version = ssa.definition(Vec::new());
+            ssa.bind(key, version);
+        }
+
+        let checkpoint = ssa.checkpoint();
+        let version = ssa.definition(Vec::new());
+        ssa.bind(500, version);
+        let branch = ssa.capture_and_rollback(checkpoint);
+
+        assert_eq!(branch.bindings.len(), 1);
+        assert_eq!(branch.bindings[&500], version);
+    }
+
+    #[test]
+    fn nested_rollback_preserves_the_outer_transaction() {
+        let mut ssa = SsaStore::default();
+        let base = ssa.definition(Vec::new());
+        ssa.bind("outer", base);
+
+        let outer_checkpoint = ssa.checkpoint();
+        let outer_definition = ssa.definition(Vec::new());
+        ssa.bind("outer", outer_definition);
+
+        let inner_checkpoint = ssa.checkpoint();
+        let inner_definition = ssa.definition(Vec::new());
+        ssa.bind("inner", inner_definition);
+        let inner_state = ssa.capture_and_rollback(inner_checkpoint);
+
+        assert_eq!(ssa.read("outer"), outer_definition);
+        assert_ne!(ssa.read("inner"), inner_definition);
+
+        ssa.merge(&[BranchState::unchanged(), inner_state]);
+        let merged_inner = ssa.read("inner");
+        let outer_state = ssa.capture_and_rollback(outer_checkpoint);
+
+        assert_eq!(ssa.read("outer"), base);
+        assert_ne!(ssa.read("inner"), merged_inner);
+        assert_eq!(outer_state.bindings["outer"], outer_definition);
+        assert_eq!(outer_state.bindings["inner"], merged_inner);
+    }
 }