From e9a9ed0453e0fe375bc0ac8972cd68efed2dab8d Mon Sep 17 00:00:00 2001
From: Luke Street <luke.street@encounterpc.com>
Date: Wed, 3 Apr 2024 00:11:58 -0600
Subject: [PATCH] Fixes for switch analysis

- When finalizing a block, restart analysis at
  the first block
- Check for default case in jump table entries
---
 src/analysis/cfa.rs    | 12 ++++++++++--
 src/analysis/slices.rs |  2 +-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/analysis/cfa.rs b/src/analysis/cfa.rs
index b3de0b4..ad77c27 100644
--- a/src/analysis/cfa.rs
+++ b/src/analysis/cfa.rs
@@ -248,8 +248,12 @@ impl AnalyzerState {
         }
         if self.functions.iter().any(|(_, i)| i.is_unfinalized()) {
             log::error!("Failed to finalize functions:");
-            for (addr, _) in self.functions.iter().filter(|(_, i)| i.is_unfinalized()) {
-                log::error!("  {:#010X}", addr);
+            for (addr, info) in self.functions.iter().filter(|(_, i)| i.is_unfinalized()) {
+                log::error!(
+                    "  {:#010X}: blocks [{:?}]",
+                    addr,
+                    info.slices.as_ref().unwrap().possible_blocks.keys()
+                );
             }
             bail!("Failed to finalize functions");
         }
@@ -298,11 +302,15 @@ impl AnalyzerState {
                             &self.functions,
                             Some(vm),
                         )?;
+                        // Start at the beginning of the function again
+                        current = SectionAddress::new(addr.section, 0);
                     }
                     TailCallResult::Is => {
                         log::trace!("Finalized tail call @ {:#010X}", block);
                         slices.possible_blocks.remove(&block);
                         slices.function_references.insert(block);
+                        // Start at the beginning of the function again
+                        current = SectionAddress::new(addr.section, 0);
                     }
                     TailCallResult::Possible => {
                         if finalize {
diff --git a/src/analysis/slices.rs b/src/analysis/slices.rs
index ac02d41..85b95c4 100644
--- a/src/analysis/slices.rs
+++ b/src/analysis/slices.rs
@@ -328,7 +328,7 @@ impl FunctionSlices {
                         function_end.or_else(|| self.end()),
                     )?;
                     log::debug!("-> size {}: {:?}", size, entries);
-                    if entries.contains(&next_address)
+                    if (entries.contains(&next_address) || self.blocks.contains_key(&next_address))
                         && !entries.iter().any(|&addr| {
                             self.is_known_function(known_functions, addr)
                                 .is_some_and(|fn_addr| fn_addr != function_start)