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

eytan-starkware · eytan-starkware · commit e1163346695e · 2026-01-25T13:39:45.000+02:00
diff --git a/crates/cairo-lang-lowering/src/analysis/backward.rs b/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, 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,100 @@ 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: 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: 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>,
+    ) {
+        self.analyzer.apply_remapping(info, statement_location, target_block_id, remapping);
+    }
+
+    fn merge_match(
+        &mut self,
+        statement_location: StatementLocation,
+        match_info: &'a MatchInfo<'db>,
+        infos: impl Iterator<Item = Self::Info>,
+    ) -> Self::Info {
+        self.analyzer.merge_match(statement_location, match_info, infos)
+    }
+
+    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;
+        self.analyzer.initial_info(statement_location.0, block_end)
+    }
+
+    fn info_from_panic(
+        &mut self,
+        statement_location: StatementLocation,
+        _var: &VarUsage<'db>,
+    ) -> Self::Info {
+        let block_end = &self.lowered.blocks[statement_location.0].end;
+        self.analyzer.initial_info(statement_location.0, block_end)
+    }
+}
diff --git a/crates/cairo-lang-lowering/src/analysis/mod.rs b/crates/cairo-lang-lowering/src/analysis/mod.rs
@@ -6,10 +6,9 @@
 pub mod backward;
 pub mod core;
 
-// Re-export commonly used types at the module level for convenience.
 pub use core::{DataflowAnalyzer, Direction};
 
-pub use backward::BackAnalysis;
+pub use backward::{BackAnalysis, DataflowBackAnalysis};
 
 use crate::{Block, BlockId, MatchInfo, Statement, VarRemapping, VarUsage};
 
diff --git a/crates/cairo-lang-lowering/src/optimizations/early_unsafe_panic.rs b/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,16 +30,17 @@ 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 ctx = UnsafePanicContext { db, libfuncs_with_sideffect };
+    let analysis = DataflowBackAnalysis::new(lowered, ctx);
+    let mut result = analysis.run();
+
+    // If the entry point itself is unreachable, add a fix for it.
+    if let Reachability::Unreachable(location) = result.reachability {
+        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 {
+    for ((block_id, statement_idx), location) in result.fixes {
         let block = &mut lowered.blocks[block_id];
         block.statements.truncate(statement_idx);
 
@@ -59,9 +58,6 @@ pub fn early_unsafe_panic<'db>(db: &'db dyn Database, lowered: &mut Lowered<'db>
 pub struct UnsafePanicContext<'db> {
     db: &'db dyn Database,
 
-    /// The list of blocks where we can insert unsafe_panic.
-    fixes: Vec<(StatementLocation, LocationId<'db>)>,
-
     /// libfuncs with side effects that we need to ignore.
     libfuncs_with_sideffect: HashSet<ExternFunctionId<'db>>,
 }
@@ -83,58 +79,104 @@ impl<'db> UnsafePanicContext<'db> {
     }
 }
 
-/// Can this state lead to a return or a statement with side effect.
-#[derive(Clone, Default, PartialEq, Debug)]
-pub enum ReachableSideEffects<'db> {
+/// Reachability state for a point in the program.
+#[derive(Clone, Copy, Default, PartialEq, Debug)]
+pub enum Reachability<'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> {
-    type Info = ReachableSideEffects<'db>;
+/// Analysis info containing reachability state and accumulated fixes.
+#[derive(Clone, Default, Debug)]
+pub struct AnalysisInfo<'db> {
+    /// The reachability state at this program point.
+    pub reachability: Reachability<'db>,
+    /// Locations where we need to insert unsafe_panic.
+    pub fixes: Vec<(StatementLocation, LocationId<'db>)>,
+}
+
+impl<'db, 'a> DataflowAnalyzer<'db, 'a> for UnsafePanicContext<'db> {
+    type Info = AnalysisInfo<'db>;
+    const DIRECTION: Direction = Direction::Backward;
+
+    fn initial_info(&mut self, _block_id: BlockId, _block_end: &'a BlockEnd<'db>) -> Self::Info {
+        AnalysisInfo::default()
+    }
 
-    fn visit_stmt(
+    fn merge(
         &mut self,
-        info: &mut Self::Info,
-        statement_location: StatementLocation,
-        stmt: &Statement<'db>,
-    ) {
-        if self.has_side_effects(stmt)
-            && let ReachableSideEffects::Unreachable(locations) = *info
-        {
-            self.fixes.push((statement_location, locations));
-            *info = ReachableSideEffects::Reachable
+        _statement_location: StatementLocation,
+        infos: impl Iterator<Item = (BlockId, Self::Info)>,
+    ) -> Self::Info {
+        let mut result = AnalysisInfo::default();
+        let mut all_unreachable = true;
+        let mut unreachable_location = None;
+
+        for (src, info) in infos {
+            result.fixes.extend(info.fixes);
+            if let Reachability::Unreachable(loc) = info.reachability {
+                // Fix at the entry of this unreachable branch.
+                result.fixes.push(((src, 0), loc));
+                unreachable_location.get_or_insert(loc);
+            } else {
+                all_unreachable = false;
+            }
+        }
+
+        if all_unreachable && let Some(loc) = unreachable_location {
+            result.reachability = Reachability::Unreachable(loc);
         }
+
+        result
     }
 
     fn merge_match(
         &mut self,
         statement_location: StatementLocation,
-        match_info: &MatchInfo<'db>,
+        match_info: &'a MatchInfo<'db>,
         infos: impl Iterator<Item = Self::Info>,
     ) -> Self::Info {
-        let mut res = ReachableSideEffects::Unreachable(*match_info.location());
-        for (arm, info) in zip_eq(match_info.arms(), infos) {
-            match info {
-                ReachableSideEffects::Reachable => {
-                    res = ReachableSideEffects::Reachable;
+        let mut result = AnalysisInfo::default();
+        let mut all_unreachable = true;
+
+        for (arm, info) in match_info.arms().iter().zip(infos) {
+            result.fixes.extend(info.fixes);
+            match info.reachability {
+                Reachability::Reachable => {
+                    all_unreachable = false;
+                }
+                Reachability::Unreachable(loc) => {
+                    // Fix at the entry of this unreachable arm.
+                    result.fixes.push(((arm.block_id, 0), loc));
                 }
-                ReachableSideEffects::Unreachable(l) => self.fixes.push(((arm.block_id, 0), l)),
             }
         }
 
-        if let ReachableSideEffects::Unreachable(location) = res {
-            self.fixes.push((statement_location, location));
+        if all_unreachable {
+            let loc = *match_info.location();
+            result.reachability = Reachability::Unreachable(loc);
+            // Fix at the match statement itself.
+            result.fixes.push((statement_location, loc));
         }
 
-        res
+        result
     }
 
-    fn info_from_return(&mut self, _: StatementLocation, _vars: &[VarUsage<'db>]) -> Self::Info {
-        ReachableSideEffects::Reachable
+    fn transfer_stmt(
+        &mut self,
+        info: &mut Self::Info,
+        statement_location: StatementLocation,
+        stmt: &'a Statement<'db>,
+    ) {
+        if self.has_side_effects(stmt)
+            && let Reachability::Unreachable(loc) = info.reachability
+        {
+            info.fixes.push((statement_location, loc));
+            info.reachability = Reachability::Reachable;
+        }
     }
 }
