feat: implement Milestone 2 Protocol Adapters and Uniswap v4 Hook memory extraction

Author: intelliDean

Cargo.lock

@@ -1,13 +1,71 @@
 pub trait ProtocolAdapter {
+    /// The name of the protocol (e.g., "Uniswap v4").
     fn name(&self) -> &str;
-    fn process(&self, trace: String) -> anyhow::Result<String>;
+    
+    /// Resolves a combination of target address and function selector into a human-readable label.
+    fn resolve_label(&self, address: Option<&str>, selector: Option<&str>) -> Option<String>;
 }
 
+/// Adapter specifically for identifying Uniswap v4 Hooks
 pub struct UniswapV4Adapter;
+
 impl ProtocolAdapter for UniswapV4Adapter {
-    fn name(&self) -> &str { "Uniswap v4" }
-    fn process(&self, trace: String) -> anyhow::Result<String> {
-        // Implementation for Uniswap v4 Hook profiling will go here
-        Ok(format!("Processing Uniswap v4 trace: {}", trace))
+    fn name(&self) -> &str {
+        "Uniswap v4"
+    }
+
+    fn resolve_label(&self, _address: Option<&str>, selector: Option<&str>) -> Option<String> {
+        let sel = selector?;
+        // Uniswap v4 Hook standard interface selectors
+        let label = match sel {
+            "0x18a9d381" => "beforeInitialize",
+            "0x999dea5d" => "afterInitialize",
+            "0x910746f2" => "beforeAddLiquidity",
+            "0xefd81287" => "afterAddLiquidity",
+            "0xd7386be3" => "beforeRemoveLiquidity",
+            "0x1efe5f9e" => "afterRemoveLiquidity",
+            "0xe82c3b75" => "beforeSwap",
+            "0x14d6eaec" => "afterSwap",
+            "0xa3d03227" => "beforeDonate",
+            "0x0df2d576" => "afterDonate",
+            _ => return None,
+        };
+        
+        Some(format!("Uniswapv4: {}", label))
+    }
+}
+
+/// The registry holding all known protocol adapters.
+pub struct AdapterRegistry {
+    adapters: Vec<Box<dyn ProtocolAdapter>>,
+}
+
+impl AdapterRegistry {
+    /// Initialize a new registry pre-loaded with all supported adapters.
+    pub fn new() -> Self {
+        let mut registry = Self { adapters: Vec::new() };
+        registry.register(Box::new(UniswapV4Adapter));
+        registry
+    }
+
+    /// Register a custom adapter
+    pub fn register(&mut self, adapter: Box<dyn ProtocolAdapter>) {
+        self.adapters.push(adapter);
+    }
+
+    /// Iterates through adapters to find a descriptive label for the call.
+    pub fn resolve(&self, address: Option<&str>, selector: Option<&str>) -> Option<String> {
+        for adapter in &self.adapters {
+            if let Some(label) = adapter.resolve_label(address, selector) {
+                return Some(label);
+            }
+        }
+        None
+    }
+}
+
+impl Default for AdapterRegistry {
+    fn default() -> Self {
+        Self::new()
     }
 }

crates/atupa-adapters/src/lib.rs

@@ -51,6 +51,8 @@ pub struct CollapsedStack {
     #[serde(default)]
     pub target_address: Option<String>,
     #[serde(default)]
+    pub resolved_label: Option<String>,
+    #[serde(default)]
     pub reverted: bool,
 }
 

crates/atupa-core/src/lib.rs

@@ -43,7 +43,9 @@ impl SvgGenerator {
             let leaf_name = stack.stack.split(';').last().unwrap_or("unknown");
 
             let mut label = format!("{} ({} gas)", leaf_name, stack.weight);
-            if let Some(addr) = &stack.target_address {
+            if let Some(r_label) = &stack.resolved_label {
+                label = format!("{} ({} gas)", r_label, stack.weight);
+            } else if let Some(addr) = &stack.target_address {
                 label = format!("{} [{}] ({} gas)", leaf_name, addr, stack.weight);
             }
             if stack.reverted {

crates/atupa-output/src/lib.rs

@@ -15,6 +15,7 @@ categories = { workspace = true }
 [dependencies]
 atupa-core = { workspace = true }
 atupa-rpc = { workspace = true }
+atupa-adapters = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 anyhow = { workspace = true }

crates/atupa-parser/Cargo.toml

@@ -15,8 +15,10 @@ impl Aggregator {
     pub fn build_collapsed_stacks(steps: &[TraceStep]) -> Vec<CollapsedStack> {
         debug!("Building collapsed stacks from {} execution steps", steps.len());
 
-        // Map to aggregate stacks: stack_string -> (total_gas, last_pc, target_address, reverted)
-        let mut stack_map: HashMap<String, (u64, u64, Option<String>, bool)> = HashMap::new();
+        let registry = atupa_adapters::AdapterRegistry::new();
+
+        // Map to aggregate stacks: stack_string -> (total_gas, last_pc, target_address, resolved_label, reverted)
+        let mut stack_map: HashMap<String, (u64, u64, Option<String>, Option<String>, bool)> = HashMap::new();
 
         // Current call stack
         let mut call_stack: Vec<String> = Vec::new();
@@ -35,21 +37,75 @@ impl Aggregator {
                 call_stack.push("CALL".to_string());
             }
 
-            // Extract Target Address if this is a Call opcode
+            // Extract Target Address & Parse Function Selector if this is a Call opcode
             let mut target_address = None;
+            let mut resolved_label = None;
+            
             if operation == "CALL" || operation == "STATICCALL" || operation == "DELEGATECALL" || operation == "CALLCODE" {
                 if let Some(stack) = &step.stack {
                     if stack.len() >= 2 {
-                        // In Geth/Anvil trace stack array, the end of the array is the top of the stack.
-                        // CALL takes: gas, address, value, argsOffset, argsLength, retOffset, retLength
+                        // Extract target address (second item from top)
                         let hex_addr = &stack[stack.len() - 2];
                         let clean_hex = hex_addr.trim_start_matches("0x");
-                        // EVM addresses are exactly 40 chars, padded to 64 chars in stack elements
                         if clean_hex.len() >= 40 {
                             let extracted = &clean_hex[clean_hex.len() - 40..];
                             target_address = Some(format!("0x{}", extracted));
                         }
                     }
+
+                    // Attempt to extract the 4-byte selector from Memory using Offset & Length
+                    let mut args_offset_idx = None;
+                    let mut args_length_idx = None;
+
+                    if operation == "CALL" || operation == "CALLCODE" {
+                        if stack.len() >= 5 {
+                            args_offset_idx = Some(stack.len() - 4);
+                            args_length_idx = Some(stack.len() - 5);
+                        }
+                    } else if (operation == "DELEGATECALL" || operation == "STATICCALL") && stack.len() >= 4 {
+                        args_offset_idx = Some(stack.len() - 3);
+                        args_length_idx = Some(stack.len() - 4);
+                    }
+
+                    if let (Some(off_idx), Some(len_idx)) = (args_offset_idx, args_length_idx) {
+                        let offset_str = stack[off_idx].trim_start_matches("0x");
+                        let len_str = stack[len_idx].trim_start_matches("0x");
+
+                        if let (Ok(offset), Ok(length)) = (
+                            usize::from_str_radix(offset_str, 16),
+                            usize::from_str_radix(len_str, 16)
+                        ) {
+                            if length >= 4 {
+                                if let Some(mem) = &step.memory {
+                                    let word_idx = offset / 32;
+                                    let byte_offset = offset % 32;
+                                    let hex_offset = byte_offset * 2; // Each byte is 2 hex chars
+
+                                    if let Some(word) = mem.get(word_idx) {
+                                        let clean_word = word.trim_start_matches("0x");
+                                        let selector_opt = if clean_word.len() >= hex_offset + 8 {
+                                            let selector = &clean_word[hex_offset..hex_offset + 8];
+                                            Some(format!("0x{}", selector))
+                                        } else if word_idx + 1 < mem.len() {
+                                            // The 4-byte selector spans across two memory boundary words
+                                            let p1 = &clean_word[hex_offset..];
+                                            let needed = 8 - p1.len();
+                                            let next_word = mem[word_idx + 1].trim_start_matches("0x");
+                                            if next_word.len() >= needed {
+                                                let p2 = &next_word[..needed];
+                                                Some(format!("0x{}{}", p1, p2))
+                                            } else { None }
+                                        } else { None };
+
+                                        // Try resolving the label
+                                        if let Some(sel) = selector_opt {
+                                            resolved_label = registry.resolve(target_address.as_deref(), Some(&sel));
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
                 }
             }
 
@@ -61,27 +117,28 @@ impl Aggregator {
             };
 
             // Accumulate gas cost and flags
-            let entry = stack_map.entry(stack_str).or_insert((0, 0, None, false));
+            let entry = stack_map.entry(stack_str).or_insert((0, 0, None, None, false));
             entry.0 += step.gas_cost;
             entry.1 = step.pc;
             if target_address.is_some() {
                 entry.2 = target_address;
             }
+            if resolved_label.is_some() {
+                entry.3 = resolved_label;
+            }
             if step.reverted {
-                entry.3 = true;
+                entry.4 = true;
             }
-            
-            // NOTE: Reverts naturally bubble up visually because if an internal call hits REVERT, 
-            // the specific reverting stack path gets the `entry.3 = true` flag. 
         }
 
         let mut stacks: Vec<CollapsedStack> = stack_map
             .into_iter()
-            .map(|(stack, (weight, pc, target_address, reverted))| CollapsedStack {
+            .map(|(stack, (weight, pc, target_address, resolved_label, reverted))| CollapsedStack {
                 stack,
                 weight,
                 last_pc: Some(pc),
                 target_address,
+                resolved_label,
                 reverted,
             })
             .collect();
@@ -144,5 +201,46 @@ mod tests {
         assert!(revert_stack.reverted);
         assert_eq!(revert_stack.weight, 200);
     }
-}
 
+    #[test]
+    fn test_aggregator_memory_selector_extraction() {
+        // Stack for CALL:
+        // gas, address, value, argsOffset, argsLength, retOffset, retLength
+        // Top of stack is at the end.
+        // We want argsOffset to be "0x20" (32 bytes), argsLength to be "0x04" (4 bytes)
+        // stack[len-4] = argsOffset
+        // stack[len-5] = argsLength
+        
+        let stack = vec![
+            "0x0".to_string(), // retLength
+            "0x0".to_string(), // retOffset
+            "0x4".to_string(), // argsLength
+            "0x20".to_string(), // argsOffset (byte 32)
+            "0x0".to_string(), // value
+            "0x0000000000000000000000001111111111111111111111111111111111111111".to_string(), // target address
+            "0x1000".to_string(), // gas
+        ];
+
+        // Memory array (32-byte chunks as 64-char hex strings)
+        // We set argsOffset = 32, so it looks in mem[1].
+        // "beforeInitialize" selector is 0x18a9d381. We'll pad the rest with zeroes.
+        let memory = vec![
+            "0000000000000000000000000000000000000000000000000000000000000000".to_string(), // word 0
+            "18a9d38100000000000000000000000000000000000000000000000000000000".to_string(), // word 1
+        ];
+
+        let steps = vec![
+            TraceStep { pc: 0, op: "CALL".into(), gas: 1000, gas_cost: 50, depth: 1, stack: Some(stack), memory: Some(memory), error: None, reverted: false },
+            TraceStep { pc: 1, op: "STOP".into(), gas: 900, gas_cost: 0, depth: 1, stack: None, memory: None, error: None, reverted: false },
+        ];
+
+        let stacks = Aggregator::build_collapsed_stacks(&steps);
+        let call_stack = stacks.iter().find(|s| s.stack == "CALL;CALL").expect("Should find CALL");
+        
+        // Ensure that the target address was resolved successfully
+        assert_eq!(call_stack.target_address.as_deref(), Some("0x1111111111111111111111111111111111111111"));
+        
+        // Ensure that the specific Uniswap v4 Hook was decoded
+        assert_eq!(call_stack.resolved_label.as_deref(), Some("Uniswapv4: beforeInitialize"));
+    }
+}