rapx/analysis/path/

mod.rs

pub mod default;
pub mod graph;

use crate::utils::source::get_fn_name_byid;
use rustc_hir::def_id::DefId;
use std::fmt::{self, Display};

use crate::compat::{FxHashMap, FxHashSet};
use graph::{InlineBinding, PathGraph};

/// Format a path slice with cleanup-block annotations.
///
/// Cleanup blocks (MIR unwind/drop paths) are marked with a `*` suffix.
/// Example: `[0, 1, 2*, 3]` where block 2 is a cleanup block.
pub fn format_path_annotated(path: &[usize], graph: &PathGraph<'_>) -> String {
    let blocks: Vec<String> = path
        .iter()
        .map(|&b| {
            if graph.is_cleanup_block(b) {
                format!("{}*", b)
            } else {
                b.to_string()
            }
        })
        .collect();
    format!("[{}]", blocks.join(", "))
}

/// A prefix-tree (trie) of whole-CFG paths sharing common prefixes.
///
/// Paths are sequences of MIR block indices.  The trie compresses shared
/// prefixes — if two paths `[0,1,2,5]` and `[0,1,2,6]` are both inserted,
/// blocks 0→1→2 are stored once and branch at block 2.
///
/// Each node is a [`PathNode`]; a node with `is_path_end == true` marks the
/// end of a complete path that exists in the tree.  The `len` field tracks
/// the number of complete paths stored.
///
/// # Invariants
/// - All paths inserted into a given tree start with the same root block
///   (by construction, block 0 — the CFG entry).
/// - A path is only stored if it passed reachability filtering
///   (see [`PathGraph::check_transition`]).
#[derive(Debug, Clone)]
pub struct PathTree {
    root: Option<PathNode>,
    len: usize,
    /// Per-global-block function ownership: `(def_id, local_index)`. Empty for
    /// single-function trees; populated after CFG inlining so consumers can
    /// resolve each block's owning function and its MIR block index.
    block_fn: Vec<(DefId, usize)>,
    /// Argument/return bindings for inlined callee entry blocks.
    inline_bindings: FxHashMap<usize, InlineBinding>,
    /// Per-inlined-callee parent: entry block (global index) → the `DefId` of
    /// the function that called it. Lets the boundary injector distinguish a
    /// *nested* callee (its body is split around a further-inlined callee) from
    /// a *sibling* call when both appear as consecutive non-caller def_ids.
    inline_parents: FxHashMap<usize, DefId>,
    /// Caller blocks whose `Call` terminator was inlined.
    inlined_call_blocks: FxHashSet<usize>,
    /// Set when enumeration stopped at a path or depth limit, so the tree
    /// holds only some of the paths.
    truncated: bool,
}

/// A node in a [`PathTree`] trie.
///
/// `block` is the MIR block index for this node.  `children` holds
/// successor blocks that appear after `block` in at least one stored
/// path.  `is_path_end` is `true` when some path terminates at this
/// node (i.e. this block is a CFG terminator for that path).
#[derive(Debug, Clone)]
pub struct PathNode {
    pub block: usize,
    pub children: Vec<PathNode>,
    pub is_path_end: bool,
}

impl PathNode {
    fn from_path(path: &[usize]) -> Self {
        let mut node = PathNode {
            block: path[0],
            children: Vec::new(),
            is_path_end: path.len() == 1,
        };
        if path.len() > 1 {
            node.children.push(PathNode::from_path(&path[1..]));
        }
        node
    }
}

impl PathTree {
    pub fn new() -> Self {
        PathTree {
            root: None,
            len: 0,
            block_fn: Vec::new(),
            inline_bindings: FxHashMap::default(),
            inline_parents: FxHashMap::default(),
            inlined_call_blocks: FxHashSet::default(),
            truncated: false,
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_truncated(&self) -> bool {
        self.truncated
    }

    pub fn mark_truncated(&mut self) {
        self.truncated = true;
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn root(&self) -> Option<&PathNode> {
        self.root.as_ref()
    }

    /// Set the per-global-block function ownership map `(def_id, local_index)`
    /// and the inlined callee argument bindings.
    pub fn set_block_fn(
        &mut self,
        map: Vec<(DefId, usize)>,
        bindings: FxHashMap<usize, InlineBinding>,
        parents: FxHashMap<usize, DefId>,
        inlined_calls: FxHashSet<usize>,
    ) {
        self.block_fn = map;
        self.inline_bindings = bindings;
        self.inline_parents = parents;
        self.inlined_call_blocks = inlined_calls;
    }

    /// Resolve `block` to `(def_id, local_index)`, or `None` for single-function
    /// trees (where `block` is already a local index of the single function).
    pub fn block_fn_of(&self, block: usize) -> Option<(DefId, usize)> {
        self.block_fn.get(block).copied()
    }

    /// All `(def_id, local_index)` entries, indexed by global block number.
    pub fn block_fns(&self) -> &[(DefId, usize)] {
        &self.block_fn
    }

    /// Argument/return binding for an inlined callee entry block.
    pub fn inline_binding(&self, block: usize) -> Option<&InlineBinding> {
        self.inline_bindings.get(&block)
    }

    /// The `DefId` of the function that called an inlined callee entry block.
    pub fn inline_parent(&self, block: usize) -> Option<DefId> {
        self.inline_parents.get(&block).copied()
    }

    /// Whether `block` (a caller block) had its `Call` terminator inlined.
    pub fn is_inlined_call(&self, block: usize) -> bool {
        self.inlined_call_blocks.contains(&block)
    }

    /// Insert a path into the tree. Returns `true` if the path was
    /// newly added (not already present as a terminal path).
    pub fn insert(&mut self, path: &[usize]) -> bool {
        if path.is_empty() {
            return false;
        }

        match &mut self.root {
            None => {
                self.root = Some(PathNode::from_path(path));
                self.len = 1;
                true
            }
            Some(root) => {
                if root.block != path[0] {
                    return false;
                }
                if Self::insert_into(root, &path[1..]) {
                    self.len += 1;
                    true
                } else {
                    false
                }
            }
        }
    }

    fn insert_into(node: &mut PathNode, suffix: &[usize]) -> bool {
        if suffix.is_empty() {
            if node.is_path_end {
                return false;
            }
            node.is_path_end = true;
            return true;
        }

        let target = suffix[0];
        for child in &mut node.children {
            if child.block == target {
                return Self::insert_into(child, &suffix[1..]);
            }
        }

        node.children.push(PathNode::from_path(suffix));
        true
    }

    /// Check whether the given path exists as a complete path in the tree.
    pub fn contains(&self, path: &[usize]) -> bool {
        let mut node = match self.root.as_ref() {
            Some(n) => n,
            None => return false,
        };
        if node.block != path[0] {
            return false;
        }
        for &block in &path[1..] {
            node = match node.children.iter().find(|c| c.block == block) {
                Some(n) => n,
                None => return false,
            };
        }
        node.is_path_end
    }

    /// Enumerate all paths as owned `Vec<usize>`.
    pub fn iter(&self) -> PathTreeIter<'_> {
        PathTreeIter {
            stack: self
                .root
                .as_ref()
                .map(|r| vec![(r, vec![r.block])])
                .unwrap_or_default(),
        }
    }

    /// Collect all paths into a flat `Vec<Vec<usize>>`.
    pub fn to_vecs(&self) -> Vec<Vec<usize>> {
        self.iter().collect()
    }

    /// Walk the tree and call `f` with each unique prefix that ends at
    /// `target_block`. The walk stops at `target_block` (does not recurse
    /// into its children), so the callback receives the path from the root
    /// up to and including `target_block`.
    ///
    /// Returns `Ok(())` if the walk completed, or `Err(())` if `f` returned
    /// `false` to request early termination.
    pub fn walk_prefixes<F>(&self, target_block: usize, f: &mut F) -> Result<(), ()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        let Some(root) = self.root.as_ref() else {
            return Ok(());
        };
        let mut path = Vec::new();
        Self::walk_prefixes_impl(root, &mut path, target_block, false, f)
    }

    /// Like [`walk_prefixes`] but continues past the target block into
    /// children, finding ALL occurrences (e.g. multiple iterations of the
    /// same checkpoint block in a loop).
    pub fn walk_all_prefixes<F>(&self, target_block: usize, f: &mut F) -> Result<(), ()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        let Some(root) = self.root.as_ref() else {
            return Ok(());
        };
        let mut path = Vec::new();
        Self::walk_prefixes_impl(root, &mut path, target_block, true, f)
    }

    fn walk_prefixes_impl<F>(
        node: &PathNode,
        path: &mut Vec<usize>,
        target_block: usize,
        continue_past_target: bool,
        f: &mut F,
    ) -> Result<(), ()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        path.push(node.block);
        if node.block == target_block {
            let cont = f(path);
            if !cont {
                path.pop();
                return Err(());
            }
            if !continue_past_target {
                path.pop();
                return Ok(());
            }
        }
        for child in &node.children {
            Self::walk_prefixes_impl(child, path, target_block, continue_past_target, f)?;
        }
        path.pop();
        Ok(())
    }
}

impl Default for PathTree {
    fn default() -> Self {
        Self::new()
    }
}

/// DFS iterator over all complete paths in a [`PathTree`].
///
/// Yields each path as an owned `Vec<usize>`.  Internal nodes (where
/// `is_path_end == false`) are skipped; only terminal path nodes are
/// emitted.
pub struct PathTreeIter<'a> {
    stack: Vec<(&'a PathNode, Vec<usize>)>,
}

impl<'a> Iterator for PathTreeIter<'a> {
    type Item = Vec<usize>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let (node, path) = self.stack.pop()?;
            for child in node.children.iter().rev() {
                let mut child_path = path.clone();
                child_path.push(child.block);
                self.stack.push((child, child_path));
            }
            if node.is_path_end {
                return Some(path);
            }
        }
    }
}

/// Display wrapper that prints all paths for every function, annotated
/// with cleanup-block markers via [`format_path_annotated`].
pub struct PathMapWrapper<'a, 'tcx>(
    pub FxHashMap<DefId, PathTree>,
    pub &'a FxHashMap<DefId, PathGraph<'tcx>>,
);

impl Display for PathMapWrapper<'_, '_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "=== Print path analysis results ===")?;
        for (def_id, paths) in &self.0 {
            let fn_name = get_fn_name_byid(def_id);
            if fn_name.contains("__raw_ptr_deref_dummy") {
                continue;
            }
            writeln!(f, "Function: {:?}:", fn_name)?;
            let graph = self.1.get(def_id);
            for path in paths.iter() {
                if let Some(g) = graph {
                    writeln!(f, "  Path {}", format_path_annotated(&path, g))?;
                } else {
                    writeln!(f, "  Path {:?}", path)?;
                }
            }
        }
        Ok(())
    }
}