rapx 0.7.65

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::PathGraph;

#[derive(Clone, Debug)]
pub struct CallBinding {
    pub arg_locals: Vec<usize>,
    pub arg_consts: Vec<(usize, CallConstArg)>,
    pub arg_field_offsets: Vec<usize>,
    pub dest_local: usize,
}

#[derive(Clone, Debug)]
pub enum CallConstArg {
    Literal(u64),
    Symbol(String),
}

/// 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).
/// - Reachability is filtered during standard enumeration; verification can
///   retain structural CFG paths and check feasibility in the symbolic VM.
#[derive(Debug, Clone)]
pub struct PathTree {
    root: Option<PathNode>,
    len: usize,
    /// Per-block function ownership for composed verification paths.
    block_fn: Vec<(DefId, usize)>,
    /// Argument and return bindings for callee entry blocks.
    call_bindings: FxHashMap<usize, CallBinding>,
    /// Parent function of each callee entry block.
    call_parents: FxHashMap<usize, DefId>,
    /// Caller blocks whose callee paths were composed into the tree.
    expanded_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,
    structural: 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(),
            call_bindings: FxHashMap::default(),
            call_parents: FxHashMap::default(),
            expanded_call_blocks: FxHashSet::default(),
            truncated: false,
            structural: false,
        }
    }

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

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

    pub fn is_structural(&self) -> bool {
        self.structural
    }

    pub fn mark_structural(&mut self) {
        self.structural = true;
    }

    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 function ownership and callsite metadata for composed paths.
    pub fn set_block_metadata(
        &mut self,
        map: Vec<(DefId, usize)>,
        bindings: FxHashMap<usize, CallBinding>,
        parents: FxHashMap<usize, DefId>,
        expanded_calls: FxHashSet<usize>,
    ) {
        self.block_fn = map;
        self.call_bindings = bindings;
        self.call_parents = parents;
        self.expanded_call_blocks = expanded_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 and return binding for a callee entry block.
    pub fn call_binding(&self, block: usize) -> Option<&CallBinding> {
        self.call_bindings.get(&block)
    }

    /// Parent function of a callee entry block.
    pub fn call_parent(&self, block: usize) -> Option<DefId> {
        self.call_parents.get(&block).copied()
    }

    /// Whether the call at `block` has callee paths in this tree.
    pub fn is_expanded_call(&self, block: usize) -> bool {
        self.expanded_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 `Some(())` if the walk completed, or `None` if `f` returned
    /// `false` to request early termination.
    pub fn walk_prefixes<F>(&self, target_block: usize, f: &mut F) -> Option<()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        let Some(root) = self.root.as_ref() else {
            return Some(());
        };
        let mut path = Vec::new();
        Self::walk_prefixes_impl(root, &mut path, target_block, false, f)
    }

    /// Build a subtree holding only the prefixes that end at `target_block`,
    /// preserving the block-function and inlining metadata so a sliced subtree
    /// still resolves inlined-callee blocks.
    pub fn subtree_to(&self, target_block: usize) -> PathTree {
        let mut sub = PathTree::new();
        sub.block_fn = self.block_fn.clone();
        sub.call_bindings = self.call_bindings.clone();
        sub.call_parents = self.call_parents.clone();
        sub.expanded_call_blocks = self.expanded_call_blocks.clone();
        sub.truncated = self.truncated;
        sub.structural = self.structural;
        let _ = self.walk_prefixes(target_block, &mut |prefix| {
            sub.insert(prefix);
            true
        });
        sub
    }

    /// 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) -> Option<()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        let Some(root) = self.root.as_ref() else {
            return Some(());
        };
        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,
    ) -> Option<()>
    where
        F: FnMut(&[usize]) -> bool,
    {
        path.push(node.block);
        if node.block == target_block {
            let cont = f(path);
            if !cont {
                path.pop();
                return None;
            }
            if !continue_past_target {
                path.pop();
                return Some(());
            }
        }
        for child in &node.children {
            Self::walk_prefixes_impl(child, path, target_block, continue_past_target, f)?;
        }
        path.pop();
        Some(())
    }
}

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(())
    }
}