traverse_graph/
cg.rs

1use crate::interface_resolver::BindingRegistry; // Added
2use crate::manifest::Manifest; // Added
3use crate::parser::get_node_text;
4use anyhow::{anyhow, Context, Result}; // Add anyhow!
5use std::collections::{HashMap, HashSet}; // Import HashSet
6use std::iter;
7use streaming_iterator::StreamingIterator;
8use tree_sitter::{Node as TsNode, Query, QueryCursor, Tree};
9use tracing::debug;
10
11// Constants for synthetic node names
12pub(crate) const EVM_NODE_NAME: &str = "EVM";
13pub(crate) const EVENT_LISTENER_NODE_NAME: &str = "EventListener";
14pub(crate) const _REQUIRE_NODE_NAME: &str = "Require"; // Added for require statements
15pub(crate) const IF_CONDITION_NODE_NAME: &str = "IfCondition"; // Added for if statements
16pub(crate) const THEN_BLOCK_NODE_NAME: &str = "ThenBlock"; // Added for then blocks
17pub(crate) const ELSE_BLOCK_NODE_NAME: &str = "ElseBlock"; // Added for else blocks
18pub(crate) const WHILE_CONDITION_NODE_NAME: &str = "WhileCondition"; // Added for while statements
19pub(crate) const WHILE_BLOCK_NODE_NAME: &str = "WhileBlock"; // Added for while blocks
20pub(crate) const FOR_CONDITION_NODE_NAME: &str = "ForCondition"; // Added for for statements
21pub(crate) const FOR_BLOCK_NODE_NAME: &str = "ForBlock"; // Added for for blocks
22
23#[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd, serde::Serialize, serde::Deserialize)]
24pub enum EdgeType {
25    Call,
26    Return,
27    StorageRead,  // Represents reading from a storage variable
28    StorageWrite, // Represents writing to a storage variable
29    Require,      // Represents a require check
30    IfConditionBranch, // Represents the edge to an if-condition
31    ThenBranch,   // Represents the true branch of an if statement
32    ElseBranch,   // Represents the false branch of an if statement
33    WhileConditionBranch, // Represents the edge to a while-condition
34    WhileBodyBranch, // Represents the edge from a while-condition to its body
35    ForConditionBranch, // Represents the edge to a for-condition
36    ForBodyBranch, // Represents the edge from a for-condition to its body
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd, serde::Serialize, serde::Deserialize)]
40pub enum NodeType {
41    Function,
42    Interface,
43    Constructor,
44    Modifier,
45    Library,       // Added Library type
46    StorageVariable, // Represents a state variable in storage
47    Evm,           // Synthetic node for EVM interaction
48    EventListener, // Synthetic node for event listeners
49    RequireCondition, // Synthetic node for require checks
50    IfStatement,   // Synthetic node for an if condition
51    ThenBlock,     // Synthetic node for a then block
52    ElseBlock,     // Synthetic node for an else block
53    WhileStatement, // Synthetic node for a while condition
54    WhileBlock,    // Synthetic node for a while block
55    ForCondition, // Synthetic node for a for condition
56    ForBlock,     // Synthetic node for a for block
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd, serde::Serialize, serde::Deserialize)]
60pub enum Visibility {
61    Public,
62    Private,
63    Internal,
64    External,
65    Default,
66}
67
68#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
69pub struct Node {
70    pub id: usize,
71    pub name: String,
72    pub node_type: NodeType,
73    pub contract_name: Option<String>,
74    pub visibility: Visibility,
75    pub span: (usize, usize),
76    pub has_explicit_return: bool, // Added flag
77    pub declared_return_type: Option<String>, // Added: Store declared return type of the function
78    pub parameters: Vec<ParameterInfo>, // Added: Store parameters for function-like nodes
79    pub revert_message: Option<String>, // Added: Store revert message for RequireCondition nodes
80    pub condition_expression: Option<String>, // Added: Store raw condition expression for RequireCondition nodes
81}
82
83#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq, PartialOrd, Ord)] // Added derive for ParameterInfo
84pub struct ParameterInfo {
85    pub name: String,
86    pub param_type: String,
87    pub description: Option<String>, // Keep description if used by codegen
88}
89
90/// Information about a state variable mapping.
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct MappingInfo {
93    pub name: String,
94    pub visibility: Visibility,
95    pub key_types: Vec<String>, // For mapping(k1 => mapping(k2 => v)), this would be [k1_type_str, k2_type_str]
96    pub value_type: String,   // The final value type string
97    pub span: (usize, usize), // Span of the full state variable declaration
98    pub full_type_str: String, // e.g., "mapping(address => mapping(uint => bool))"
99}
100
101// --- DOT Label Implementation ---
102
103impl crate::cg_dot::ToDotLabel for Node {
104    fn to_dot_label(&self) -> String {
105        format!(
106            "{}{}\n({})", // Use a literal newline character here
107            self.contract_name
108                .as_deref()
109                .map(|c| format!("{}.", c))
110                .unwrap_or_default(),
111            self.name,
112            format!("{:?}", self.node_type).to_lowercase()
113        )
114    }
115}
116
117#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
118pub struct Edge {
119    pub source_node_id: usize,
120    pub target_node_id: usize,
121    pub edge_type: EdgeType,
122    pub call_site_span: (usize, usize),
123    pub return_site_span: Option<(usize, usize)>,
124    pub sequence_number: usize,
125    pub returned_value: Option<String>,
126    pub argument_names: Option<Vec<String>>, // Added: Store argument names/texts
127    pub event_name: Option<String>,          // Added: Store event name for emits
128    pub declared_return_type: Option<String>, // Added: Store declared return type for Return edges
129}
130
131impl crate::cg_dot::ToDotLabel for Edge {
132    fn to_dot_label(&self) -> String {
133        // Default label based on type, can be overridden in cg_dot.rs formatter
134        match self.edge_type {
135            EdgeType::Call => self.sequence_number.to_string(),
136            EdgeType::Return => "ret".to_string(),
137            EdgeType::StorageRead => "read".to_string(),
138            EdgeType::StorageWrite => "write".to_string(),
139            EdgeType::Require => "require".to_string(), // Add label for Require
140            EdgeType::IfConditionBranch => "if_cond".to_string(),
141            EdgeType::ThenBranch => "then".to_string(),
142            EdgeType::ElseBranch => "else".to_string(),
143            EdgeType::WhileConditionBranch => "while_cond".to_string(),
144            EdgeType::WhileBodyBranch => "while_body".to_string(),
145            EdgeType::ForConditionBranch => "for_cond".to_string(),
146            EdgeType::ForBodyBranch => "for_body".to_string(),
147        }
148    }
149}
150
151// --- Call Graph ---
152
153#[derive(Debug, Clone, Default)]
154pub struct CallGraph {
155    pub nodes: Vec<Node>,
156    pub edges: Vec<Edge>,
157    pub(crate) node_lookup: HashMap<(Option<String>, String), usize>,
158}
159
160impl CallGraph {
161    pub fn new() -> Self {
162        Default::default()
163    }
164
165    pub fn add_node(
166        &mut self,
167        name: String,
168        node_type: NodeType,
169        contract_name: Option<String>,
170        visibility: Visibility,
171        span: (usize, usize),
172    ) -> usize {
173        let id = self.nodes.len();
174        let node = Node {
175            id,
176            name: name.clone(),
177            node_type,
178            contract_name: contract_name.clone(),
179            visibility,
180            span,
181            has_explicit_return: false,
182            declared_return_type: None, // Initialize with None
183            parameters: Vec::new(), // Initialize parameters as empty
184            revert_message: None, // Initialize new field
185            condition_expression: None, // Initialize new field
186        };
187        self.nodes.push(node);
188
189        self.node_lookup.insert((contract_name, name), id);
190        id
191    }
192
193    pub fn add_edge(
194        &mut self,
195        source_node_id: usize,
196        target_node_id: usize,
197        edge_type: EdgeType,                      // Added type
198        call_site_span: (usize, usize),           // Span of the call site (or function for return)
199        return_site_span: Option<(usize, usize)>, // Span of the return site
200        sequence_number: usize,                   // Sequence for calls, maybe 0 for returns
201        returned_value: Option<String>,           // Added returned value
202        argument_names: Option<Vec<String>>,      // Added: Argument names/texts
203        event_name: Option<String>,               // Added: Event name for emits
204        declared_return_type: Option<String>,     // Added: Declared return type
205    ) {
206        debug!(
207            "Attempting to add edge: {} -> {} (Type: {:?}, Seq: {}, RetVal: {:?}, Args: {:?}, DeclRetType: {:?})",
208            source_node_id, target_node_id, edge_type, sequence_number, returned_value, argument_names, declared_return_type
209        );
210        let edge = Edge {
211            source_node_id,
212            target_node_id,
213            edge_type, // Store type
214            call_site_span,
215            return_site_span, // Store return span
216            sequence_number,
217            returned_value, // Store returned value
218            argument_names, // Store argument names
219            event_name,
220            declared_return_type, // Store declared return type
221        };
222        self.edges.push(edge);
223    }
224
225    /// Finds a node ID based on its name and potential contract scope.
226    /// Resolution logic:
227    /// 1. Look for the name within the `current_contract` scope.
228    /// 2. Look for the name as a free function (no contract scope).
229    /// 3. **Fallback:** Search across all known contract scopes.
230    /// TODO: Enhance resolution with proper type analysis for member access, inheritance, imports, etc.
231    ///       This current fallback is a simplification and may be ambiguous if multiple
232
233    pub fn iter_nodes(&self) -> impl Iterator<Item = &Node> {
234        self.nodes.iter()
235    }
236
237    pub fn iter_edges(&self) -> impl Iterator<Item = &Edge> {
238        self.edges.iter()
239    }
240
241    // Remove 'a, change ctx type, add input parameter
242    pub fn add_explicit_return_edges(
243        &mut self,
244        input: &CallGraphGeneratorInput, // Add input parameter containing source, tree, lang
245        ctx: &CallGraphGeneratorContext, // Remove 'a
246    ) -> Result<()> {
247        debug!(
248            "[Address Debug add_explicit_return_edges START] Graph: {:p}, Total Edges: {}",
249            self,
250            self.edges.len()
251        );
252        let mut found_node4_edges = false;
253        for (idx, edge) in self.edges.iter().enumerate() {
254            if edge.source_node_id == 4 {
255                // Node ID for UniswapV2ERC20._transfer
256                let target_node_name =
257                    self.nodes
258                        .get(edge.target_node_id)
259                        .map_or("?".to_string(), |n| {
260                            format!(
261                                "{}.{}",
262                                n.contract_name.as_deref().unwrap_or("Global"),
263                                n.name
264                            )
265                        });
266                debug!(
267                    "Edge Index {}: {} -> {} (Target: '{}'), Type: {:?}, Seq: {}",
268                    idx, edge.source_node_id, edge.target_node_id, target_node_name, edge.edge_type, edge.sequence_number
269                );
270                found_node4_edges = true;
271            }
272        }
273        if !found_node4_edges {
274            debug!("No edges found originating from Node 4.");
275        }
276
277        // 1. Build a map of callee -> Vec<(caller_id, call_sequence_number)> from existing Call edges
278        let mut callee_to_callers_and_seq: HashMap<usize, Vec<(usize, usize)>> = HashMap::new();
279        for edge in self.edges.iter() {
280            if edge.edge_type == EdgeType::Call {
281                callee_to_callers_and_seq
282                    .entry(edge.target_node_id)
283                    .or_default()
284                    .push((edge.source_node_id, edge.sequence_number)); // Store caller and sequence
285            }
286        }
287
288        // 2. Prepare return statement query (Revised)
289        // This query looks for return statements possibly containing an expression.
290        let return_query_str = r#"
291            (return_statement
292                (expression)? @return_value ; Optional: capture the returned expression node
293            ) @return ; Capture the whole return statement node
294        "#;
295        // Use language from input
296        let return_query = Query::new(&input.solidity_lang, return_query_str)
297            .context("Failed to create return statement query")?;
298        let mut return_cursor = QueryCursor::new();
299        // Use source from input
300        let source_bytes = input.source.as_bytes();
301
302        // 3. Iterate through definition nodes to find return statements within them
303        let mut new_return_edges: Vec<Edge> = Vec::new(); // Collect new edges separately
304        let mut total_returns_found_by_query = 0; // DEBUG
305        let mut total_returns_processed = 0; // DEBUG
306        let mut _nodes_with_explicit_return_set = 0; // DEBUG - Mark as unused for now
307
308        // Iterate through the definition nodes stored in the context (contains NodeInfo)
309        for (callee_node_id, callee_node_info, _caller_contract_name_opt) in
310            &ctx.definition_nodes_info
311        {
312            // Get the actual TsNode for the callee's definition using span from NodeInfo
313            let definition_ts_node = input.tree.root_node()
314                .descendant_for_byte_range(callee_node_info.span.0, callee_node_info.span.1)
315                .ok_or_else(|| anyhow!("Failed to find definition TsNode for span {:?} in add_explicit_return_edges", callee_node_info.span))?;
316
317            // Retrieve the corresponding Node struct from the graph
318            let callee_node_exists = self.nodes.get(*callee_node_id).is_some();
319            if !callee_node_exists {
320                debug!(
321                    "Warning: Node ID {} found in definition_nodes_info but not in graph.nodes",
322                    callee_node_id
323                );
324                continue; // Skip if node not found in graph
325            }
326
327            // Use immutable borrow here for logging before potential mutable borrow
328            if let Some(callee_node_for_log) = self.nodes.get(*callee_node_id) {
329                debug!(
330                    "Processing callee: Node ID {}, Name: {}.{}, Type: {:?}",
331                    callee_node_for_log.id,
332                    callee_node_for_log
333                        .contract_name
334                        .as_deref()
335                        .unwrap_or("Global"),
336                    callee_node_for_log.name,
337                    callee_node_for_log.node_type
338                );
339            }
340
341            if *callee_node_id == 21 {
342                debug!(
343                    "S-expression for Node ID {}:\n{}",
344                    callee_node_id,
345                    definition_ts_node.to_sexp()
346                );
347            }
348
349            // Check node type and update Node struct (has_explicit_return and declared_return_type)
350            let node_type_for_check = self.nodes[*callee_node_id].node_type.clone();
351            match node_type_for_check {
352                NodeType::Function
353                | NodeType::Modifier
354                | NodeType::Constructor
355                | NodeType::Library => {
356                    // Get declared return type ONCE for this callee, ONLY if it's a callable type
357                    // This is used to update the Node and also for any Return Edges created later.
358                    let actual_declared_ret_type = get_function_return_type(*callee_node_id, ctx, input);
359
360                    // Update the Node struct in the graph
361                    if let Some(node_mut) = self.nodes.get_mut(*callee_node_id) {
362                        // Set has_explicit_return
363                        if !node_mut.has_explicit_return {
364                            let mut return_check_matches = return_cursor.matches(
365                                &return_query,
366                                definition_ts_node,
367                                |node: TsNode| iter::once(&source_bytes[node.byte_range()]),
368                            );
369                            return_check_matches.advance(); // Advance once
370                            if return_check_matches.get().is_some() {
371                                node_mut.has_explicit_return = true;
372                                _nodes_with_explicit_return_set += 1;
373                                debug!(
374                                     "Set has_explicit_return=true for Node ID {}",
375                                     *callee_node_id
376                                );
377                            }
378                        }
379                        // Set declared_return_type (clone from the one fetched above)
380                        node_mut.declared_return_type = actual_declared_ret_type.clone();
381                        if actual_declared_ret_type.is_some() {
382                             debug!(
383                                "Set declared_return_type='{:?}' for Node ID {}",
384                                actual_declared_ret_type,
385                                *callee_node_id
386                            );
387                        }
388                    }
389
390                    // Find callers and their call sequences for this callee
391                    if let Some(callers_info) = callee_to_callers_and_seq.get(callee_node_id) {
392                        let mut matches = return_cursor.matches(
393                            &return_query,
394                            definition_ts_node,
395                            |node: TsNode| iter::once(&source_bytes[node.byte_range()]),
396                        );
397                        matches.advance(); // Need to advance once initially
398
399                        while let Some(match_) = matches.get() {
400                            total_returns_found_by_query += 1;
401
402                            let mut return_node_opt: Option<TsNode> = None;
403                            let mut return_value_node_opt: Option<TsNode> = None;
404
405                            for capture in match_.captures {
406                                let capture_name =
407                                    &return_query.capture_names()[capture.index as usize];
408                                match capture_name.as_ref() {
409                                    "return" => return_node_opt = Some(capture.node),
410                                    "return_value" => return_value_node_opt = Some(capture.node),
411                                    _ => {}
412                                }
413                            }
414
415                            if let Some(return_node) = return_node_opt {
416                                total_returns_processed += 1;
417                                let return_span =
418                                    (return_node.start_byte(), return_node.end_byte());
419                                let return_kind = return_node.kind();
420
421                                debug!(
422                                    "  Found return statement within definition: Kind='{}', Span={:?}. => ACCEPTED (within definition node)",
423                                    return_kind, return_span
424                                );
425
426                                let returned_value_text = return_value_node_opt
427                                    .map(|n| get_node_text(&n, &input.source).to_string());
428
429                                for (caller_id, call_sequence) in callers_info {
430                                    let callee_node_span = self.nodes[*callee_node_id].span;
431                                    // Use actual_declared_ret_type obtained earlier for the edge
432                                    new_return_edges.push(Edge {
433                                        source_node_id: *callee_node_id,
434                                        target_node_id: *caller_id,
435                                        edge_type: EdgeType::Return,
436                                        call_site_span: callee_node_span, // Span of the function definition itself
437                                        return_site_span: Some(return_span), // Span of the return statement
438                                        sequence_number: *call_sequence, // Use the sequence number from the original call
439                                        returned_value: returned_value_text.clone(), // Add the returned value text
440                                        argument_names: None, // Return edges don't have call arguments
441                                        event_name: None,
442                                        declared_return_type: actual_declared_ret_type.clone(), // Use the fetched type
443                                    });
444                                }
445                            } else {
446                                // This case should not happen if the query is correct and finds a match
447                                debug!("  Warning: Query matched but failed to extract @return capture.");
448                            }
449                            matches.advance(); // Advance after processing captures for the current match
450                        }
451                    }
452                }
453                _ => {} // Ignore other node types like Interface, ContractDefinition etc.
454            }
455        }
456
457        // 4. Add the collected return edges to the graph
458        debug!(
459            "Total returns found by query: {}", // DEBUG Summary
460            total_returns_found_by_query
461        );
462        debug!(
463            "Total returns processed (found within definition): {}", // DEBUG Summary
464            total_returns_processed
465        );
466        debug!(
467            "Total return edges generated: {}", // DEBUG Summary
468            new_return_edges.len()
469        );
470        debug!(
471            "Edge count BEFORE extend: {}",
472            self.edges.len()
473        );
474        debug!(
475            "[Address Debug add_explicit_return_edges] Graph: {:p}",
476            self
477        ); // Added address log
478        self.edges.extend(new_return_edges);
479        debug!(
480            "Edge count AFTER extend: {}",
481            self.edges.len()
482        );
483
484        Ok(())
485    }
486
487    /// Retrieves the list of ancestor contracts for a given contract name,
488    /// including the contract itself, ordered from most specific to most general.
489    /// Uses the inheritance information stored in the context.
490    pub(crate) fn get_ancestor_contracts(
491        &self,
492        contract_name: &str,
493        ctx: &CallGraphGeneratorContext,
494    ) -> Vec<String> {
495        let mut ancestors = Vec::new();
496        let mut queue = std::collections::VecDeque::new();
497        let mut visited = HashSet::new(); // Prevent cycles and redundant lookups
498
499        queue.push_back(contract_name.to_string());
500        visited.insert(contract_name.to_string());
501
502        while let Some(current_contract) = queue.pop_front() {
503            // Add the current contract to the front of the ancestors list
504            // to maintain the order from specific to general after reversal.
505            ancestors.push(current_contract.clone());
506
507            // Look up direct parents (contracts and interfaces)
508            if let Some(parents) = ctx.contract_inherits.get(&current_contract) {
509                for parent_name in parents {
510                    if visited.insert(parent_name.clone()) {
511                        queue.push_back(parent_name.clone());
512                    }
513                }
514            }
515            // Also consider inherited interfaces if necessary, though storage vars are usually in contracts.
516            // If interfaces could define constants used in storage access, this might be needed.
517            // if let Some(parents) = ctx.interface_inherits.get(&current_contract) { ... }
518        }
519
520        // The queue processing naturally explores breadth-first, but the order
521        // we add to `ancestors` and the final reversal ensures linearization
522        // (though Solidity's C3 linearization is more complex, this provides
523        // a reasonable lookup order: self -> parents -> grandparents...).
524        ancestors.reverse(); // Reverse to get the order: self, parent, grandparent...
525        ancestors
526    }
527}
528
529#[derive(Debug)]
530pub struct CallGraphGeneratorInput {
531    pub source: String,
532    pub tree: Tree,
533    pub solidity_lang: tree_sitter::Language,
534}
535
536// Implement Clone manually to ensure proper ownership transfer
537impl Clone for CallGraphGeneratorInput {
538    fn clone(&self) -> Self {
539        Self {
540            source: self.source.clone(),
541            tree: self.tree.clone(),
542            solidity_lang: self.solidity_lang.clone(),
543        }
544    }
545}
546
547/// Stores lifetime-independent information about a tree-sitter node.
548#[derive(Debug, Clone)]
549pub struct NodeInfo {
550    pub span: (usize, usize),
551    pub kind: String,
552    // Add other relevant lifetime-independent info if needed (e.g., text?)
553}
554
555
556// Remove 'a lifetime - Context now stores NodeInfo which owns its data
557#[derive(Debug, Clone, Default)] // Add Default
558pub struct CallGraphGeneratorContext {
559    pub state_var_types: HashMap<(String, String), String>,
560    pub using_for_directives: HashMap<(Option<String>, String), Vec<String>>,
561    // Store NodeInfo instead of TsNode<'a>
562    pub definition_nodes_info: Vec<(usize, NodeInfo, Option<String>)>,
563    pub all_contracts: HashMap<String, NodeInfo>,
564    pub contracts_with_explicit_constructors: HashSet<String>,
565    pub all_libraries: HashMap<String, NodeInfo>,
566    pub all_interfaces: HashMap<String, NodeInfo>,
567    pub interface_functions: HashMap<String, Vec<String>>,
568    pub contract_implements: HashMap<String, Vec<String>>, // Contract -> List of Interfaces it implements
569    pub interface_inherits: HashMap<String, Vec<String>>, // Interface -> List of Interfaces it inherits from
570    pub contract_inherits: HashMap<String, Vec<String>>, // Contract -> List of Contracts/Interfaces it inherits from
571    pub storage_var_nodes: HashMap<(Option<String>, String), usize>, // (ContractScope, VarName) -> Node ID
572    /// Stores detailed information about declared mappings.
573    /// Key: (contract_name, mapping_variable_name)
574    pub contract_mappings: HashMap<(String, String), MappingInfo>,
575    // Added for interface binding resolution
576    pub manifest: Option<Manifest>,
577    pub binding_registry: Option<BindingRegistry>,
578}
579
580
581pub trait CallGraphGeneratorStep {
582    fn name(&self) -> &'static str;
583    /// Configure the step with settings.
584    fn config(&mut self, config: &HashMap<String, String>);
585    /// Generate part of the call graph.
586    fn generate(
587        &self,
588        input: CallGraphGeneratorInput,
589        ctx: &mut CallGraphGeneratorContext, // Remove 'a
590        graph: &mut CallGraph,
591    ) -> Result<()>;
592}
593
594
595pub(crate) fn extract_function_parameters(
596    fn_like_ts_node: TsNode, // This is function_definition or constructor_definition
597    source: &str,
598) -> Vec<ParameterInfo> {
599    let mut parameters = Vec::new();
600    debug!("[cg::extract_function_parameters DEBUG] Analyzing TsNode kind: '{}', text snippet: '{}'", fn_like_ts_node.kind(), get_node_text(&fn_like_ts_node, source).chars().take(70).collect::<String>());
601
602    let mut child_cursor = fn_like_ts_node.walk();
603    // Iterate through direct children of fn_like_ts_node (constructor_definition or function_definition)
604    // The 'parameter' nodes are direct children, not nested under a 'parameter_list' node.
605    for child_node in fn_like_ts_node.children(&mut child_cursor) {
606        debug!("[cg::extract_function_parameters DEBUG] Child of fn_like_ts_node: Kind='{}', Text='{}'", child_node.kind(), get_node_text(&child_node, source).chars().take(30).collect::<String>());
607        if child_node.kind() == "parameter" { // These are the actual parameter definitions
608            let type_node = child_node.child_by_field_name("type");
609            let name_node = child_node.child_by_field_name("name");
610
611            // Both type and name are expected for function/constructor parameters
612            if let (Some(tn), Some(nn)) = (type_node, name_node) {
613                let param_name = crate::parser::get_node_text(&nn, source).to_string();
614                let param_type = crate::parser::get_node_text(&tn, source).to_string();
615                debug!("[cg::extract_function_parameters DEBUG]   Extracted param: name='{}', type='{}'", param_name, param_type);
616                parameters.push(ParameterInfo {
617                    name: param_name,
618                    param_type: param_type,
619                    description: None,
620                });
621            } else {
622                debug!("[cg::extract_function_parameters DEBUG]   Found 'parameter' node (Kind: '{}', Text: '{}') but missing type or name field.", child_node.kind(), get_node_text(&child_node, source).chars().take(30).collect::<String>());
623            }
624        }
625    }
626
627    // Log if no parameters were found, which might indicate an issue if parameters were expected.
628    if parameters.is_empty() {
629        // Check if the function/constructor signature actually has parameters in the source
630        // This is a heuristic based on text, as we've already failed to parse them structurally.
631        let signature_text = get_node_text(&fn_like_ts_node, source);
632        if signature_text.contains('(') && signature_text.contains(')') && !signature_text.contains("()") {
633             // A more robust check would be to see if there are any 'parameter' kind nodes between '(' and ')'
634            let mut has_parameter_nodes_in_signature = false;
635            let mut temp_cursor = fn_like_ts_node.walk();
636            for child in fn_like_ts_node.children(&mut temp_cursor) {
637                if child.kind() == "parameter" {
638                    has_parameter_nodes_in_signature = true;
639                    break;
640                }
641            }
642
643            if has_parameter_nodes_in_signature {
644                 debug!("[cg::extract_function_parameters DEBUG] No parameters extracted, but 'parameter' nodes were found as direct children. This indicates the loop logic might be incorrect or tree-sitter grammar differs from expectation.");
645            } else {
646                 debug!("[cg::extract_function_parameters DEBUG] No parameters extracted, and no 'parameter' nodes found as direct children. This might be a function/constructor with no parameters, or an issue with identifying 'parameter' nodes.");
647            }
648        } else {
649            debug!("[cg::extract_function_parameters DEBUG] No parameters extracted. This appears to be a function/constructor with no parameters based on signature text: '{}'", signature_text.chars().take(50).collect::<String>());
650        }
651    }
652
653    debug!("[cg::extract_function_parameters DEBUG] For node kind '{}', extracted {} parameters: {:?}", fn_like_ts_node.kind(), parameters.len(), parameters);
654    parameters
655}
656
657pub(crate) fn extract_arguments<'a>(
658    call_expr_node: TsNode<'a>,
659    input: &CallGraphGeneratorInput,
660) -> Vec<String> {
661    let mut argument_texts: Vec<String> = Vec::new();
662    debug!(
663        "[cg::extract_arguments DEBUG] Extracting argument texts for Call Expr Node: Kind='{}', Span=({:?})",
664        call_expr_node.kind(),
665        (call_expr_node.start_byte(), call_expr_node.end_byte())
666    );
667
668    if let Some(arguments_field_node) = call_expr_node.child_by_field_name("arguments") {
669        debug!("[cg::extract_arguments DEBUG]   Found 'arguments' field. Iterating its children.");
670        let mut arg_cursor = arguments_field_node.walk();
671        for arg_node in arguments_field_node.children(&mut arg_cursor) { // Iterate children of 'arguments'
672            // These children should be 'call_argument' nodes, which are expressions
673            if arg_node.kind() == "call_argument" { // Ensure it's a call_argument
674                let arg_text = get_node_text(&arg_node, &input.source).to_string();
675                debug!("[cg::extract_arguments DEBUG]     Extracted argument text (from 'arguments' field, child is call_argument): '{}'", arg_text);
676                argument_texts.push(arg_text);
677            }
678        }
679    } else {
680        // If 'arguments' field is not found, iterate direct children of call_expr_node
681        // and look for 'call_argument' nodes. This handles cases like `require(arg)`.
682        debug!("[cg::extract_arguments DEBUG]   No 'arguments' field found. Iterating direct children of call_expression for 'call_argument' nodes.");
683        let mut direct_child_cursor = call_expr_node.walk();
684        for child_node in call_expr_node.children(&mut direct_child_cursor) {
685            if child_node.kind() == "call_argument" {
686                // child_node is the 'call_argument', which is an _expression node.
687                let arg_text = get_node_text(&child_node, &input.source).to_string();
688                debug!("[cg::extract_arguments DEBUG]     Extracted argument text (direct child call_argument): '{}'", arg_text);
689                argument_texts.push(arg_text);
690            }
691        }
692        if argument_texts.is_empty() {
693             debug!("[cg::extract_arguments DEBUG]   No 'call_argument' nodes found as direct children either (after checking 'arguments' field).");
694        }
695    }
696    argument_texts
697}
698
699pub(crate) fn extract_argument_nodes<'a>(call_expr_node: TsNode<'a>) -> Vec<TsNode<'a>> {
700    let mut argument_nodes_ts: Vec<TsNode<'a>> = Vec::new();
701    debug!(
702        "[cg::extract_argument_nodes DEBUG] Extracting argument nodes for Call Expr Node: Kind='{}', Span=({:?})",
703        call_expr_node.kind(),
704        (call_expr_node.start_byte(), call_expr_node.end_byte())
705    );
706
707    if let Some(arguments_field_node) = call_expr_node.child_by_field_name("arguments") {
708        debug!("[cg::extract_argument_nodes DEBUG]   Found 'arguments' field. Iterating its children.");
709        let mut arg_cursor = arguments_field_node.walk();
710        for arg_node in arguments_field_node.children(&mut arg_cursor) { // Iterate children of 'arguments'
711            if arg_node.kind() == "call_argument" { // Ensure it's a call_argument
712                debug!("[cg::extract_argument_nodes DEBUG]     Extracted argument node (from 'arguments' field, child is call_argument): Kind='{}', Span=({:?})", arg_node.kind(), (arg_node.start_byte(), arg_node.end_byte()));
713                argument_nodes_ts.push(arg_node);
714            }
715        }
716    } else {
717        debug!("[cg::extract_argument_nodes DEBUG]   No 'arguments' field found. Iterating direct children of call_expression for 'call_argument' nodes.");
718        let mut direct_child_cursor = call_expr_node.walk();
719        for child_node in call_expr_node.children(&mut direct_child_cursor) {
720            if child_node.kind() == "call_argument" {
721                // child_node is the 'call_argument', which is an _expression node.
722                debug!("[cg::extract_argument_nodes DEBUG]     Extracted argument node (direct child call_argument): Kind='{}', Span=({:?})", child_node.kind(), (child_node.start_byte(), child_node.end_byte()));
723                argument_nodes_ts.push(child_node);
724            }
725        }
726        if argument_nodes_ts.is_empty() {
727            debug!("[cg::extract_argument_nodes DEBUG]   No 'call_argument' nodes found as direct children either (after checking 'arguments' field).");
728        }
729    }
730    argument_nodes_ts
731}
732
733/// Attempts to parse the return type string from a function definition node.
734/// TODO: Handle multiple return types (tuples). Currently returns the first found or the tuple itself as string.
735// Remove 'a, change ctx type, input is already present
736pub(crate) fn get_function_return_type(
737    target_node_id: usize,
738    ctx: &CallGraphGeneratorContext,
739    input: &CallGraphGeneratorInput,
740) -> Option<String> {
741    // Find the NodeInfo corresponding to the target_node_id
742    let definition_node_info_opt = ctx
743        .definition_nodes_info
744        .iter()
745        .find(|(id, _, _)| *id == target_node_id)
746        .map(|(_, node_info, _)| node_info); // Extract the NodeInfo
747
748    if let Some(definition_node_info) = definition_node_info_opt {
749         // Get the actual TsNode using the span from NodeInfo and the input tree
750        let definition_ts_node = match input.tree.root_node()
751            .descendant_for_byte_range(definition_node_info.span.0, definition_node_info.span.1) {
752            Some(node) => node,
753            None => {
754                debug!("[Return Type Parse DEBUG] Failed to find TsNode for span {:?} for node ID {}", definition_node_info.span, target_node_id);
755                return None;
756            }
757        };
758
759        debug!("[Return Type Parse DEBUG] Analyzing definition node ID {} for return type. S-Expr:\n{}", target_node_id, definition_ts_node.to_sexp()); // Use retrieved node
760
761        // Query to find the actual type name node within the standard return structure.
762        // NOTE: This simplified query only handles `returns (type)`.
763        // It does NOT handle named return parameters like `returns (uint value)`
764        // or returns from modifiers. This is sufficient for the current test case.
765        // Added pattern for interface function definitions.
766        let return_type_query_str = r#"
767            [
768              (function_definition
769                return_type: (return_type_definition
770                  (parameter
771                    type: (type_name) @return_type_name_node
772                  )
773                )
774              )
775              (interface_declaration (_                       ;; Same for interfaces
776                (function_definition
777                  return_type: (return_type_definition
778                    (parameter
779                      type: (type_name) @return_type_name_node
780                    )
781                  )
782                )
783              ))
784            ]
785        "#;
786        // Note: Using a static query cache might be more efficient if this runs often.
787        let return_type_query = match Query::new(&input.solidity_lang, return_type_query_str) {
788            Ok(q) => q,
789            Err(e) => {
790                debug!(
791                    "[Return Type Parse DEBUG] Failed to create query for node ID {}: {}",
792                    target_node_id, e
793                );
794                return None;
795            }
796        };
797        debug!(
798            "[Return Type Parse DEBUG] Query created successfully for node ID {}.",
799            target_node_id
800        );
801
802        let mut cursor = QueryCursor::new();
803        let source_bytes = input.source.as_bytes();
804
805        // Run the query only on the specific definition node
806        debug!("[Return Type Parse DEBUG] Running matches on definition node...");
807        let mut matches = cursor.matches(
808            &return_type_query,
809            definition_ts_node, // Use retrieved node
810            |node: TsNode| iter::once(&source_bytes[node.byte_range()]),
811        );
812        matches.advance(); // Advance once
813
814        debug!("[Return Type Parse DEBUG] Checking for match result..."); // New log
815
816        if let Some(match_) = matches.get() {
817            debug!(
818                "[Return Type Parse DEBUG] Match found! Processing {} captures...",
819                match_.captures.len()
820            ); // New log
821               // Find the @return_type_name_node capture
822            let mut found_expected_capture = false; // Flag to track if we found the right capture
823            for (cap_idx, capture) in match_.captures.iter().enumerate() {
824                // Added index for logging
825                let capture_name = &return_type_query.capture_names()[capture.index as usize];
826                debug!(
827                    "[Return Type Parse DEBUG]   Capture {}: Name='{}', Node Kind='{}', Text='{}'",
828                    cap_idx,
829                    capture_name,
830                    capture.node.kind(),
831                    get_node_text(&capture.node, &input.source)
832                ); // New log
833
834                if *capture_name == "return_type_name_node" {
835                    // Use the new capture name
836                    debug!("[Return Type Parse DEBUG]     Match on '@return_type_name_node'!"); // New log
837                    found_expected_capture = true; // Mark that we found it
838                    let type_name_node = capture.node;
839                    // Get the text of the type_name node directly
840                    let type_name_text = get_node_text(&type_name_node, &input.source).to_string();
841
842                    if !type_name_text.is_empty() {
843                        debug!(
844                            "[Return Type Parse DEBUG] Found single return type name: '{}'",
845                            type_name_text
846                        );
847                        return Some(type_name_text); // Return the captured type name text
848                    } else {
849                        debug!("[Return Type Parse DEBUG] Found empty return type name node.");
850                        // Let it fall through to return None at the end if text is empty
851                    }
852                    // TODO: Handle multiple return types if the query is extended later
853                } else {
854                    debug!("[Return Type Parse DEBUG]     Capture name '{}' did not match expected '@return_type_name_node'.", capture_name);
855                    // New log
856                }
857            }
858            // If the loop finishes without returning, check the flag
859            if !found_expected_capture {
860                debug!("[Return Type Parse DEBUG] Loop finished. Query matched but the '@return_type_name_node' capture was not found among the captures for node ID {}.", target_node_id);
861            } else {
862                debug!("[Return Type Parse DEBUG] Loop finished. Found '@return_type_name_node' capture but it resulted in empty text or didn't return for node ID {}.", target_node_id);
863            }
864        } else {
865            debug!(
866                "[Return Type Parse DEBUG] Query found no return type match for node ID {}.",
867                target_node_id
868            ); // Refined log
869        }
870    } else {
871        debug!(
872            "[Return Type Parse DEBUG] Definition TsNode not found for node ID {}",
873            target_node_id
874        );
875    }
876
877    debug!(
878        "[Return Type Parse DEBUG] Function returning None for node ID {}.",
879        target_node_id
880    ); // New log before returning None
881    None // Return type not found or parsing failed
882}
883
884// Remove 'a lifetime
885pub struct CallGraphGeneratorPipeline {
886    steps: Vec<Box<dyn CallGraphGeneratorStep>>, // Remove 'a
887    enabled_steps: HashSet<String>, // Track enabled step names
888    // Remove _marker
889}
890
891// Implement Default for Pipeline
892impl Default for CallGraphGeneratorPipeline {
893    fn default() -> Self {
894        Self::new()
895    }
896}
897
898
899// Remove 'a lifetime
900impl CallGraphGeneratorPipeline {
901    pub fn new() -> Self {
902        Self {
903            steps: Vec::new(),
904            enabled_steps: HashSet::new(), // Initialize the set
905            // Remove _marker
906        }
907    }
908
909    /// Adds a step to the pipeline. Steps are enabled by default.
910    pub fn add_step(&mut self, step: Box<dyn CallGraphGeneratorStep>) { // Remove 'a
911        self.enabled_steps.insert(step.name().to_string()); // Enable by default
912        self.steps.push(step);
913    }
914
915    /// Enables a step by its name.
916    pub fn enable_step(&mut self, name: &str) {
917        self.enabled_steps.insert(name.to_string());
918    }
919
920    /// Disables a step by its name.
921    pub fn disable_step(&mut self, name: &str) {
922        self.enabled_steps.remove(name);
923    }
924
925    pub fn run(
926        &mut self,
927        input: CallGraphGeneratorInput,
928        ctx: &mut CallGraphGeneratorContext, // Remove 'a
929        graph: &mut CallGraph,
930        config: &HashMap<String, String>, // Keep config map parameter for run
931    ) -> Result<()> {
932        // First pass: configure enabled steps
933        for step in self.steps.iter_mut() {
934            if self.enabled_steps.contains(step.name()) {
935                step.config(config);
936            }
937        }
938        // Second pass: generate using enabled steps
939        for step in &self.steps {
940            if self.enabled_steps.contains(step.name()) {
941                debug!(
942                    "[Address Debug Pipeline] Running step '{}', Graph: {:p}",
943                    step.name(),
944                    graph
945                ); // Added address log
946                   // Call generate only if the step is enabled
947                let step_input = input.clone(); // Clone input for each step
948                step.generate(step_input, ctx, graph)?;
949            }
950        }
951        Ok(())
952    }
953}
traverse_graph/cg.rs

traverse_graph/
cg.rs
