Skip to main content

cairo_lang_lowering/analysis/
use_sites.rs

1//! Mutable use-site tracking for lowered IR variables.
2//!
3//! Tracks where each variable is used (statement inputs and block-end inputs).
4//! Updated incrementally as forwarding runs, so later checks see current use counts.
5//!
6//! Each `(variable, UseLocation)` pair stores a multiplicity: the number of input
7//! slots consuming the variable at that location (e.g. `felt252_add(v, v)` records `v`
8//! at that location with count 2). `use_count` sums these, so it is the total number
9//! of consuming slots. A rename rewrites every slot at a location at once, so moving a
10//! variable's uses to another (see `move_uses`) transfers the whole per-location count.
11
12use cairo_lang_utils::ordered_hash_map::{Entry, OrderedHashMap};
13
14use crate::analysis::UseLocation;
15use crate::{BlockEnd, Lowered, VariableId};
16
17/// Tracks use sites for each variable, indexed by variable arena index.
18pub struct UseSites {
19    sites: Vec<OrderedHashMap<UseLocation, usize>>,
20}
21
22impl UseSites {
23    /// Builds use-site information by scanning all statements and block ends.
24    pub fn analyze(lowered: &Lowered<'_>) -> Self {
25        let mut sites: Vec<OrderedHashMap<UseLocation, usize>> =
26            (0..lowered.variables.len()).map(|_| OrderedHashMap::default()).collect();
27        for (block_id, block) in lowered.blocks.iter() {
28            for (stmt_idx, stmt) in block.statements.iter().enumerate() {
29                let loc = UseLocation::Statement((block_id, stmt_idx));
30                for input in stmt.inputs() {
31                    *sites[input.var_id.index()].entry(loc).or_default() += 1;
32                }
33            }
34            let end_loc = UseLocation::BlockEnd(block_id);
35            match &block.end {
36                BlockEnd::Return(returns, _) => {
37                    for ret in returns {
38                        *sites[ret.var_id.index()].entry(end_loc).or_default() += 1;
39                    }
40                }
41                BlockEnd::Panic(var) => {
42                    *sites[var.var_id.index()].entry(end_loc).or_default() += 1;
43                }
44                BlockEnd::Goto(_, remapping) => {
45                    for (_, src) in remapping.iter() {
46                        *sites[src.var_id.index()].entry(end_loc).or_default() += 1;
47                    }
48                }
49                BlockEnd::Match { info } => {
50                    for input in info.inputs() {
51                        *sites[input.var_id.index()].entry(end_loc).or_default() += 1;
52                    }
53                }
54                BlockEnd::NotSet => {}
55            }
56        }
57        Self { sites }
58    }
59
60    /// Returns the total number of consuming slots remaining for `var` (the summed
61    /// multiplicity across all of its use-site locations).
62    pub fn use_count(&self, var: VariableId) -> usize {
63        self.sites[var.index()].iter().map(|(_, count)| *count).sum()
64    }
65
66    /// Moves all uses of `from` at `loc` to `to`, mirroring a rename that rewrites
67    /// every consuming slot at that location in one go: the whole per-location count is
68    /// transferred and added to any uses `to` already has there. No-op if `from` has no
69    /// use at `loc`.
70    pub fn move_uses(&mut self, from: VariableId, to: VariableId, loc: UseLocation) {
71        let Entry::Occupied(entry) = self.sites[from.index()].entry(loc) else {
72            return;
73        };
74        let count = entry.swap_remove();
75        *self.sites[to.index()].entry(loc).or_default() += count;
76    }
77
78    /// Returns each use-site location for `var` together with its multiplicity (the
79    /// number of consuming slots at that location).
80    pub fn use_locs(
81        &self,
82        var: VariableId,
83    ) -> impl ExactSizeIterator<Item = (UseLocation, usize)> + '_ {
84        self.sites[var.index()].iter().map(|(loc, count)| (*loc, *count))
85    }
86}
87
88impl std::fmt::Debug for UseSites {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        let mut list = f.debug_list();
91        for (idx, sites) in self.sites.iter().enumerate() {
92            let locs: Vec<_> = sites.iter().collect();
93            list.entry(&format_args!("v{idx}: {locs:?}"));
94        }
95        list.finish()
96    }
97}