neo-decompiler 0.10.1

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

use crate::decompiler::cfg::{BlockId, Terminator};
use crate::decompiler::ir::{ControlFlow, Stmt};

use super::StructCtx;

impl<'a> StructCtx<'a> {
    /// Recover a `try` / `catch` / `finally` from a `TryEntry` terminator.
    pub(super) fn handle_try(
        &self,
        out: &mut crate::decompiler::ir::Block,
        bid: BlockId,
        body_target: BlockId,
        catch_target: Option<BlockId>,
        finally_target: Option<BlockId>,
        visited: &mut HashSet<BlockId>,
    ) -> Option<BlockId> {
        let continuation = self.find_endtry_continuation(body_target);

        let mut body_stop: HashSet<BlockId> = HashSet::new();
        if let Some(c) = catch_target {
            body_stop.insert(c);
        }
        if let Some(f) = finally_target {
            body_stop.insert(f);
        }
        if let Some(c) = continuation {
            body_stop.insert(c);
        }
        let try_body = self.structure_set(body_target, &body_stop, visited);

        let catch_body = catch_target.map(|c| {
            let mut stop = HashSet::new();
            if let Some(f) = finally_target {
                stop.insert(f);
            }
            if let Some(cont) = continuation {
                stop.insert(cont);
            }
            self.structure_set(c, &stop, visited)
        });
        let finally_body = finally_target.map(|f| {
            let mut stop = HashSet::new();
            if let Some(cont) = continuation {
                stop.insert(cont);
            }
            self.structure_set(f, &stop, visited)
        });

        out.push(Stmt::ControlFlow(Box::new(ControlFlow::try_catch(
            try_body,
            None,
            catch_body,
            finally_body,
        ))));
        let _ = bid;
        continuation
    }

    fn find_endtry_continuation(&self, start: BlockId) -> Option<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) {
                if let Terminator::EndTry { continuation } = &block.terminator {
                    return Some(*continuation);
                }
                for s in block.terminator.successors() {
                    stack.push(s);
                }
            }
        }
        None
    }
}