unluac 1.3.0

Multi-dialect Lua decompiler written in Rust.
Documentation
//! CFG 构图层的稳定类型与基础查询。
//!
//! 这一层只表达 basic block、edge 和可达性等稳定图结构,不夹带 GraphFacts /
//! Dataflow 的派生语义。后续如果 StructureFacts/HIR 需要稳定 block/edge 查询,
//! 也应优先补在这里,而不是再在后层各自包一层小型 CFG API。

use std::collections::{BTreeSet, VecDeque};
use std::fmt;

use crate::transformer::{InstrRef, LowInstr};

/// 一个 proto 的控制流图,以及它的子 proto 图。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CfgGraph {
    pub cfg: Cfg,
    pub children: Vec<CfgGraph>,
}

/// 单个 proto 的基础控制流图。
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct Cfg {
    pub blocks: Vec<BasicBlock>,
    pub edges: Vec<CfgEdge>,
    pub entry_block: BlockRef,
    pub exit_block: BlockRef,
    pub block_order: Vec<BlockRef>,
    pub instr_to_block: Vec<BlockRef>,
    pub preds: Vec<Vec<EdgeRef>>,
    pub succs: Vec<Vec<EdgeRef>>,
    pub reachable_blocks: BTreeSet<BlockRef>,
}

/// 去重后的 reachable successor 数量形态。
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum ReachableSuccessorShape {
    Empty,
    Single(BlockRef),
    Multiple,
}

impl ReachableSuccessorShape {
    pub const fn unique(self) -> Option<BlockRef> {
        match self {
            Self::Single(block) => Some(block),
            Self::Empty | Self::Multiple => None,
        }
    }
}

/// 边的稳定引用。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
pub struct EdgeRef(pub usize);

impl EdgeRef {
    pub const fn index(self) -> usize {
        self.0
    }
}

impl fmt::Display for EdgeRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "#{}", self.0)
    }
}

/// block 的稳定引用。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
pub struct BlockRef(pub usize);

impl BlockRef {
    pub const fn index(self) -> usize {
        self.0
    }
}

impl fmt::Display for BlockRef {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "#{}", self.0)
    }
}

/// 一个 basic block。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Default)]
pub struct BasicBlock {
    pub kind: BlockKind,
    pub instrs: InstrRange,
}

/// block 的类别。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, Default)]
pub enum BlockKind {
    #[default]
    Normal,
    SyntheticExit,
}

/// 指令线性区间。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct InstrRange {
    pub start: InstrRef,
    pub len: usize,
}

impl InstrRange {
    pub const fn new(start: InstrRef, len: usize) -> Self {
        Self { start, len }
    }

    pub const fn end(self) -> usize {
        self.start.index() + self.len
    }

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

    pub const fn last(self) -> Option<InstrRef> {
        if self.len == 0 {
            None
        } else {
            Some(InstrRef(self.start.index() + self.len - 1))
        }
    }
}

impl Default for InstrRange {
    fn default() -> Self {
        Self::new(InstrRef(0), 0)
    }
}

/// CFG 边。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct CfgEdge {
    pub from: BlockRef,
    pub to: BlockRef,
    pub kind: EdgeKind,
}

/// CFG 原生边类别。
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum EdgeKind {
    Fallthrough,
    Jump,
    BranchTrue,
    BranchFalse,
    LoopBody,
    LoopExit,
    Return,
    TailCall,
}

impl Cfg {
    /// block 末尾指令通常决定了边形态,所以这里提供统一入口避免各层重复取尾。
    pub fn terminator<'a>(&self, instrs: &'a [LowInstr], block: BlockRef) -> Option<&'a LowInstr> {
        self.blocks
            .get(block.index())
            .and_then(|basic_block| basic_block.instrs.last())
            .and_then(|instr| instrs.get(instr.index()))
    }

    /// 读取 branch block 的真假边。
    ///
    /// 这个查询只依赖 CFG 边的 kind,不应散在 Structure/HIR 各自维护一份 helper。
    pub fn branch_edges(&self, block: BlockRef) -> Option<(EdgeRef, EdgeRef)> {
        let succs = &self.succs[block.index()];
        if succs.len() != 2 {
            return None;
        }

        let then_edge = succs
            .iter()
            .find(|edge_ref| matches!(self.edges[edge_ref.index()].kind, EdgeKind::BranchTrue))?;
        let else_edge = succs
            .iter()
            .find(|edge_ref| matches!(self.edges[edge_ref.index()].kind, EdgeKind::BranchFalse))?;

        Some((*then_edge, *else_edge))
    }

    /// 返回去重后的 reachable successors。
    pub fn reachable_successors(&self, block: BlockRef) -> Vec<BlockRef> {
        let mut succs = self.succs[block.index()]
            .iter()
            .map(|edge_ref| self.edges[edge_ref.index()].to)
            .filter(|succ| self.reachable_blocks.contains(succ))
            .collect::<Vec<_>>();
        succs.sort();
        succs.dedup();
        succs
    }

    /// 返回去重后的 reachable predecessors。
    pub fn reachable_predecessors(&self, block: BlockRef) -> Vec<BlockRef> {
        let mut preds = self.preds[block.index()]
            .iter()
            .map(|edge_ref| self.edges[edge_ref.index()].from)
            .filter(|pred| self.reachable_blocks.contains(pred))
            .collect::<Vec<_>>();
        preds.sort();
        preds.dedup();
        preds
    }

    /// 如果 block 在筛选后只有一个去重 reachable predecessor,返回它。
    pub fn unique_reachable_predecessor_matching(
        &self,
        block: BlockRef,
        mut matches: impl FnMut(BlockRef) -> bool,
    ) -> Option<BlockRef> {
        let mut unique = None;
        for edge_ref in &self.preds[block.index()] {
            let pred = self.edges[edge_ref.index()].from;
            if !self.reachable_blocks.contains(&pred) || !matches(pred) {
                continue;
            }
            match unique {
                None => unique = Some(pred),
                Some(existing) if existing == pred => {}
                Some(_) => return None,
            }
        }
        unique
    }

    /// 如果 block 去重后只有一个 reachable successor,返回它。
    pub fn unique_reachable_successor(&self, block: BlockRef) -> Option<BlockRef> {
        self.reachable_successor_shape(block).unique()
    }

    /// 返回去重后 reachable successor 的数量形态,不为常见线性查询分配临时 Vec。
    pub fn reachable_successor_shape(&self, block: BlockRef) -> ReachableSuccessorShape {
        let mut unique = None;
        for edge_ref in &self.succs[block.index()] {
            let succ = self.edges[edge_ref.index()].to;
            if !self.reachable_blocks.contains(&succ) {
                continue;
            }
            match unique {
                None => unique = Some(succ),
                Some(existing) if existing == succ => {}
                Some(_) => return ReachableSuccessorShape::Multiple,
            }
        }

        match unique {
            Some(succ) => ReachableSuccessorShape::Single(succ),
            None => ReachableSuccessorShape::Empty,
        }
    }

    pub fn can_reach(&self, from: BlockRef, to: BlockRef) -> bool {
        self.can_reach_within(from, to, &self.reachable_blocks)
    }

    pub fn can_reach_within(
        &self,
        from: BlockRef,
        to: BlockRef,
        allowed_blocks: &BTreeSet<BlockRef>,
    ) -> bool {
        if from == to {
            return true;
        }

        let mut visited = BTreeSet::new();
        let mut worklist = VecDeque::from([from]);

        while let Some(block) = worklist.pop_front() {
            if !self.reachable_blocks.contains(&block)
                || !allowed_blocks.contains(&block)
                || !visited.insert(block)
            {
                continue;
            }

            for edge_ref in &self.succs[block.index()] {
                let succ = self.edges[edge_ref.index()].to;
                if succ == to {
                    return true;
                }
                if allowed_blocks.contains(&succ) {
                    worklist.push_back(succ);
                }
            }
        }

        false
    }
}