From 3385f583413f6bb0f8c7e6c99084ef289fa6d4e6 Mon Sep 17 00:00:00 2001
From: Mark Langen <stravant@gmail.com>
Date: Thu, 17 Jul 2025 15:04:56 -0700
Subject: [PATCH] Fix data flow analysis for multiple text sections (#220)

* Fix data flow analysis for multiple text sections

* Data flow analysis results were only keyed by the symbol (function)
  address. That doen't work if there are multiple text sections, the
  result from the first function in one section will stomp the result
  from the first function in another because both have address zero.

* Remove the ambiguity by keying off of the section address as well.

* Formatting

* Satisfy wasm build

* Clippy

* Formatting again

* Thought that section was the section address not the section number.

---------

Co-authored-by: Luke Street <luke.street@encounterpc.com>
---
 objdiff-core/src/diff/display.rs |  4 ++--
 objdiff-core/src/obj/mod.rs      | 14 ++++++++++++++
 objdiff-core/src/obj/read.rs     | 17 ++++++++++++-----
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/objdiff-core/src/diff/display.rs b/objdiff-core/src/diff/display.rs
index 99b678d..8e35939 100644
--- a/objdiff-core/src/diff/display.rs
+++ b/objdiff-core/src/diff/display.rs
@@ -189,7 +189,7 @@ pub fn display_row(
     let mut arg_idx = 0;
     let mut displayed_relocation = false;
     let analysis_result = if diff_config.show_data_flow {
-        obj.flow_analysis_results.get(&resolved.symbol.address)
+        obj.get_flow_analysis_result(resolved.symbol)
     } else {
         None
     };
@@ -217,7 +217,7 @@ pub fn display_row(
             }
             let data_flow_value =
             analysis_result.and_then(|result|
-                result.as_ref().get_argument_value_at_address(
+                result.get_argument_value_at_address(
                     ins_ref.address, (arg_idx - 1) as u8));
             match (arg, data_flow_value, resolved.ins_ref.branch_dest) {
                 // If we have a flow analysis result, always use that over anything else.
diff --git a/objdiff-core/src/obj/mod.rs b/objdiff-core/src/obj/mod.rs
index b754c4c..22b909f 100644
--- a/objdiff-core/src/obj/mod.rs
+++ b/objdiff-core/src/obj/mod.rs
@@ -338,6 +338,20 @@ impl Object {
         self.symbols.iter().position(|symbol| symbol.section.is_some() && symbol.name == name)
     }
 
+    pub fn get_flow_analysis_result(&self, symbol: &Symbol) -> Option<&dyn FlowAnalysisResult> {
+        let key = symbol.section.unwrap_or_default() as u64 * 1024 * 1024 * 1024 + symbol.address;
+        self.flow_analysis_results.get(&key).map(|result| result.as_ref())
+    }
+
+    pub fn add_flow_analysis_result(
+        &mut self,
+        symbol: &Symbol,
+        result: Box<dyn FlowAnalysisResult>,
+    ) {
+        let key = symbol.section.unwrap_or_default() as u64 * 1024 * 1024 * 1024 + symbol.address;
+        self.flow_analysis_results.insert(key, result);
+    }
+
     pub fn has_flow_analysis_result(&self) -> bool { !self.flow_analysis_results.is_empty() }
 }
 
diff --git a/objdiff-core/src/obj/read.rs b/objdiff-core/src/obj/read.rs
index e194e2f..8196e37 100644
--- a/objdiff-core/src/obj/read.rs
+++ b/objdiff-core/src/obj/read.rs
@@ -1,4 +1,5 @@
 use alloc::{
+    boxed::Box,
     collections::BTreeMap,
     format,
     string::{String, ToString},
@@ -13,8 +14,8 @@ use crate::{
     arch::{Arch, new_arch},
     diff::DiffObjConfig,
     obj::{
-        Object, Relocation, RelocationFlags, Section, SectionData, SectionFlag, SectionKind,
-        Symbol, SymbolFlag, SymbolKind,
+        FlowAnalysisResult, Object, Relocation, RelocationFlags, Section, SectionData, SectionFlag,
+        SectionKind, Symbol, SymbolFlag, SymbolKind,
         split_meta::{SPLITMETA_SECTION, SplitMeta},
     },
     util::{align_data_slice_to, align_u64_to, read_u16, read_u32},
@@ -439,6 +440,7 @@ fn perform_data_flow_analysis(obj: &mut Object, config: &DiffObjConfig) -> Resul
     }
 
     let mut generated_relocations = Vec::<(usize, Vec<Relocation>)>::new();
+    let mut generated_flow_results = Vec::<(Symbol, Box<dyn FlowAnalysisResult>)>::new();
     for (section_index, section) in obj.sections.iter().enumerate() {
         if section.kind != SectionKind::Code {
             continue;
@@ -474,12 +476,17 @@ fn perform_data_flow_analysis(obj: &mut Object, config: &DiffObjConfig) -> Resul
 
             // Optional full data flow analysis
             if config.analyze_data_flow {
-                obj.arch.data_flow_analysis(obj, symbol, code, &section.relocations).and_then(
-                    |flow_result| obj.flow_analysis_results.insert(symbol.address, flow_result),
-                );
+                if let Some(flow_result) =
+                    obj.arch.data_flow_analysis(obj, symbol, code, &section.relocations)
+                {
+                    generated_flow_results.push((symbol.clone(), flow_result));
+                }
             }
         }
     }
+    for (symbol, flow_result) in generated_flow_results {
+        obj.add_flow_analysis_result(&symbol, flow_result);
+    }
     for (section_index, mut relocations) in generated_relocations {
         obj.sections[section_index].relocations.append(&mut relocations);
     }