Recover cleanup state bindings and shared handler finally

dexdec/src/analysis/value_recovery.rs

@@ -18,7 +18,7 @@ mod source;
 mod ssa_constants;
 
 use crate::analysis::SemanticTransform;
-use crate::ir::analysis::{DominanceError, SsaVar};
+use crate::ir::analysis::{DominanceError, PhiMerge, SsaVar};
 use crate::ir::{
     BlockId, SemanticFoldError, SemanticMethod, SourceVariableContext, SsaSemantics,
     ValueSemantics, CFG,
@@ -258,26 +258,41 @@ impl SemanticTransform<SsaSemantics> for ValueRecovery {
             "value.ssa.specialization_apply",
             specialization_schedule.apply(method.body_mut())
         )?;
-        let retained =
-            crate::profile_scope!("value.ssa.cleanup_roots", cleanup_state_values(&method));
+        let retained = crate::profile_scope!(
+            "value.ssa.cleanup_roots",
+            cleanup_state_values(&method, &self.ssa_constants)
+        );
         let constants = SsaValueSolver::new(&recovered_phis, &self.ssa_constants, &retained)
             .solve_profiled(&mut method)?;
         Ok(method.into_values(constants, recovered_phis))
     }
 }
 
-fn cleanup_state_values(method: &SemanticMethod<SsaSemantics>) -> BTreeSet<SsaVar> {
-    let mut values = method
+fn cleanup_state_values(
+    method: &SemanticMethod<SsaSemantics>,
+    constants: &BTreeMap<SsaVar, crate::ir::InsnArg>,
+) -> BTreeSet<SsaVar> {
+    let roots = method
         .state()
         .regions()
         .cleanup_value_bindings()
         .iter()
-        .flat_map(|(handler, normal)| [*handler, *normal])
+        .flat_map(|(handler, normal)| [*handler, *normal]);
+    cleanup_state_closure(roots, method.state().values().phis(), constants)
+}
+
+fn cleanup_state_closure(
+    roots: impl IntoIterator<Item = SsaVar>,
+    phis: &[PhiMerge],
+    constants: &BTreeMap<SsaVar, crate::ir::InsnArg>,
+) -> BTreeSet<SsaVar> {
+    // Shared null/zero roots can be inlined, but a cleanup root that is a
+    // non-trivial constant (for example a catch-path `-1`) must stay live.
+    let mut values = roots
+        .into_iter()
+        .filter(|value| !trivial_cleanup_constant(constants, *value))
         .collect::<BTreeSet<_>>();
-    let phis = method
-        .state()
-        .values()
-        .phis()
+    let phis = phis
         .iter()
         .map(|phi| (phi.result, phi))
         .collect::<BTreeMap<_, _>>();
@@ -287,14 +302,120 @@ fn cleanup_state_values(method: &SemanticMethod<SsaSemantics>) -> BTreeSet<SsaVa
             continue;
         };
         for input in &phi.inputs {
-            if values.insert(input.value) {
+            if !canonical_constant(constants, input.value) && values.insert(input.value) {
                 pending.push(input.value);
             }
         }
     }
     values
 }
 
+fn canonical_constant(constants: &BTreeMap<SsaVar, crate::ir::InsnArg>, mut value: SsaVar) -> bool {
+    let mut visited = BTreeSet::new();
+    while visited.insert(value) {
+        match constants.get(&value) {
+            Some(crate::ir::InsnArg::Reg(register)) => {
+                let Some(next) = SsaVar::from_reg(register) else {
+                    return false;
+                };
+                value = next;
+            }
+            Some(crate::ir::InsnArg::Lit(_)) => return true,
+            Some(crate::ir::InsnArg::Wrapped(instruction)) => {
+                return matches!(
+                    instruction.insn_type,
+                    crate::ir::InsnType::Const | crate::ir::InsnType::ConstStr
+                );
+            }
+            None => return false,
+        }
+    }
+    false
+}
+
+fn trivial_cleanup_constant(
+    constants: &BTreeMap<SsaVar, crate::ir::InsnArg>,
+    mut value: SsaVar,
+) -> bool {
+    let mut visited = BTreeSet::new();
+    while visited.insert(value) {
+        match constants.get(&value) {
+            Some(crate::ir::InsnArg::Reg(register)) => {
+                let Some(next) = SsaVar::from_reg(register) else {
+                    return false;
+                };
+                value = next;
+            }
+            Some(crate::ir::InsnArg::Lit(literal)) => return literal.is_zero(),
+            Some(crate::ir::InsnArg::Wrapped(instruction)) => {
+                return matches!(
+                    instruction.insn_type,
+                    crate::ir::InsnType::Const | crate::ir::InsnType::ConstStr
+                ) && instruction.args.iter().any(|argument| match argument {
+                    crate::ir::InsnArg::Lit(literal) => literal.is_zero(),
+                    _ => false,
+                });
+            }
+            None => return false,
+        }
+    }
+    false
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ir::analysis::PhiInput;
+    use crate::ir::{ArgType, EdgeKind, InsnArg, InstructionId};
+
+    #[test]
+    fn cleanup_state_closure_does_not_retain_shared_null_constant() {
+        let result = SsaVar::new(1, 26);
+        let live_input = SsaVar::new(1, 25);
+        let normal = SsaVar::new(12, 3);
+        let null = SsaVar::new(1, 8);
+        let phis = [PhiMerge {
+            block: BlockId::new(118),
+            instruction: InstructionId::new(0),
+            result,
+            inputs: vec![
+                PhiInput {
+                    predecessor: BlockId::new(74),
+                    edge_kind: EdgeKind::Normal,
+                    value: null,
+                },
+                PhiInput {
+                    predecessor: BlockId::new(117),
+                    edge_kind: EdgeKind::Normal,
+                    value: live_input,
+                },
+            ],
+        }];
+        let constants = BTreeMap::from([(null, InsnArg::lit(0, ArgType::unknown_object()))]);
+
+        assert_eq!(
+            cleanup_state_closure([result, normal], &phis, &constants),
+            BTreeSet::from([result, live_input, normal])
+        );
+    }
+
+    #[test]
+    fn cleanup_state_closure_retains_a_nonzero_constant_root() {
+        let catch_result = SsaVar::new(1, 13);
+        let normal = SsaVar::new(1, 6);
+        let zero = SsaVar::new(1, 0);
+        let constants = BTreeMap::from([
+            (catch_result, InsnArg::lit(-1, ArgType::INT)),
+            (zero, InsnArg::lit(0, ArgType::INT)),
+        ]);
+
+        assert_eq!(
+            cleanup_state_closure([catch_result, normal, zero], &[], &constants),
+            BTreeSet::from([catch_result, normal])
+        );
+    }
+}
+
 struct SsaValueSolver<'a> {
     recovered_phis: &'a BTreeSet<SsaVar>,
     constants: &'a BTreeMap<SsaVar, crate::ir::InsnArg>,

dexdec/src/ir/exception.rs

@@ -103,6 +103,58 @@ impl TryRegion {
     }
 }
 
+/// Canonical handler entries referenced by more than one protected region.
+fn shared_handler_entries(handlers: impl IntoIterator<Item = (u32, BlockId)>) -> BTreeSet<BlockId> {
+    let mut owners = BTreeMap::<BlockId, BTreeSet<u32>>::new();
+    for (region, entry) in handlers {
+        owners.entry(entry).or_default().insert(region);
+    }
+    owners
+        .into_iter()
+        .filter_map(|(entry, regions)| (regions.len() > 1).then_some(entry))
+        .collect()
+}
+
+/// Catch entries shared by try regions that do not share a cleanup identity.
+///
+/// Split fragments of one source `try` reuse the same catch and finally
+/// entries; those must still recover structured `finally`. Two independent
+/// tries that only share a catch tail must not elide that tail on behalf of
+/// one region.
+fn conflicting_shared_catch_entries(regions: &[TryRegion]) -> BTreeSet<BlockId> {
+    let mut owners = BTreeMap::<BlockId, BTreeSet<(u32, Option<BlockId>)>>::new();
+    for region in regions {
+        let cleanup = region
+            .handlers
+            .iter()
+            .find(|handler| handler.kind != HandlerKind::Catch)
+            .map(|handler| handler.canonical_entry);
+        for handler in region.catch_handlers() {
+            owners
+                .entry(handler.canonical_entry)
+                .or_default()
+                .insert((region.id, cleanup));
+        }
+    }
+    owners
+        .into_iter()
+        .filter_map(|(entry, identities)| {
+            let region_count = identities
+                .iter()
+                .map(|(region, _)| *region)
+                .collect::<BTreeSet<_>>()
+                .len();
+            let cleanup_identities = identities
+                .iter()
+                .map(|(_, cleanup)| *cleanup)
+                .collect::<BTreeSet<_>>();
+            let conflicting_cleanup =
+                cleanup_identities.len() > 1 || cleanup_identities == BTreeSet::from([None]);
+            (region_count > 1 && conflicting_cleanup).then_some(entry)
+        })
+        .collect()
+}
+
 #[derive(Debug, Clone)]
 pub struct CleanupContraction {
     pub entry: BlockId,
@@ -409,6 +461,7 @@ impl<'a> ExceptionAnalyzer<'a> {
         HandlerDomains::assign(self.cfg, &mut regions);
         let nested_handlers = NestedHandlerDomains::analyze(&regions);
         let recovery_order = Self::cleanup_recovery_order(&regions);
+        let shared_handler_entries = conflicting_shared_catch_entries(&regions);
 
         let mut elided_instructions = BTreeSet::new();
         let mut cleanup_contractions = Vec::new();
@@ -422,7 +475,13 @@ impl<'a> ExceptionAnalyzer<'a> {
                 .ok_or(ExceptionInvariantError::MissingExceptionScope(region_id))?;
             let nested_cleanup = nested_handlers.cleanup(region.id);
             let nested_all = nested_handlers.all(region.id);
-            let cleanup = CleanupRecovery::new(self.cfg, self.values, &normal_dominators).recover(
+            let cleanup = CleanupRecovery::new(
+                self.cfg,
+                self.values,
+                &normal_dominators,
+                &shared_handler_entries,
+            )
+            .recover(
                 region,
                 &nested_cleanup,
                 &nested_all,
@@ -4342,6 +4401,43 @@ mod tests {
         }
     }
 
+    #[test]
+    fn identifies_handler_entry_shared_across_protected_regions() {
+        let shared = BlockId::new(7);
+        let private = BlockId::new(8);
+        let entries = shared_handler_entries([(1, shared), (1, shared), (2, shared), (2, private)]);
+
+        assert_eq!(entries, BTreeSet::from([shared]));
+    }
+
+    #[test]
+    fn split_try_fragments_that_share_finally_are_not_conflicting_catches() {
+        let mut cleanup = catch_handler(4, &[4]);
+        cleanup.kind = HandlerKind::Cleanup;
+        let left = try_region(
+            1,
+            0,
+            4,
+            &[0, 1],
+            vec![catch_handler(3, &[3]), cleanup.clone()],
+        );
+        let right = try_region(2, 4, 8, &[2], vec![catch_handler(3, &[3]), cleanup]);
+
+        assert!(conflicting_shared_catch_entries(&[left, right]).is_empty());
+    }
+
+    #[test]
+    fn independent_tries_sharing_only_a_catch_are_conflicting() {
+        let catch = BlockId::new(3);
+        let left = try_region(1, 0, 4, &[0], vec![catch_handler(3, &[3])]);
+        let right = try_region(2, 4, 8, &[1], vec![catch_handler(3, &[3])]);
+
+        assert_eq!(
+            conflicting_shared_catch_entries(&[left, right]),
+            BTreeSet::from([catch])
+        );
+    }
+
     #[test]
     fn shared_noop_catch_continuation_is_not_inherited_by_outer_handler() {
         let mut cfg = CFG::new("shared_noop_catch_continuation");