feat(decompile): detect and resolve internal function calls

crates/common/src/ether/signatures.rs

@@ -710,7 +710,7 @@ mod tests {
                                           // num_params = 1, num_dyn_params = 2 + 1 = 3
                                           // Without saturating_sub: 1 - 3 would underflow
                                           // With saturating_sub: 1.saturating_sub(3) = 0
-        let score = score_signature(&signature, Some(0));
+        let score = score_signature(signature, Some(0));
         // Should not panic, should return a valid score (greater than 0)
         assert!(score > 0);
     }
@@ -722,7 +722,7 @@ mod tests {
                                   // num_params = 1, num_dyn_params = 0, num_static_params = 1
                                   // If num_words = 10, then num_words - num_static_params = 9
                                   // This would reduce score by 90
-        let score = score_signature(&signature, Some(10));
+        let score = score_signature(signature, Some(10));
         // Should not panic and should be reduced appropriately
         // Initial score calculation:
         // - Start: 1000
@@ -742,7 +742,7 @@ mod tests {
         // num_dyn_params = bytes(3) + string(1) + [(1) = 5
         // Without saturating_sub: 3 - 5 would underflow
         // With saturating_sub: 3.saturating_sub(5) = 0
-        let score = score_signature(&signature, Some(2));
+        let score = score_signature(signature, Some(2));
         // Should not panic and should return a valid score
         // The score should be positive since we add 10 per param
         assert!(score > 0);

crates/decompile/src/core/analyze.rs

@@ -1,6 +1,7 @@
-use std::{fmt::Display, time::Instant};
+use std::{fmt::Display, sync::Arc, time::Instant};
 
 use futures::future::BoxFuture;
+use hashbrown::HashMap;
 use heimdall_vm::ext::exec::VMTrace;
 use tracing::debug;
 
@@ -59,6 +60,10 @@ pub(crate) struct AnalyzerState {
     pub analyzer_type: AnalyzerType,
     /// Whether to skip resolving internal calls
     pub skip_resolving: bool,
+    /// Map of selector -> argument count for internal call resolution
+    pub selector_arg_counts: Arc<HashMap<String, usize>>,
+    /// Map of selector -> resolved function name for internal call resolution
+    pub selector_names: Arc<HashMap<String, String>>,
 }
 
 /// The analyzer, which will analyze a [`VMTrace`] generated by symbolic execution and build an
@@ -74,12 +79,29 @@ pub(crate) struct Analyzer {
     function: AnalyzedFunction,
     /// A list of registered heuristics with the Heuristic Trait
     heuristics: Vec<Heuristic>,
+    /// Map of selector -> argument count for internal call resolution
+    selector_arg_counts: Arc<HashMap<String, usize>>,
+    /// Map of selector -> resolved function name for internal call resolution
+    selector_names: Arc<HashMap<String, String>>,
 }
 
 impl Analyzer {
     /// Build a new analyzer with the given type, function, and trace
-    pub(crate) fn new(typ: AnalyzerType, skip_resolving: bool, function: AnalyzedFunction) -> Self {
-        Self { typ, function, skip_resolving, heuristics: Vec::new() }
+    pub(crate) fn new(
+        typ: AnalyzerType,
+        skip_resolving: bool,
+        function: AnalyzedFunction,
+        selector_arg_counts: Arc<HashMap<String, usize>>,
+        selector_names: Arc<HashMap<String, String>>,
+    ) -> Self {
+        Self {
+            typ,
+            function,
+            skip_resolving,
+            heuristics: Vec::new(),
+            selector_arg_counts,
+            selector_names,
+        }
     }
 
     /// Register heuristics for the given function and trace
@@ -126,6 +148,8 @@ impl Analyzer {
             conditional_stack: Vec::new(),
             analyzer_type: self.typ,
             skip_resolving: self.skip_resolving,
+            selector_arg_counts: Arc::clone(&self.selector_arg_counts),
+            selector_names: Arc::clone(&self.selector_names),
         };
 
         // Perform analysis
@@ -158,6 +182,46 @@ impl Analyzer {
                 }
             }
 
+            // Check for internal function call at the end of this trace
+            if let Some(internal_call) = &branch.internal_call {
+                // Get the expected argument count for the target function
+                let arg_count =
+                    analyzer_state.selector_arg_counts.get(&internal_call.selector).copied();
+
+                // Format arguments using the known arg count
+                // Stack captured top-to-bottom: [argN-1, argN-2, ..., arg0, return_addr, ...]
+                // Args are at the top, take first 'count' items and reverse for correct order
+                let args = match arg_count {
+                    Some(count) if count > 0 && internal_call.arguments.len() >= count => {
+                        // Take [0..count], reverse to get [arg0, arg1, ..., argN-1]
+                        // Use itertools-style join without collecting to Vec first
+                        let mut result = String::new();
+                        for (i, arg) in internal_call.arguments[..count].iter().rev().enumerate() {
+                            if i > 0 {
+                                result.push_str(", ");
+                            }
+                            result.push_str(arg);
+                        }
+                        result
+                    }
+                    _ => String::new(),
+                };
+
+                // Use the resolved function name if available, otherwise use Unresolved_SELECTOR
+                // Format directly into the output to avoid intermediate String allocation
+                match analyzer_state.selector_names.get(&internal_call.selector) {
+                    Some(name) => {
+                        self.function.logic.push(format!("return {}({});", name, args));
+                    }
+                    None => {
+                        self.function.logic.push(format!(
+                            "return Unresolved_{}({});",
+                            internal_call.selector, args
+                        ));
+                    }
+                }
+            }
+
             // recurse into the children of the current trace branch
             for child in &branch.children {
                 self.analyze_inner(child, analyzer_state).await?;

crates/decompile/src/core/mod.rs

@@ -22,9 +22,12 @@ use heimdall_common::{
 use heimdall_disassembler::{disassemble, DisassemblerArgsBuilder};
 use heimdall_vm::{
     core::vm::VM,
-    ext::selectors::{find_function_selectors, resolve_selectors},
+    ext::selectors::{find_function_selectors, resolve_internal_body, resolve_selectors},
+};
+use std::{
+    sync::Arc,
+    time::{Duration, Instant},
 };
-use std::time::{Duration, Instant};
 
 use crate::{
     core::{
@@ -150,6 +153,30 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
     let selectors = find_function_selectors(&evm, &assembly);
     debug!("finding function selectors took {:?}", start_selectors_time.elapsed());
 
+    // Resolve internal function bodies for each selector
+    // Internal bodies are where the actual function logic lives (after calldata parsing)
+    let mut internal_bodies: HashMap<u128, String> = HashMap::new();
+    for (selector, entry_point) in &selectors {
+        evm.reset();
+        let internal_body = resolve_internal_body(&mut evm, selector, *entry_point);
+        if internal_body > 0 {
+            internal_bodies.insert(internal_body, selector.clone());
+            debug!(
+                "selector {} (entry {}) has internal body at {}",
+                selector, entry_point, internal_body
+            );
+        }
+    }
+
+    // Build a map combining entry points and internal bodies for detection
+    // When tracing a function, we detect internal calls to other functions'
+    // internal bodies (not just their dispatcher entries)
+    let known_entry_points: HashMap<u128, String> = selectors
+        .iter()
+        .map(|(sel, ep)| (*ep, sel.clone()))
+        .chain(internal_bodies.iter().map(|(ib, sel)| (*ib, sel.clone())))
+        .collect();
+
     // resolve selectors (if enabled)
     let resolved_selectors = match args.skip_resolving {
         true => HashMap::new(),
@@ -176,12 +203,13 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
     }
 
     let overall_sym_exec_time = Instant::now();
-    for (selector, entry_point) in selectors {
+    for (selector, entry_point) in &selectors {
         let start_sym_exec_time = Instant::now();
         evm.reset();
         let (map, jumpdest_count) = match evm.symbolic_exec_selector(
-            &selector,
-            entry_point,
+            selector,
+            *entry_point,
+            &known_entry_points,
             Instant::now()
                 .checked_add(Duration::from_millis(args.timeout))
                 .expect("invalid timeout"),
@@ -199,14 +227,48 @@ pub async fn decompile(args: DecompilerArgs) -> Result<DecompileResult, Error> {
     debug!("symbolic execution took {:?}", overall_sym_exec_time.elapsed());
     info!("symbolically executed {} selectors", symbolic_execution_maps.len());
 
+    // Build a map of selector -> argument count from resolved signatures
+    // e.g., "transferFrom(address,address,uint256)" -> 3 arguments
+    let selector_arg_counts: Arc<HashMap<String, usize>> = Arc::new(
+        resolved_selectors
+            .iter()
+            .filter_map(|(selector, funcs)| {
+                funcs.first().map(|f| {
+                    // Count commas in the signature to determine arg count
+                    // e.g., "transfer(address,uint256)" has 1 comma = 2 args
+                    let sig = &f.signature;
+                    let start = sig.find('(').unwrap_or(0);
+                    let end = sig.rfind(')').unwrap_or(sig.len());
+                    let params = &sig[start + 1..end];
+                    let count = if params.is_empty() { 0 } else { params.matches(',').count() + 1 };
+                    (selector.clone(), count)
+                })
+            })
+            .collect(),
+    );
+
+    // Build a map of selector -> function name from resolved signatures
+    let selector_names: Arc<HashMap<String, String>> = Arc::new(
+        resolved_selectors
+            .iter()
+            .filter_map(|(selector, funcs)| {
+                funcs.first().map(|f| (selector.clone(), f.name.clone()))
+            })
+            .collect(),
+    );
+
     let start_analysis_time = Instant::now();
     let handles = symbolic_execution_maps.into_iter().map(|(selector, trace_root)| {
         let mut evm_clone = evm.clone();
+        let arg_counts = Arc::clone(&selector_arg_counts);
+        let names = Arc::clone(&selector_names);
         async move {
             let mut analyzer = Analyzer::new(
                 analyzer_type,
                 args.skip_resolving,
                 AnalyzedFunction::new(&selector, selector == "fallback"),
+                arg_counts,
+                names,
             );
 
             // analyze the symbolic execution trace

crates/decompile/src/utils/heuristics/extcall.rs

@@ -39,9 +39,9 @@ pub(crate) fn extcall_heuristic<'a>(
                         instruction.instruction,
                         opcode_name(instruction.opcode)
                     );
-                    function.logic.push(format!(
-                        "(bool success, bytes memory ret0) = address({address}).transfer({value_solidified});"
-                    ));
+                    function
+                        .logic
+                        .push(format!("address({address}).transfer({value_solidified});"));
                     return Ok(());
                 }
                 if extcalldata.is_empty() {
@@ -50,9 +50,9 @@ pub(crate) fn extcall_heuristic<'a>(
                         instruction.instruction,
                         opcode_name(instruction.opcode)
                     );
-                    function.logic.push(format!(
-                        "(bool success, bytes memory ret0) = address({address}).transfer({value_solidified});"
-                    ));
+                    function
+                        .logic
+                        .push(format!("address({address}).transfer({value_solidified});"));
                     return Ok(());
                 }