rapx/analysis/path_analysis/
mod.rs

1pub mod default;
2pub mod graph;
3
4use crate::utils::source::get_fn_name_byid;
5use rustc_hir::def_id::DefId;
6use std::fmt::{self, Display};
7
8use crate::compat::FxHashMap;
9use graph::PathGraph;
10
11/// Format a path slice with cleanup-block annotations.
12///
13/// Cleanup blocks (MIR unwind/drop paths) are marked with a `*` suffix.
14/// Example: `[0, 1, 2*, 3]` where block 2 is a cleanup block.
15pub fn format_path_annotated(path: &[usize], graph: &PathGraph<'_>) -> String {
16    let blocks: Vec<String> = path
17        .iter()
18        .map(|&b| {
19            if graph.is_cleanup_block(b) {
20                format!("{}*", b)
21            } else {
22                b.to_string()
23            }
24        })
25        .collect();
26    format!("[{}]", blocks.join(", "))
27}
28
29/// A prefix-tree (trie) of whole-CFG paths sharing common prefixes.
30///
31/// Paths are sequences of MIR block indices.  The trie compresses shared
32/// prefixes — if two paths `[0,1,2,5]` and `[0,1,2,6]` are both inserted,
33/// blocks 0→1→2 are stored once and branch at block 2.
34///
35/// Each node is a [`PathNode`]; a node with `is_path_end == true` marks the
36/// end of a complete path that exists in the tree.  The `len` field tracks
37/// the number of complete paths stored.
38///
39/// # Invariants
40/// - All paths inserted into a given tree start with the same root block
41///   (by construction, block 0 — the CFG entry).
42/// - A path is only stored if it passed reachability filtering
43///   (see [`PathGraph::check_transition`]).
44#[derive(Debug, Clone)]
45pub struct PathTree {
46    root: Option<PathNode>,
47    len: usize,
48}
49
50/// A node in a [`PathTree`] trie.
51///
52/// `block` is the MIR block index for this node.  `children` holds
53/// successor blocks that appear after `block` in at least one stored
54/// path.  `is_path_end` is `true` when some path terminates at this
55/// node (i.e. this block is a CFG terminator for that path).
56#[derive(Debug, Clone)]
57pub struct PathNode {
58    pub block: usize,
59    pub children: Vec<PathNode>,
60    pub is_path_end: bool,
61}
62
63impl PathNode {
64    fn from_path(path: &[usize]) -> Self {
65        let mut node = PathNode {
66            block: path[0],
67            children: Vec::new(),
68            is_path_end: path.len() == 1,
69        };
70        if path.len() > 1 {
71            node.children.push(PathNode::from_path(&path[1..]));
72        }
73        node
74    }
75}
76
77impl PathTree {
78    pub fn new() -> Self {
79        PathTree { root: None, len: 0 }
80    }
81
82    pub fn len(&self) -> usize {
83        self.len
84    }
85
86    pub fn is_empty(&self) -> bool {
87        self.len == 0
88    }
89
90    pub fn root(&self) -> Option<&PathNode> {
91        self.root.as_ref()
92    }
93
94    /// Insert a path into the tree. Returns `true` if the path was
95    /// newly added (not already present as a terminal path).
96    pub fn insert(&mut self, path: &[usize]) -> bool {
97        if path.is_empty() {
98            return false;
99        }
100
101        match &mut self.root {
102            None => {
103                self.root = Some(PathNode::from_path(path));
104                self.len = 1;
105                true
106            }
107            Some(root) => {
108                if root.block != path[0] {
109                    return false;
110                }
111                if Self::insert_into(root, &path[1..]) {
112                    self.len += 1;
113                    true
114                } else {
115                    false
116                }
117            }
118        }
119    }
120
121    fn insert_into(node: &mut PathNode, suffix: &[usize]) -> bool {
122        if suffix.is_empty() {
123            if node.is_path_end {
124                return false;
125            }
126            node.is_path_end = true;
127            return true;
128        }
129
130        let target = suffix[0];
131        for child in &mut node.children {
132            if child.block == target {
133                return Self::insert_into(child, &suffix[1..]);
134            }
135        }
136
137        node.children.push(PathNode::from_path(suffix));
138        true
139    }
140
141    /// Check whether the given path exists as a complete path in the tree.
142    pub fn contains(&self, path: &[usize]) -> bool {
143        let mut node = match self.root.as_ref() {
144            Some(n) => n,
145            None => return false,
146        };
147        if node.block != path[0] {
148            return false;
149        }
150        for &block in &path[1..] {
151            node = match node.children.iter().find(|c| c.block == block) {
152                Some(n) => n,
153                None => return false,
154            };
155        }
156        node.is_path_end
157    }
158
159    /// Enumerate all paths as owned `Vec<usize>`.
160    pub fn iter(&self) -> PathTreeIter<'_> {
161        PathTreeIter {
162            stack: self
163                .root
164                .as_ref()
165                .map(|r| vec![(r, vec![r.block])])
166                .unwrap_or_default(),
167        }
168    }
169
170    /// Collect all paths into a flat `Vec<Vec<usize>>`.
171    pub fn to_vecs(&self) -> Vec<Vec<usize>> {
172        self.iter().collect()
173    }
174
175    /// Walk the tree and call `f` with each unique prefix that ends at
176    /// `target_block`. The walk stops at `target_block` (does not recurse
177    /// into its children), so the callback receives the path from the root
178    /// up to and including `target_block`.
179    ///
180    /// Returns `Ok(())` if the walk completed, or `Err(())` if `f` returned
181    /// `false` to request early termination.
182    pub fn walk_prefixes<F>(&self, target_block: usize, f: &mut F) -> Result<(), ()>
183    where
184        F: FnMut(&[usize]) -> bool,
185    {
186        let Some(root) = self.root.as_ref() else {
187            return Ok(());
188        };
189        let mut path = Vec::new();
190        Self::walk_prefixes_impl(root, &mut path, target_block, false, f)
191    }
192
193    /// Like [`walk_prefixes`] but continues past the target block into
194    /// children, finding ALL occurrences (e.g. multiple iterations of the
195    /// same checkpoint block in a loop).
196    pub fn walk_all_prefixes<F>(&self, target_block: usize, f: &mut F) -> Result<(), ()>
197    where
198        F: FnMut(&[usize]) -> bool,
199    {
200        let Some(root) = self.root.as_ref() else {
201            return Ok(());
202        };
203        let mut path = Vec::new();
204        Self::walk_prefixes_impl(root, &mut path, target_block, true, f)
205    }
206
207    fn walk_prefixes_impl<F>(
208        node: &PathNode,
209        path: &mut Vec<usize>,
210        target_block: usize,
211        continue_past_target: bool,
212        f: &mut F,
213    ) -> Result<(), ()>
214    where
215        F: FnMut(&[usize]) -> bool,
216    {
217        path.push(node.block);
218        if node.block == target_block {
219            let cont = f(path);
220            if !cont {
221                path.pop();
222                return Err(());
223            }
224            if !continue_past_target {
225                path.pop();
226                return Ok(());
227            }
228        }
229        for child in &node.children {
230            Self::walk_prefixes_impl(child, path, target_block, continue_past_target, f)?;
231        }
232        path.pop();
233        Ok(())
234    }
235}
236
237impl Default for PathTree {
238    fn default() -> Self {
239        Self::new()
240    }
241}
242
243/// DFS iterator over all complete paths in a [`PathTree`].
244///
245/// Yields each path as an owned `Vec<usize>`.  Internal nodes (where
246/// `is_path_end == false`) are skipped; only terminal path nodes are
247/// emitted.
248pub struct PathTreeIter<'a> {
249    stack: Vec<(&'a PathNode, Vec<usize>)>,
250}
251
252impl<'a> Iterator for PathTreeIter<'a> {
253    type Item = Vec<usize>;
254
255    fn next(&mut self) -> Option<Self::Item> {
256        loop {
257            let (node, path) = self.stack.pop()?;
258            for child in node.children.iter().rev() {
259                let mut child_path = path.clone();
260                child_path.push(child.block);
261                self.stack.push((child, child_path));
262            }
263            if node.is_path_end {
264                return Some(path);
265            }
266        }
267    }
268}
269
270/// Display wrapper that prints all paths for every function, annotated
271/// with cleanup-block markers via [`format_path_annotated`].
272pub struct PathMapWrapper<'a, 'tcx>(
273    pub FxHashMap<DefId, PathTree>,
274    pub &'a FxHashMap<DefId, PathGraph<'tcx>>,
275);
276
277impl Display for PathMapWrapper<'_, '_> {
278    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
279        writeln!(f, "=== Print path analysis results ===")?;
280        for (def_id, paths) in &self.0 {
281            let fn_name = get_fn_name_byid(def_id);
282            if fn_name.contains("__raw_ptr_deref_dummy") {
283                continue;
284            }
285            writeln!(f, "Function: {:?}:", fn_name)?;
286            let graph = self.1.get(def_id);
287            for path in paths.iter() {
288                if let Some(g) = graph {
289                    writeln!(f, "  Path {}", format_path_annotated(&path, g))?;
290                } else {
291                    writeln!(f, "  Path {:?}", path)?;
292                }
293            }
294        }
295        Ok(())
296    }
297}
rapx/analysis/path_analysis/mod.rs

rapx/analysis/path_analysis/
mod.rs
