(refactor): moving unsafe panic optimization into new dataflow format

Author: eytan-starkware

crates/cairo-lang-lowering/src/analysis/backward.rs

@@ -3,8 +3,8 @@
 
 use std::collections::HashMap;
 
-use crate::analysis::Analyzer;
-use crate::{BlockEnd, BlockId, Lowered};
+use crate::analysis::{Analyzer, DataflowAnalyzer, Direction, Edge, StatementLocation};
+use crate::{Block, BlockEnd, BlockId, Lowered, MatchInfo, Statement, VarRemapping, VarUsage};
 
 /// Main analysis type that allows traversing the flow backwards.
 pub struct BackAnalysis<'db, 'a, TAnalyzer: Analyzer<'db, 'a>> {
@@ -103,3 +103,118 @@ impl<'db, 'a, TAnalyzer: Analyzer<'db, 'a>> BackAnalysis<'db, 'a, TAnalyzer> {
         }
     }
 }
+
+/// Backward analysis runner using `DataflowAnalyzer`.
+///
+/// This is an adapter that wraps `BackAnalysis` internally, translating
+/// between the new `DataflowAnalyzer` trait and the legacy `Analyzer` trait.
+/// Once all analyses are migrated, this can be simplified to inline the
+/// traversal logic directly.
+pub struct DataflowBackAnalysis<'db, 'a, TAnalyzer: DataflowAnalyzer<'db, 'a>> {
+    inner: BackAnalysis<'db, 'a, AnalyzerAdapter<'db, 'a, TAnalyzer>>,
+}
+
+impl<'db, 'a, TAnalyzer: DataflowAnalyzer<'db, 'a>> DataflowBackAnalysis<'db, 'a, TAnalyzer> {
+    /// Creates a new DataflowBackAnalysis instance.
+    pub fn new(lowered: &'a Lowered<'db>, analyzer: &'a mut TAnalyzer) -> Self {
+        assert!(
+            TAnalyzer::DIRECTION == Direction::Backward,
+            "DataflowBackAnalysis requires a backward analyzer"
+        );
+        let adapter = AnalyzerAdapter { analyzer, lowered };
+        Self { inner: BackAnalysis::new(lowered, adapter) }
+    }
+
+    /// Runs the analysis and returns the result.
+    ///
+    /// For backward analysis, returns the info at the function entry (root block).
+    pub fn run(mut self) -> TAnalyzer::Info {
+        self.inner.get_root_info()
+    }
+}
+
+/// Adapter that implements the legacy `Analyzer` trait by delegating to `DataflowAnalyzer`.
+pub struct AnalyzerAdapter<'db, 'a, TAnalyzer: DataflowAnalyzer<'db, 'a>> {
+    pub analyzer: &'a mut TAnalyzer,
+    lowered: &'a Lowered<'db>,
+}
+
+impl<'db, 'a, TAnalyzer: DataflowAnalyzer<'db, 'a>> Analyzer<'db, 'a>
+    for AnalyzerAdapter<'db, 'a, TAnalyzer>
+{
+    type Info = TAnalyzer::Info;
+
+    fn visit_block_start(&mut self, info: &mut Self::Info, block_id: BlockId, _block: &Block<'db>) {
+        // Get block from lowered with correct lifetime 'a.
+        let block = &self.lowered.blocks[block_id];
+        // First apply transfer_block (which processes statements in reverse for backward).
+        self.analyzer.transfer_block(info, block_id, block);
+        // Then call the block start hook.
+        self.analyzer.visit_block_start(info, block_id, block);
+    }
+
+    fn visit_stmt(
+        &mut self,
+        _info: &mut Self::Info,
+        _statement_location: StatementLocation,
+        _stmt: &'a Statement<'db>,
+    ) {
+        // Statements are handled by transfer_block in visit_block_start.
+        // This is intentionally empty.
+    }
+
+    fn visit_goto(
+        &mut self,
+        info: &mut Self::Info,
+        _statement_location: StatementLocation,
+        target_block_id: BlockId,
+        remapping: &'a VarRemapping<'db>,
+    ) {
+        let edge = Edge::Goto { target: target_block_id, remapping };
+        *info = self.analyzer.transfer_edge(info, &edge);
+    }
+
+    fn merge_match(
+        &mut self,
+        statement_location: StatementLocation,
+        match_info: &'a MatchInfo<'db>,
+        infos: impl Iterator<Item = Self::Info>,
+    ) -> Self::Info {
+        // Transfer each arm's info through its edge, then merge.
+        let transferred_infos: Vec<_> = match_info
+            .arms()
+            .iter()
+            .zip(infos)
+            .map(|(arm, info)| {
+                let edge = Edge::MatchArm { arm, match_info };
+                let transferred = self.analyzer.transfer_edge(&info, &edge);
+                (arm.block_id, transferred)
+            })
+            .collect();
+        self.analyzer.merge(self.lowered, statement_location, transferred_infos.into_iter())
+    }
+
+    fn info_from_return(
+        &mut self,
+        statement_location: StatementLocation,
+        vars: &'a [VarUsage<'db>],
+    ) -> Self::Info {
+        let block_end = &self.lowered.blocks[statement_location.0].end;
+        let BlockEnd::Return(_, location) = block_end else {
+            unreachable!("Unexpected block end type")
+        };
+        let location = *location;
+        let info = self.analyzer.initial_info(statement_location.0, block_end);
+        self.analyzer.transfer_edge(&info, &Edge::Return { vars, location })
+    }
+
+    fn info_from_panic(
+        &mut self,
+        statement_location: StatementLocation,
+        var: &VarUsage<'db>,
+    ) -> Self::Info {
+        let block_end = &self.lowered.blocks[statement_location.0].end;
+        let info = self.analyzer.initial_info(statement_location.0, block_end);
+        self.analyzer.transfer_edge(&info, &Edge::Panic { var: *var })
+    }
+}

crates/cairo-lang-lowering/src/analysis/core.rs

@@ -25,7 +25,6 @@ use crate::{
 pub type StatementLocation = (BlockId, usize);
 
 /// The direction of dataflow analysis.
-#[expect(dead_code)]
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Direction {
     Forward,
@@ -37,7 +36,6 @@ pub enum Direction {
 /// Each variant captures the specific information needed for that edge type,
 /// enabling analyzers to handle variable introductions and remappings.
 #[derive(Debug)]
-#[expect(dead_code)]
 pub enum Edge<'db, 'a> {
     /// A goto edge with variable remapping.
     Goto { target: BlockId, remapping: &'a VarRemapping<'db> },
@@ -77,7 +75,6 @@ pub enum Edge<'db, 'a> {
 ///     }
 /// }
 /// ```
-#[expect(dead_code)]
 pub trait DataflowAnalyzer<'db, 'a> {
     /// The analysis state/info type.
     type Info: Clone;

crates/cairo-lang-lowering/src/analysis/mod.rs

@@ -3,10 +3,12 @@
 //! This module provides generic analysis frameworks that can be used by various
 //! optimization passes and semantic checks.
 
-mod backward;
-mod core;
+pub mod backward;
+pub mod core;
 
-pub use backward::BackAnalysis;
+pub use core::{DataflowAnalyzer, Direction, Edge};
+
+pub use backward::{BackAnalysis, DataflowBackAnalysis};
 
 use crate::{Block, BlockId, MatchInfo, Statement, VarRemapping, VarUsage};
 

crates/cairo-lang-lowering/src/optimizations/early_unsafe_panic.rs

@@ -7,14 +7,12 @@ use std::collections::HashSet;
 use cairo_lang_defs::ids::ExternFunctionId;
 use cairo_lang_filesystem::flag::FlagsGroup;
 use cairo_lang_semantic::helper::ModuleHelper;
-use itertools::zip_eq;
 use salsa::Database;
 
-use crate::analysis::{Analyzer, BackAnalysis, StatementLocation};
+use crate::analysis::core::StatementLocation;
+use crate::analysis::{DataflowAnalyzer, DataflowBackAnalysis, Direction};
 use crate::ids::{LocationId, SemanticFunctionIdEx};
-use crate::{
-    BlockEnd, BlockId, Lowered, MatchExternInfo, MatchInfo, Statement, StatementCall, VarUsage,
-};
+use crate::{BlockEnd, BlockId, Lowered, MatchExternInfo, MatchInfo, Statement, StatementCall};
 
 /// Adds an early unsafe_panic when we detect that `return` is unreachable from a certain point in
 /// the code. This step is needed to avoid issues with undroppable references in Sierra to CASM.
@@ -32,13 +30,14 @@ pub fn early_unsafe_panic<'db>(db: &'db dyn Database, lowered: &mut Lowered<'db>
         core.submodule("internal").extern_function_id("trace"),
     ]);
 
-    let ctx = UnsafePanicContext { db, fixes: vec![], libfuncs_with_sideffect };
-    let mut analysis = BackAnalysis::new(lowered, ctx);
-    let fixes = if let ReachableSideEffects::Unreachable(location) = analysis.get_root_info() {
-        vec![((BlockId::root(), 0), location)]
-    } else {
-        analysis.analyzer.fixes
-    };
+    let mut ctx = UnsafePanicContext { db, libfuncs_with_sideffect, fixes: Vec::new() };
+    let result = DataflowBackAnalysis::new(lowered, &mut ctx).run();
+    let UnsafePanicContext { mut fixes, .. } = ctx;
+
+    // If the entry point itself is unreachable, add a fix for it.
+    if let ReachableSideEffects::Unreachable(location) = result {
+        fixes.push(((BlockId::root(), 0), location));
+    }
 
     let panic_func_id = core.submodule("panics").function_id("unsafe_panic", vec![]).lowered(db);
     for ((block_id, statement_idx), location) in fixes {
@@ -60,7 +59,7 @@ pub struct UnsafePanicContext<'db> {
     db: &'db dyn Database,
 
     /// The list of blocks where we can insert unsafe_panic.
-    fixes: Vec<(StatementLocation, LocationId<'db>)>,
+    pub fixes: Vec<(StatementLocation, LocationId<'db>)>,
 
     /// libfuncs with side effects that we need to ignore.
     libfuncs_with_sideffect: HashSet<ExternFunctionId<'db>>,
@@ -84,57 +83,68 @@ impl<'db> UnsafePanicContext<'db> {
 }
 
 /// Can this state lead to a return or a statement with side effect.
-#[derive(Clone, Default, PartialEq, Debug)]
+#[derive(Clone, Copy, Default, PartialEq, Debug)]
 pub enum ReachableSideEffects<'db> {
     /// Some return statement or statement with side effect is reachable.
     #[default]
     Reachable,
     /// No return statement or statement with side effect is reachable.
-    /// holds the location of the closest match with no returning arms.
+    /// Holds the location of the closest match with no returning arms.
     Unreachable(LocationId<'db>),
 }
 
-impl<'db> Analyzer<'db, '_> for UnsafePanicContext<'db> {
+impl<'db, 'a> DataflowAnalyzer<'db, 'a> for UnsafePanicContext<'db> {
     type Info = ReachableSideEffects<'db>;
+    const DIRECTION: Direction = Direction::Backward;
 
-    fn visit_stmt(
+    fn transfer_stmt(
         &mut self,
         info: &mut Self::Info,
         statement_location: StatementLocation,
-        stmt: &Statement<'db>,
+        stmt: &'a Statement<'db>,
     ) {
         if self.has_side_effects(stmt)
-            && let ReachableSideEffects::Unreachable(locations) = *info
+            && let ReachableSideEffects::Unreachable(loc) = *info
         {
-            self.fixes.push((statement_location, locations));
-            *info = ReachableSideEffects::Reachable
+            self.fixes.push((statement_location, loc));
+            *info = ReachableSideEffects::Reachable;
         }
     }
 
-    fn merge_match(
+    fn merge(
         &mut self,
+        lowered: &Lowered<'db>,
         statement_location: StatementLocation,
-        match_info: &MatchInfo<'db>,
-        infos: impl Iterator<Item = Self::Info>,
+        infos: impl Iterator<Item = (BlockId, Self::Info)>,
     ) -> Self::Info {
-        let mut res = ReachableSideEffects::Unreachable(*match_info.location());
-        for (arm, info) in zip_eq(match_info.arms(), infos) {
-            match info {
+        let mut result = ReachableSideEffects::default();
+        let mut all_unreachable = true;
+
+        for (src, reachability) in infos {
+            match reachability {
                 ReachableSideEffects::Reachable => {
-                    res = ReachableSideEffects::Reachable;
+                    all_unreachable = false;
+                }
+                ReachableSideEffects::Unreachable(loc) => {
+                    // Fix at the entry of this unreachable branch.
+                    self.fixes.push(((src, 0), loc));
                 }
-                ReachableSideEffects::Unreachable(l) => self.fixes.push(((arm.block_id, 0), l)),
             }
         }
 
-        if let ReachableSideEffects::Unreachable(location) = res {
+        // All branches are unreachable (or there are no branches, e.g., match on `never`).
+        // Use the match location as the source of unreachability.
+        if all_unreachable {
+            // In a backward runner location is always available
+            let location = self.block_entry_location(lowered, statement_location.0);
+            result = ReachableSideEffects::Unreachable(location);
             self.fixes.push((statement_location, location));
         }
 
-        res
+        result
     }
 
-    fn info_from_return(&mut self, _: StatementLocation, _vars: &[VarUsage<'db>]) -> Self::Info {
-        ReachableSideEffects::Reachable
+    fn initial_info(&mut self, _block_id: BlockId, _block_end: &'a BlockEnd<'db>) -> Self::Info {
+        ReachableSideEffects::default()
     }
 }