neo-decompiler 0.10.2

use std::collections::{BTreeSet, HashSet};

use crate::decompiler::cfg::{BlockId, Terminator};
use crate::decompiler::ir::{BinOp, Block as IrBlock, ControlFlow, Expr, Stmt};

use super::StructCtx;
use crate::decompiler::cfg::ssa::{ssa_expr_to_ir, SsaExpr, SsaForm, SsaStmt};

impl<'a> StructCtx<'a> {
    /// Recover an `if` / `if-else` from a `Branch` terminator.
    pub(super) fn handle_branch(
        &self,
        out: &mut IrBlock,
        bid: BlockId,
        then_target: BlockId,
        else_target: BlockId,
        outer_boundary: Option<BlockId>,
        visited: &mut HashSet<BlockId>,
    ) -> Option<BlockId> {
        let cond = Self::comparison_condition_for_block(self.ssa, bid)
            .unwrap_or_else(|| self.condition_for_block(bid));

        if then_target == else_target {
            return Some(then_target);
        }

        if let Some(switch) = self.try_switch(bid, then_target, else_target, visited) {
            out.push(Stmt::ControlFlow(Box::new(ControlFlow::Switch {
                expr: switch.scrutinee,
                cases: switch.cases,
                default: switch.default,
            })));
            return match switch.merge {
                Some(m) if Some(m) != outer_boundary => Some(m),
                _ => None,
            };
        }

        let merge = self.find_merge(then_target, else_target);
        let then_block = self.structure_region(then_target, merge, visited, true);
        let else_block = self.structure_region(else_target, merge, visited, true);

        let cf = if else_block.is_empty() && then_block.is_empty() {
            out.push(Stmt::ExprStmt(cond));
            return merge.or(Some(else_target));
        } else if else_block.is_empty() {
            ControlFlow::if_then(cond, then_block)
        } else {
            ControlFlow::if_else(cond, then_block, else_block)
        };
        out.push(Stmt::ControlFlow(Box::new(cf)));

        match merge {
            Some(m) if Some(m) != outer_boundary => Some(m),
            _ => None,
        }
    }

    /// Inline the comparison expression driving a branch when available.
    pub(super) fn comparison_condition_for_block(ssa: &SsaForm, bid: BlockId) -> Option<Expr> {
        let value = ssa.block(bid)?.stmts.last().and_then(|s| match s {
            SsaStmt::Assign { value, .. } => Some(value),
            _ => None,
        })?;
        let SsaExpr::Binary { op, .. } = value else {
            return None;
        };
        let is_comparison = matches!(
            op,
            BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge
        );
        is_comparison.then(|| ssa_expr_to_ir(value))
    }

    /// The last assignment's raw SSA value expression in `bid`, if any.
    pub(super) fn condition_expression(&self, bid: BlockId) -> Option<SsaExpr> {
        let block = self.ssa.block(bid)?;
        for stmt in block.stmts.iter().rev() {
            if let SsaStmt::Assign { value, .. } = stmt {
                return Some(value.clone());
            }
        }
        None
    }

    /// Find the merge of two branch arms.
    pub(super) fn find_merge(&self, then_target: BlockId, else_target: BlockId) -> Option<BlockId> {
        let from_then = self.reachable(then_target);
        let from_else = self.reachable(else_target);
        let common: BTreeSet<BlockId> = from_then.intersection(&from_else).copied().collect();
        common
            .into_iter()
            .find(|b| self.cfg.predecessors(*b).len() >= 2)
    }

    /// All blocks reachable from `start` (inclusive).
    pub(super) fn reachable(&self, start: BlockId) -> BTreeSet<BlockId> {
        let mut seen = BTreeSet::new();
        let mut stack = vec![start];
        while let Some(b) = stack.pop() {
            if !seen.insert(b) {
                continue;
            }
            if let Some(block) = self.cfg.block(b) {
                for s in block.terminator.successors() {
                    stack.push(s);
                }
            }
        }
        seen
    }

    pub(super) fn structure_set(
        &self,
        entry: BlockId,
        stop: &HashSet<BlockId>,
        visited: &mut HashSet<BlockId>,
    ) -> IrBlock {
        let mut out = IrBlock::new();
        let mut cur = Some(entry);
        while let Some(bid) = cur {
            if stop.contains(&bid) {
                break;
            }
            if !visited.insert(bid) {
                break;
            }
            self.emit_body(&mut out, bid);
            cur = match self.terminator(bid) {
                Terminator::Return
                | Terminator::Throw
                | Terminator::Abort
                | Terminator::Unknown
                | Terminator::EndTry { .. }
                | Terminator::TryEntry { .. } => None,
                Terminator::Fallthrough { target } | Terminator::Jump { target } => Some(target),
                Terminator::Branch {
                    then_target,
                    else_target,
                } => {
                    if self.loop_headers.contains(&bid) {
                        self.emit_while_loop(&mut out, bid, then_target, else_target, visited)
                    } else {
                        self.handle_branch(&mut out, bid, then_target, else_target, None, visited)
                    }
                }
            };
            if matches!(cur, Some(c) if stop.contains(&c)) {
                break;
            }
        }
        out
    }
}