neo-decompiler 0.10.1

use std::collections::HashSet;

use crate::decompiler::cfg::ssa::{ssa_expr_to_ir, SsaExpr};
use crate::decompiler::cfg::{BlockId, Terminator};
use crate::decompiler::ir::{Block as IrBlock, Expr};

use super::util::{is_literal, strip_version, var_name};
use super::StructCtx;

pub(super) struct SwitchResult {
    pub(super) scrutinee: Expr,
    pub(super) cases: Vec<(Expr, IrBlock)>,
    pub(super) default: Option<IrBlock>,
    pub(super) merge: Option<BlockId>,
}

impl<'a> StructCtx<'a> {
    /// Try to recognise a `switch` equality-cascade starting at `bid`.
    pub(super) fn try_switch(
        &self,
        bid: BlockId,
        then_target: BlockId,
        else_target: BlockId,
        visited: &mut HashSet<BlockId>,
    ) -> Option<SwitchResult> {
        let first = self.condition_expression(bid)?;
        let (scrut_base, first_val) = extract_eq_cond(&first)?;

        let mut cases: Vec<(Expr, BlockId)> = vec![(ssa_expr_to_ir(&first_val), then_target)];
        let mut cur = else_target;
        let default_entry;
        loop {
            if self.cfg.predecessors(cur).len() >= 2 {
                default_entry = cur;
                break;
            }
            match self.terminator(cur) {
                Terminator::Branch {
                    then_target,
                    else_target,
                } => {
                    let cond = self.condition_expression(cur);
                    match cond.and_then(|c| extract_eq_cond(&c)) {
                        Some((base, val)) if base == scrut_base => {
                            cases.push((ssa_expr_to_ir(&val), then_target));
                            cur = else_target;
                        }
                        _ => {
                            default_entry = cur;
                            break;
                        }
                    }
                }
                _ => {
                    default_entry = cur;
                    break;
                }
            }
        }

        if cases.len() < 2 {
            return None;
        }

        let merge = self.find_merge(cases[0].1, default_entry);

        let mut case_blocks: Vec<(Expr, IrBlock)> = Vec::with_capacity(cases.len());
        for (val, body_entry) in &cases {
            let body = self.structure_region(*body_entry, merge, visited, true);
            case_blocks.push((val.clone(), body));
        }
        let default_body = self.structure_region(default_entry, merge, visited, true);
        let default = if default_body.is_empty() {
            None
        } else {
            Some(default_body)
        };

        Some(SwitchResult {
            scrutinee: Expr::Variable(scrut_base),
            cases: case_blocks,
            default,
            merge,
        })
    }
}

fn extract_eq_cond(expr: &SsaExpr) -> Option<(String, SsaExpr)> {
    use crate::decompiler::ir::BinOp;
    let SsaExpr::Binary { op, left, right } = expr else {
        return None;
    };
    if !matches!(*op, BinOp::Eq) {
        return None;
    }
    match (left.as_ref(), right.as_ref()) {
        (SsaExpr::Variable(v), lit) if is_literal(lit) => {
            Some((strip_version(&var_name(v)), lit.clone()))
        }
        (lit, SsaExpr::Variable(v)) if is_literal(lit) => {
            Some((strip_version(&var_name(v)), lit.clone()))
        }
        _ => None,
    }
}