bat-cli 0.34.0

use std::collections::HashMap;

use solar_parse::{
    ast,
    interface::{source_map::FileName, Session},
    Parser,
};

use crate::batbelt::evm::types::{EvmContract, EvmFunction};

/// Represents a resolved function call.
#[derive(Debug, Clone)]
pub struct ResolvedCall {
    pub caller_contract: String,
    pub caller_function: String,
    pub callee_contract: String,
    pub callee_function: String,
    pub is_external: bool,
    pub is_super: bool,
}

/// Resolves function calls within and across contracts.
pub struct CallResolver<'a> {
    contracts_by_name: HashMap<String, &'a EvmContract>,
}

impl<'a> CallResolver<'a> {
    pub fn new(contracts: &'a [EvmContract]) -> Self {
        let contracts_by_name = contracts.iter().map(|c| (c.name.clone(), c)).collect();
        Self { contracts_by_name }
    }

    /// Resolve all calls in a function's body using AST-based extraction.
    /// Falls back to regex if AST parsing fails.
    pub fn resolve_calls(&self, contract_name: &str, function: &EvmFunction) -> Vec<ResolvedCall> {
        let body = &function.body_source;
        if body.is_empty() {
            return Vec::new();
        }

        // Try AST-based extraction first
        let call_names = extract_calls_from_source(body);

        let mut calls = Vec::new();

        for callee_name in &call_names {
            if is_builtin(callee_name) {
                continue;
            }

            // Check for external calls: "Contract.function"
            if let Some(dot_pos) = callee_name.find('.') {
                let target = &callee_name[..dot_pos];
                let method = &callee_name[dot_pos + 1..];

                if target == "super" {
                    calls.push(ResolvedCall {
                        caller_contract: contract_name.to_string(),
                        caller_function: function.name.clone(),
                        callee_contract: contract_name.to_string(),
                        callee_function: method.to_string(),
                        is_external: false,
                        is_super: true,
                    });
                } else if self.contracts_by_name.contains_key(target) {
                    calls.push(ResolvedCall {
                        caller_contract: contract_name.to_string(),
                        caller_function: function.name.clone(),
                        callee_contract: target.to_string(),
                        callee_function: method.to_string(),
                        is_external: true,
                        is_super: false,
                    });
                }
                continue;
            }

            // Internal call
            if let Some(resolved) =
                self.resolve_single_call(contract_name, &function.name, callee_name)
            {
                calls.push(resolved);
            }
        }

        calls
    }

    fn resolve_single_call(
        &self,
        contract_name: &str,
        caller_function: &str,
        callee_name: &str,
    ) -> Option<ResolvedCall> {
        // Check internal functions in same contract
        if let Some(contract) = self.contracts_by_name.get(contract_name) {
            if contract.functions.iter().any(|f| f.name == callee_name) {
                return Some(ResolvedCall {
                    caller_contract: contract_name.to_string(),
                    caller_function: caller_function.to_string(),
                    callee_contract: contract_name.to_string(),
                    callee_function: callee_name.to_string(),
                    is_external: false,
                    is_super: false,
                });
            }
        }

        // Check inherited functions
        if let Some(contract) = self.contracts_by_name.get(contract_name) {
            for base_name in &contract.base_contracts {
                if let Some(base) = self.contracts_by_name.get(base_name.as_str()) {
                    if base.functions.iter().any(|f| f.name == callee_name) {
                        return Some(ResolvedCall {
                            caller_contract: contract_name.to_string(),
                            caller_function: caller_function.to_string(),
                            callee_contract: base_name.clone(),
                            callee_function: callee_name.to_string(),
                            is_external: false,
                            is_super: false,
                        });
                    }
                }
            }
        }

        None
    }
}

/// Extract function call names from Solidity source code using AST.
/// Falls back to regex if AST parsing fails.
pub fn extract_calls_from_source(source: &str) -> Vec<String> {
    // Wrap in a dummy function so it parses as a valid Solidity file
    let wrapped = format!("contract _C {{ function _f() {{ {} }} }}", source);

    let sess = Session::builder().with_silent_emitter(None).build();

    let result = sess.enter(|| -> Option<Vec<String>> {
        let arena = ast::Arena::new();
        let mut parser = Parser::from_source_code(
            &sess,
            &arena,
            FileName::Custom("call_resolver".into()),
            wrapped.clone(),
        )
        .ok()?;

        let file = parser.parse_file().map_err(|e| e.emit()).ok()?;

        let mut calls = Vec::new();
        // Navigate: file > contract > function > body > stmts
        for item in file.items.iter() {
            if let ast::ItemKind::Contract(c) = &item.kind {
                for body_item in c.body.iter() {
                    if let ast::ItemKind::Function(f) = &body_item.kind {
                        if let Some(block) = &f.body {
                            for stmt in block.stmts.iter() {
                                extract_calls_from_stmt(&stmt.kind, &mut calls);
                            }
                        }
                    }
                }
            }
        }
        calls.sort();
        calls.dedup();
        Some(calls)
    });

    result.unwrap_or_else(|| extract_calls_regex(source))
}

/// AST walk: extract call names from a statement kind.
fn extract_calls_from_stmt(kind: &ast::StmtKind<'_>, calls: &mut Vec<String>) {
    match kind {
        ast::StmtKind::Expr(expr) => {
            extract_calls_from_expr(&expr.kind, calls);
        }
        ast::StmtKind::Return(opt_expr) => {
            if let Some(expr) = opt_expr {
                extract_calls_from_expr(&expr.kind, calls);
            }
        }
        ast::StmtKind::Block(block) => {
            for s in block.stmts.iter() {
                extract_calls_from_stmt(&s.kind, calls);
            }
        }
        ast::StmtKind::UncheckedBlock(block) => {
            for s in block.stmts.iter() {
                extract_calls_from_stmt(&s.kind, calls);
            }
        }
        ast::StmtKind::If(cond, then_branch, else_branch) => {
            extract_calls_from_expr(&cond.kind, calls);
            extract_calls_from_stmt(&then_branch.kind, calls);
            if let Some(else_stmt) = else_branch {
                extract_calls_from_stmt(&else_stmt.kind, calls);
            }
        }
        ast::StmtKind::For {
            init,
            cond,
            next,
            body,
        } => {
            if let Some(init_stmt) = init {
                extract_calls_from_stmt(&init_stmt.kind, calls);
            }
            if let Some(cond_expr) = cond {
                extract_calls_from_expr(&cond_expr.kind, calls);
            }
            if let Some(next_expr) = next {
                extract_calls_from_expr(&next_expr.kind, calls);
            }
            extract_calls_from_stmt(&body.kind, calls);
        }
        ast::StmtKind::While(cond, body) => {
            extract_calls_from_expr(&cond.kind, calls);
            extract_calls_from_stmt(&body.kind, calls);
        }
        ast::StmtKind::DoWhile(body, cond) => {
            extract_calls_from_stmt(&body.kind, calls);
            extract_calls_from_expr(&cond.kind, calls);
        }
        ast::StmtKind::DeclSingle(var) => {
            if let Some(init) = &var.initializer {
                extract_calls_from_expr(&init.kind, calls);
            }
        }
        ast::StmtKind::DeclMulti(_, expr) => {
            extract_calls_from_expr(&expr.kind, calls);
        }
        ast::StmtKind::Try(try_stmt) => {
            extract_calls_from_expr(&try_stmt.expr.kind, calls);
            for clause in try_stmt.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    extract_calls_from_stmt(&s.kind, calls);
                }
            }
        }
        // The event or the error is not a call to anything this graph draws, but its
        // ARGUMENTS are ordinary expressions and may hold the function's only call:
        // `emit Funded(currentEpoch(), …)` was a dependency that never appeared.
        ast::StmtKind::Emit(_, args) | ast::StmtKind::Revert(_, args) => {
            for arg in args.exprs() {
                extract_calls_from_expr(&arg.kind, calls);
            }
        }
        _ => {}
    }
}

/// The contract created by `new X(...)` (or `new X{salt: s}(...)`), and the span of `X`.
///
/// Creating a contract runs its constructor, so for the call graph `new X(...)` is a call
/// to `X.constructor` — rendered that way so the same resolver that follows `Lib.fn(...)`
/// reaches it. `new bytes(n)` and `new T[](n)` allocate memory and are not calls.
pub(crate) fn created_contract(
    callee: &ast::ExprKind<'_>,
) -> Option<(String, solar_parse::interface::Span)> {
    match callee {
        ast::ExprKind::New(ty) => match &ty.kind {
            ast::TypeKind::Custom(path) => {
                let ident = path.last();
                Some((ident.as_str().to_string(), ident.span))
            }
            _ => None,
        },
        ast::ExprKind::CallOptions(inner, _) => created_contract(&inner.kind),
        _ => None,
    }
}

/// AST walk: extract call names from an expression kind.
fn extract_calls_from_expr(kind: &ast::ExprKind<'_>, calls: &mut Vec<String>) {
    match kind {
        ast::ExprKind::Call(callee, args) => {
            // `new X(...)` runs X's constructor. Checked first: its callee is not a name.
            if let Some((created, _)) = created_contract(&callee.kind) {
                calls.push(format!("{created}.constructor"));
                for arg in args.exprs() {
                    extract_calls_from_expr(&arg.kind, calls);
                }
                return;
            }
            // `recv.method{value: x, gas: y}(args)` wraps the callee in its option block, so
            // the callee underneath is what says who is being called. Without this the whole
            // call vanished — no box, no boundary, no unresolved entry — and a transfer of
            // native value is exactly the call an auditor cannot afford to lose.
            let callee = match &callee.kind {
                ast::ExprKind::CallOptions(inner, options) => {
                    for option in options.iter() {
                        extract_calls_from_expr(&option.value.kind, calls);
                    }
                    inner
                }
                _ => callee,
            };
            match &callee.kind {
                ast::ExprKind::Ident(ident) => {
                    let name = ident.as_str().to_string();
                    if !is_builtin(&name) {
                        calls.push(name);
                    }
                }
                ast::ExprKind::Member(obj_expr, method_ident) => {
                    let method_name = method_ident.as_str().to_string();
                    match &obj_expr.kind {
                        // obj.method() — extract as "obj.method"
                        ast::ExprKind::Ident(obj_ident) => {
                            let obj_name = obj_ident.as_str().to_string();
                            if !is_builtin_receiver(&obj_name) {
                                calls.push(format!("{}.{}", obj_name, method_name));
                            }
                        }
                        // IFace(addr).method() — the cast renders as `IFace()`, the same
                        // spelling the deploy extractor and analyze_body use. The address
                        // expression is a call site in its own right: in
                        // `IBeacon(_getBeacon()).implementation()`, `_getBeacon()` was lost.
                        ast::ExprKind::Call(inner_callee, inner_args) => {
                            if let ast::ExprKind::Ident(id) = &inner_callee.kind {
                                calls.push(format!("{}().{}", id.as_str(), method_name));
                            } else {
                                // A cast to a BUILT-IN type is not an `Ident`, so the
                                // receiver carries no name to qualify with — and dropping
                                // the whole call with it made a `using X for Y` helper
                                // invisible: `address(token).functionCall(data, …)`, which
                                // is how OpenZeppelin v4's SafeERC20 moves tokens, left no
                                // callee at all and the transfer was marked nowhere. The
                                // method alone is enough for the resolver, which looks a
                                // bare name up in the contracts that define it.
                                calls.push(method_name.clone());
                                extract_calls_from_expr(&inner_callee.kind, calls);
                            }
                            for arg in inner_args.exprs() {
                                extract_calls_from_expr(&arg.kind, calls);
                            }
                        }
                        // A longer chain: whatever it contains still gets walked.
                        other => extract_calls_from_expr(other, calls),
                    }
                }
                _ => {
                    // Complex callee (e.g. chained calls) — recurse into it
                    extract_calls_from_expr(&callee.kind, calls);
                }
            }
            // Also walk arguments for nested calls like foo(bar(x))
            for arg in args.exprs() {
                extract_calls_from_expr(&arg.kind, calls);
            }
        }
        ast::ExprKind::Binary(left, _op, right) => {
            extract_calls_from_expr(&left.kind, calls);
            extract_calls_from_expr(&right.kind, calls);
        }
        ast::ExprKind::Unary(_op, expr) => {
            extract_calls_from_expr(&expr.kind, calls);
        }
        ast::ExprKind::Ternary(cond, if_true, if_false) => {
            extract_calls_from_expr(&cond.kind, calls);
            extract_calls_from_expr(&if_true.kind, calls);
            extract_calls_from_expr(&if_false.kind, calls);
        }
        ast::ExprKind::Assign(left, _op, right) => {
            extract_calls_from_expr(&left.kind, calls);
            extract_calls_from_expr(&right.kind, calls);
        }
        // Both halves: `budget[currentEpoch()][token] += x` calls `currentEpoch`, and
        // walking only the thing being indexed lost it.
        ast::ExprKind::Index(expr, index) => {
            extract_calls_from_expr(&expr.kind, calls);
            match index {
                ast::IndexKind::Index(Some(at)) => extract_calls_from_expr(&at.kind, calls),
                ast::IndexKind::Range(from, to) => {
                    for part in [from, to].into_iter().flatten() {
                        extract_calls_from_expr(&part.kind, calls);
                    }
                }
                _ => {}
            }
        }
        ast::ExprKind::Tuple(elems) => {
            for elem in elems.iter() {
                if let solar_parse::interface::SpannedOption::Some(e) = elem {
                    extract_calls_from_expr(&e.kind, calls);
                }
            }
        }
        ast::ExprKind::Member(expr, _ident) => {
            extract_calls_from_expr(&expr.kind, calls);
        }
        _ => {}
    }
}

/// Regex fallback for extracting function calls when AST parsing fails.
fn extract_calls_regex(source: &str) -> Vec<String> {
    let mut calls = Vec::new();

    let identifier_pattern = regex::Regex::new(r"(\w+)\s*\(").unwrap();
    for cap in identifier_pattern.captures_iter(source) {
        let name = cap[1].to_string();
        if !is_builtin(&name) {
            calls.push(name);
        }
    }

    let external_pattern = regex::Regex::new(r"(\w+)\.(\w+)\s*\(").unwrap();
    for cap in external_pattern.captures_iter(source) {
        let target = &cap[1];
        let method = &cap[2];
        if !is_builtin(target) && !is_builtin(method) {
            calls.push(format!("{}.{}", target, method));
        }
    }

    calls.sort();
    calls.dedup();
    calls
}

/// A receiver that makes `receiver.method()` a builtin (`abi.encode`, `msg.sender`). `this` and
/// `super` are keywords too, but `this.x()` / `super.x()` call the contract's own code, and the
/// resolver follows them, so they are kept.
fn is_builtin_receiver(name: &str) -> bool {
    is_builtin(name) && name != "this" && name != "super"
}

fn is_builtin(name: &str) -> bool {
    matches!(
        name,
        "require"
            | "assert"
            | "revert"
            | "emit"
            | "keccak256"
            | "sha256"
            | "ripemd160"
            | "ecrecover"
            | "addmod"
            | "mulmod"
            | "selfdestruct"
            | "type"
            | "abi"
            | "block"
            | "msg"
            | "tx"
            | "gasleft"
            | "blockhash"
            | "address"
            | "uint256"
            | "uint"
            | "int"
            | "bool"
            | "bytes"
            | "string"
            | "if"
            | "else"
            | "for"
            | "while"
            | "do"
            | "return"
            | "delete"
            | "new"
            | "this"
            | "super"
            | "push"
            | "pop"
            | "length"
    )
}

/// A single call site: the callee name plus where it sits in the source slice.
///
/// The column matters as much as the line. Two calls on one line — say
/// `MathLib.wadMul(amount, price(asset))` — need two connectors landing on two
/// different tokens, and the AST already knows exactly where each one starts.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CallSite {
    pub name: String,
    /// 1-based line, relative to the parsed slice.
    pub line: usize,
    /// 0-based column where the callee's own name begins. For `a.b()` this is
    /// the column of `b`, not of `a`.
    pub column: usize,
    /// The identifier the connector should point at: `b` for `a.b()`.
    pub symbol: String,
    /// Number of arguments passed at this call site — used to pick the right
    /// overload when a name has several. `usize::MAX` means "unknown" (the regex
    /// fallback can't count reliably), which disables overload disambiguation.
    pub arg_count: usize,
}

/// Extract every call site — name **and** line — from a Solidity source slice.
///
/// Unlike [`extract_calls_from_source`], which dedupes into a set of names, this
/// keeps one entry per occurrence, because the Miro deployment draws one
/// connector per call site anchored to its exact line.
///
/// The wrapper puts `source` on its own lines, so a reported column is the
/// column in `source` and a reported line is one more than the line in
/// `source`. Sharing a line with the wrapper would shift every column on the
/// first line by the length of the prefix.
pub fn extract_call_sites_from_source(source: &str) -> Vec<CallSite> {
    if source.trim().is_empty() {
        return Vec::new();
    }
    let wrapped = format!("contract _C {{ function _f() {{\n{}\n}} }}", source);

    let sess = Session::builder().with_silent_emitter(None).build();

    let result = sess.enter(|| -> Option<Vec<CallSite>> {
        let arena = ast::Arena::new();
        let mut parser = Parser::from_source_code(
            &sess,
            &arena,
            FileName::Custom("call_sites".into()),
            wrapped.clone(),
        )
        .ok()?;

        let file = parser.parse_file().map_err(|e| e.emit()).ok()?;

        let mut raw: Vec<RawCallSite> = Vec::new();
        for item in file.items.iter() {
            if let ast::ItemKind::Contract(c) = &item.kind {
                for body_item in c.body.iter() {
                    if let ast::ItemKind::Function(f) = &body_item.kind {
                        if let Some(block) = &f.body {
                            for stmt in block.stmts.iter() {
                                collect_call_sites_from_stmt(&stmt.kind, &mut raw);
                            }
                        }
                    }
                }
            }
        }

        let source_map = sess.source_map();
        let mut sites: Vec<CallSite> = raw
            .into_iter()
            .map(|(name, symbol, span, arg_count)| {
                let location = source_map.lookup_char_pos(span.lo());
                CallSite {
                    name,
                    // The wrapper occupies the first line.
                    line: location.line.saturating_sub(1).max(1),
                    column: location.col.0,
                    symbol,
                    arg_count,
                }
            })
            .collect();
        sites.sort_by(|a, b| {
            a.line
                .cmp(&b.line)
                .then_with(|| a.column.cmp(&b.column))
                .then_with(|| a.name.cmp(&b.name))
        });
        sites.dedup();
        Some(sites)
    });

    result.unwrap_or_else(|| extract_call_sites_regex(source))
}

/// Line-aware regex fallback, used when the AST parse fails.
fn extract_call_sites_regex(source: &str) -> Vec<CallSite> {
    let identifier_pattern = regex::Regex::new(r"(\w+(?:\.\w+)?)\s*\(").unwrap();
    let mut sites = Vec::new();
    for (idx, line) in source.lines().enumerate() {
        // Drop line comments so a call mentioned in prose is not picked up.
        let code = line.split("//").next().unwrap_or("");
        for cap in identifier_pattern.captures_iter(code) {
            let name = cap[1].to_string();
            let head = name.split('.').next().unwrap_or(&name);
            let builtin_head = if name.contains('.') { is_builtin_receiver(head) } else { is_builtin(head) };
            if is_builtin(&name) || builtin_head {
                continue;
            }
            let symbol = name.rsplit('.').next().unwrap_or(&name).to_string();
            let column = cap
                .get(1)
                .map(|m| {
                    // Point at the last segment, matching the AST behaviour.
                    m.start() + name.len() - symbol.len()
                })
                .unwrap_or(0);
            sites.push(CallSite {
                name,
                line: idx + 1,
                column,
                symbol,
                arg_count: usize::MAX, // regex can't count args reliably
            });
        }
    }
    sites
}

type RawCallSite = (String, String, solar_parse::interface::Span, usize);

fn collect_call_sites_from_stmt(kind: &ast::StmtKind<'_>, out: &mut Vec<RawCallSite>) {
    match kind {
        ast::StmtKind::Expr(expr) => collect_call_sites_from_expr(expr, out),
        ast::StmtKind::Return(opt_expr) => {
            if let Some(expr) = opt_expr {
                collect_call_sites_from_expr(expr, out);
            }
        }
        ast::StmtKind::Block(block) | ast::StmtKind::UncheckedBlock(block) => {
            for s in block.stmts.iter() {
                collect_call_sites_from_stmt(&s.kind, out);
            }
        }
        ast::StmtKind::If(cond, then_branch, else_branch) => {
            collect_call_sites_from_expr(cond, out);
            collect_call_sites_from_stmt(&then_branch.kind, out);
            if let Some(else_stmt) = else_branch {
                collect_call_sites_from_stmt(&else_stmt.kind, out);
            }
        }
        ast::StmtKind::For {
            init,
            cond,
            next,
            body,
        } => {
            if let Some(init_stmt) = init {
                collect_call_sites_from_stmt(&init_stmt.kind, out);
            }
            if let Some(cond_expr) = cond {
                collect_call_sites_from_expr(cond_expr, out);
            }
            if let Some(next_expr) = next {
                collect_call_sites_from_expr(next_expr, out);
            }
            collect_call_sites_from_stmt(&body.kind, out);
        }
        ast::StmtKind::While(cond, body) => {
            collect_call_sites_from_expr(cond, out);
            collect_call_sites_from_stmt(&body.kind, out);
        }
        ast::StmtKind::DoWhile(body, cond) => {
            collect_call_sites_from_stmt(&body.kind, out);
            collect_call_sites_from_expr(cond, out);
        }
        ast::StmtKind::DeclSingle(var) => {
            if let Some(init) = &var.initializer {
                collect_call_sites_from_expr(init, out);
            }
        }
        ast::StmtKind::DeclMulti(_, expr) => collect_call_sites_from_expr(expr, out),
        // `emit Funded(currentEpoch(), …)` and `revert TooLate(deadline())`. The event or
        // error itself is not a call to anything drawable, but its ARGUMENTS are ordinary
        // expressions, and a function whose only call sits in one was a dependency that
        // simply never appeared on the diagram.
        ast::StmtKind::Emit(_, args) | ast::StmtKind::Revert(_, args) => {
            for arg in args.exprs() {
                collect_call_sites_from_expr(arg, out);
            }
        }
        ast::StmtKind::Try(try_stmt) => {
            collect_call_sites_from_expr(&try_stmt.expr, out);
            for clause in try_stmt.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_call_sites_from_stmt(&s.kind, out);
                }
            }
        }
        _ => {}
    }
}

fn collect_call_sites_from_expr(expr: &ast::Expr<'_>, out: &mut Vec<RawCallSite>) {
    match &expr.kind {
        ast::ExprKind::Call(callee, args) => {
            let arg_count = args.exprs().count();
            // `new X(...)` runs X's constructor; anchor the arrow on `X`.
            if let Some((created, span)) = created_contract(&callee.kind) {
                out.push((format!("{created}.constructor"), created, span, arg_count));
                for arg in args.exprs() {
                    collect_call_sites_from_expr(arg, out);
                }
                return;
            }
            // `recv.method{value: x, gas: y}(args)` wraps the callee in its option block, so
            // the callee underneath is what says who is being called. Without this the whole
            // call vanished — no box, no boundary, no unresolved entry — and a call that
            // moves native value is the last one an auditor can afford to lose.
            let callee = match &callee.kind {
                ast::ExprKind::CallOptions(inner, options) => {
                    for option in options.iter() {
                        collect_call_sites_from_expr(&option.value, out);
                    }
                    inner
                }
                _ => callee,
            };
            match &callee.kind {
                ast::ExprKind::Ident(ident) => {
                    let name = ident.as_str().to_string();
                    if !is_builtin(&name) {
                        // The identifier's own span, so the connector lands on
                        // the called name rather than on the whole expression.
                        out.push((name.clone(), name, ident.span, arg_count));
                    }
                }
                ast::ExprKind::Member(obj_expr, method_ident) => {
                    let method_name = method_ident.as_str().to_string();
                    // Render the receiver so the call is `receiver.method`. A plain
                    // variable/type is its name; an interface CAST
                    // `IFace(addr).method()` renders as `IFace()` (matching
                    // analyze_body), so the resolver can reach it — otherwise the
                    // call was dropped and neither drawn nor linked.
                    let receiver = match &obj_expr.kind {
                        ast::ExprKind::Ident(obj_ident) => Some(obj_ident.as_str().to_string()),
                        ast::ExprKind::Call(inner_callee, inner_args) => {
                            // The cast's argument is its own call site: in
                            // `IBeacon(_getBeacon()).implementation()` it is `_getBeacon()`,
                            // and skipping it dropped the function's only internal call.
                            for arg in inner_args.exprs() {
                                collect_call_sites_from_expr(arg, out);
                            }
                            match &inner_callee.kind {
                                ast::ExprKind::Ident(id) => Some(format!("{}()", id.as_str())),
                                _ => {
                                    // A cast to a BUILT-IN type is not an `Ident`, so there
                                    // is no receiver to qualify with — and dropping the
                                    // call with it lost every `using X for Y` helper called
                                    // on one. `address(token).functionCall(data, …)` is how
                                    // OpenZeppelin v4 moves tokens, and it was not even a
                                    // call site. The bare method is enough: the resolver
                                    // looks it up in the libraries the contract binds.
                                    collect_call_sites_from_expr(inner_callee, out);
                                    out.push((
                                        method_name.clone(),
                                        method_name.clone(),
                                        method_ident.span,
                                        arg_count,
                                    ));
                                    None
                                }
                            }
                        }
                        _ => None,
                    };
                    match receiver {
                        Some(obj_name) if !is_builtin_receiver(&obj_name) => {
                            // Point at the method, not at the receiver: in
                            // `MathLib.wadMul(...)` the interesting token is
                            // `wadMul`.
                            out.push((
                                format!("{}.{}", obj_name, method_name),
                                method_name,
                                method_ident.span,
                                arg_count,
                            ));
                        }
                        // A call receiver had its callee and arguments walked above;
                        // walking it again would record every nested call twice.
                        _ if matches!(obj_expr.kind, ast::ExprKind::Call(..)) => {}
                        _ => collect_call_sites_from_expr(obj_expr, out),
                    }
                }
                _ => collect_call_sites_from_expr(callee, out),
            }
            for arg in args.exprs() {
                collect_call_sites_from_expr(arg, out);
            }
        }
        ast::ExprKind::Binary(left, _op, right) => {
            collect_call_sites_from_expr(left, out);
            collect_call_sites_from_expr(right, out);
        }
        ast::ExprKind::Unary(_op, inner) => collect_call_sites_from_expr(inner, out),
        ast::ExprKind::Ternary(cond, if_true, if_false) => {
            collect_call_sites_from_expr(cond, out);
            collect_call_sites_from_expr(if_true, out);
            collect_call_sites_from_expr(if_false, out);
        }
        ast::ExprKind::Assign(left, _op, right) => {
            collect_call_sites_from_expr(left, out);
            collect_call_sites_from_expr(right, out);
        }
        // Both halves: `budget[currentEpoch()] += x` calls `currentEpoch`, and walking only
        // the thing being indexed lost it.
        ast::ExprKind::Index(inner, index) => {
            collect_call_sites_from_expr(inner, out);
            match index {
                ast::IndexKind::Index(Some(at)) => collect_call_sites_from_expr(at, out),
                ast::IndexKind::Range(from, to) => {
                    for part in [from, to].into_iter().flatten() {
                        collect_call_sites_from_expr(part, out);
                    }
                }
                _ => {}
            }
        }
        ast::ExprKind::Tuple(elems) => {
            for elem in elems.iter() {
                if let solar_parse::interface::SpannedOption::Some(e) = elem {
                    collect_call_sites_from_expr(e, out);
                }
            }
        }
        ast::ExprKind::Member(inner, _ident) => collect_call_sites_from_expr(inner, out),
        // The option block of a call that carries one, when it is met outside a `Call`.
        ast::ExprKind::CallOptions(inner, options) => {
            collect_call_sites_from_expr(inner, out);
            for option in options.iter() {
                collect_call_sites_from_expr(&option.value, out);
            }
        }
        _ => {}
    }
}

/// Extract the storage locations WRITTEN by a function body, from its source text.
///
/// Fully generic — no library/framework assumptions. A "storage write" is any
/// assignment (`x = …`, `x += …`), `delete x`, `++`/`--`, or `.push()`/`.pop()`
/// whose lvalue's ROOT is either a declared state variable (`state_vars`) OR a
/// local declared with `storage` data location — a storage pointer, e.g. the
/// diamond-storage `MyStruct storage $ = _get();` then `$.field = …`. Resolving
/// the root through `Member`/`Index` is what makes `$.reserveStable += r0` and
/// `balances[user] = v` count, on any codebase.
///
/// Returns readable lvalue paths (e.g. `"reserveStable"`, `"$.reserveStable"`,
/// `"balances[]"`), sorted and deduped. Empty when the body doesn't parse.
/// Everything the metadata needs from a function body, from ONE parse + walk.
#[derive(Default)]
pub struct BodyAnalysis {
    /// Storage locations written (state vars / storage-pointer paths).
    pub storage_writes: Vec<String>,
    /// Every storage write with the 1-based line (relative to `source`) it sits
    /// on — so a diagram can point at the exact statement, not just the function.
    pub storage_write_sites: Vec<(String, usize)>,
    /// External call targets as `(receiver, method)`.
    pub call_targets: Vec<(String, String)>,
    /// The ARGUMENTS of those calls, as written: `(receiver, method, args)`.
    ///
    /// A map that says a contract calls `safeTransfer` says half of it; the half an audit wants
    /// is WHO receives, and that is the first argument. Rendered as source text — `burnSink`,
    /// `address(treasury)`, `msg.sender` — because a name the reader recognises is worth more
    /// than a type we could not resolve: an address-typed immutable has no type to resolve to.
    pub call_arguments: Vec<(String, String, Vec<String>)>,
    /// Local (and parameter) variable declared types, `(name, type)`.
    pub local_types: Vec<(String, String)>,
    /// Callee names for the dependency graph (`foo`, `obj.method`), same format as
    /// `extract_calls_from_source`.
    pub call_names: Vec<String>,
    /// Methods this body calls on `msg.sender`: the control it hands back to its caller.
    pub callbacks: Vec<String>,
    /// Calls that carry native value, `(receiver, method, cast type)`:
    /// `distributor.fund{value: v}(…)`, `IKinetiqStaking(stakingManager).stake{value: v}()`.
    pub value_calls: Vec<(String, String, String)>,
    /// `new X(…)` assigned to something: `(the variable it is stored in, the contract)`.
    pub creations: Vec<(String, String)>,
}

/// Parse a function body ONCE and extract storage writes, external call targets,
/// and local/parameter types together. This is what sonar calls per function —
/// replacing three separate parses. `body_source` is normally the whole function
/// (so the first wrapping parses and we short-circuit); the bare-statements wrapping
/// is a fallback for snippets.
pub fn analyze_body(source: &str, state_vars: &[String], storage_params: &[String]) -> BodyAnalysis {
    for wrapped in [
        format!("contract _C {{ {} }}", source),
        format!("contract _C {{ function _f() {{ {} }} }}", source),
    ] {
        if let Some(a) = analyze_in_wrap(&wrapped, state_vars, storage_params) {
            return a;
        }
    }
    BodyAnalysis::default()
}

fn analyze_in_wrap(
    wrapped: &str,
    state_vars: &[String],
    storage_params: &[String],
) -> Option<BodyAnalysis> {
    let sess = Session::builder().with_silent_emitter(None).build();
    sess.enter(|| -> Option<BodyAnalysis> {
        let arena = ast::Arena::new();
        let mut parser = Parser::from_source_code(
            &sess,
            &arena,
            FileName::Custom("analyze".into()),
            wrapped.to_string(),
        )
        .ok()?;
        let file = parser.parse_file().map_err(|e| e.emit()).ok()?;

        let mut a = BodyAnalysis::default();
        let mut found_fn = false;
        for item in file.items.iter() {
            if let ast::ItemKind::Contract(c) = &item.kind {
                for body_item in c.body.iter() {
                    if let ast::ItemKind::Function(f) = &body_item.kind {
                        if let Some(block) = &f.body {
                            found_fn = true;
                            // Pass 1: storage pointers (params + locals) and the types
                            // of every parameter and local.
                            let mut storage_locals: Vec<String> = storage_params.to_vec();
                            for p in f.header.parameters.iter() {
                                if p.data_location == Some(ast::DataLocation::Storage) {
                                    if let Some(name) = p.name {
                                        storage_locals.push(name.as_str().to_string());
                                    }
                                }
                                if let Some(name) = p.name {
                                    a.local_types.push((
                                        name.as_str().to_string(),
                                        crate::batbelt::evm::parser::evm_file_parser::type_to_string(
                                            &sess, &p.ty,
                                        ),
                                    ));
                                }
                            }
                            for stmt in block.stmts.iter() {
                                collect_storage_locals(&stmt.kind, &mut storage_locals);
                                collect_local_types_from_stmt(&sess, &stmt.kind, &mut a.local_types);
                            }
                            // Pass 2: writes, external call targets, and callee names.
                            let mut writes: Vec<(String, solar_parse::interface::Span)> =
                                Vec::new();
                            let mut targets = Targets::default();
                            for stmt in block.stmts.iter() {
                                collect_writes_from_stmt(
                                    &stmt.kind,
                                    state_vars,
                                    &storage_locals,
                                    &mut writes,
                                );
                                collect_targets_from_stmt(&stmt.kind, &mut targets);
                                extract_calls_from_stmt(&stmt.kind, &mut a.call_names);
                            }
                            a.call_targets.extend(std::mem::take(&mut targets.pairs));
                            a.call_arguments.extend(std::mem::take(&mut targets.arguments));
                            a.callbacks.extend(std::mem::take(&mut targets.callbacks));
                            a.value_calls.extend(std::mem::take(&mut targets.value_calls));
                            a.creations.extend(std::mem::take(&mut targets.creations));
                            let source_map = sess.source_map();
                            for (path, span) in writes {
                                // analyze_body's wrapper prepends `contract _C { … `
                                // with no newline, so `source` line 1 is wrapped
                                // line 1: the reported line IS the source line.
                                let line = source_map.lookup_char_pos(span.lo()).line.max(1);
                                a.storage_writes.push(path.clone());
                                a.storage_write_sites.push((path, line));
                            }
                        }
                    }
                }
            }
        }
        if !found_fn {
            return None;
        }
        a.storage_writes.sort();
        a.storage_writes.dedup();
        a.storage_write_sites.sort();
        a.storage_write_sites.dedup();
        a.call_targets.sort();
        a.call_targets.dedup();
        a.call_arguments.sort();
        a.call_arguments.dedup();
        a.local_types.sort();
        a.local_types.dedup();
        a.call_names.sort();
        a.call_names.dedup();
        Some(a)
    })
}

/// Storage locations written by a function body. Thin wrapper over `analyze_body`
/// (kept for tests); sonar uses `analyze_body` directly to parse once.
pub fn extract_storage_writes_from_source(
    source: &str,
    state_vars: &[String],
    storage_params: &[String],
) -> Vec<String> {
    analyze_body(source, state_vars, storage_params).storage_writes
}

/// Collect the names of locals declared with `storage` data location (storage
/// pointers), recursing through nested blocks/control-flow.
fn collect_storage_locals(kind: &ast::StmtKind<'_>, out: &mut Vec<String>) {
    match kind {
        ast::StmtKind::DeclSingle(var) => {
            if var.data_location == Some(ast::DataLocation::Storage) {
                if let Some(name) = var.name {
                    out.push(name.as_str().to_string());
                }
            }
        }
        ast::StmtKind::Block(b) | ast::StmtKind::UncheckedBlock(b) => {
            for s in b.stmts.iter() {
                collect_storage_locals(&s.kind, out);
            }
        }
        ast::StmtKind::If(_, t, e) => {
            collect_storage_locals(&t.kind, out);
            if let Some(e) = e {
                collect_storage_locals(&e.kind, out);
            }
        }
        ast::StmtKind::For {
            init, body, ..
        } => {
            if let Some(i) = init {
                collect_storage_locals(&i.kind, out);
            }
            collect_storage_locals(&body.kind, out);
        }
        ast::StmtKind::While(_, body) | ast::StmtKind::DoWhile(body, _) => {
            collect_storage_locals(&body.kind, out);
        }
        ast::StmtKind::Try(t) => {
            for clause in t.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_storage_locals(&s.kind, out);
                }
            }
        }
        _ => {}
    }
}

/// Declared types of a function body's local variables + parameters, `(name, type)`.
/// Thin wrapper over `analyze_body` (kept for tests).
pub fn extract_local_types(source: &str) -> Vec<(String, String)> {
    analyze_body(source, &[], &[]).local_types
}

/// The variables a `for` declares in its own initialiser — the loop counters.
///
/// They are locals like any other to the compiler, but not to a reader: a counter's whole
/// life is the three tokens of the loop header, so following it through the function is
/// the one thing nobody needs help with. Reported separately so a caller can leave them
/// out without having to guess from the name — `i` is a counter here and a perfectly good
/// variable elsewhere.
pub fn extract_loop_variables(source: &str) -> Vec<String> {
    // The same two wrappings `analyze_body` uses: a slice is sometimes a whole function and
    // sometimes bare statements, and nesting a function inside one does not parse.
    for wrapped in [
        format!("contract _C {{ {source} }}"),
        format!("contract _C {{ function _f() {{ {source} }} }}"),
    ] {
        let found = loop_variables_in_wrap(&wrapped);
        if !found.is_empty() {
            return found;
        }
    }
    Vec::new()
}

fn loop_variables_in_wrap(wrapped: &str) -> Vec<String> {
    let mut out = Vec::new();
    let sess = Session::builder().with_silent_emitter(None).build();
    let _ = sess.enter(|| -> Option<()> {
        let arena = ast::Arena::new();
        let mut parser = Parser::from_source_code(
            &sess,
            &arena,
            FileName::Custom("loop_vars".into()),
            wrapped.to_string(),
        )
        .ok()?;
        let file = parser.parse_file().map_err(|e| e.emit()).ok()?;
        for item in file.items.iter() {
            if let ast::ItemKind::Contract(c) = &item.kind {
                for body_item in c.body.iter() {
                    if let ast::ItemKind::Function(f) = &body_item.kind {
                        if let Some(block) = &f.body {
                            for stmt in block.stmts.iter() {
                                collect_loop_variables(&stmt.kind, &mut out);
                            }
                        }
                    }
                }
            }
        }
        Some(())
    });
    out
}

fn collect_loop_variables(kind: &ast::StmtKind<'_>, out: &mut Vec<String>) {
    match kind {
        ast::StmtKind::For { init, body, .. } => {
            if let Some(init) = init {
                match &init.kind {
                    ast::StmtKind::DeclSingle(var) => {
                        if let Some(name) = var.name {
                            out.push(name.as_str().to_string());
                        }
                    }
                    ast::StmtKind::DeclMulti(vars, _) => {
                        for slot in vars.iter() {
                            if let Some(var) = slot.as_ref().unspan() {
                                if let Some(name) = var.name {
                                    out.push(name.as_str().to_string());
                                }
                            }
                        }
                    }
                    _ => {}
                }
            }
            collect_loop_variables(&body.kind, out);
        }
        ast::StmtKind::Block(b) | ast::StmtKind::UncheckedBlock(b) => {
            for s in b.stmts.iter() {
                collect_loop_variables(&s.kind, out);
            }
        }
        ast::StmtKind::If(_, t, e) => {
            collect_loop_variables(&t.kind, out);
            if let Some(e) = e {
                collect_loop_variables(&e.kind, out);
            }
        }
        ast::StmtKind::While(_, b) | ast::StmtKind::DoWhile(b, _) => {
            collect_loop_variables(&b.kind, out);
        }
        ast::StmtKind::Try(t) => {
            for clause in t.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_loop_variables(&s.kind, out);
                }
            }
        }
        _ => {}
    }
}

fn collect_local_types_from_stmt(
    sess: &Session,
    kind: &ast::StmtKind<'_>,
    out: &mut Vec<(String, String)>,
) {
    match kind {
        ast::StmtKind::DeclSingle(var) => {
            if let Some(name) = var.name {
                out.push((
                    name.as_str().to_string(),
                    crate::batbelt::evm::parser::evm_file_parser::type_to_string(sess, &var.ty),
                ));
            }
        }
        // A tuple declaration — `(uint256 p0, uint48 ts) = FLAMMStore.price($)` — declares
        // just as much as a single one, and a slot may be skipped (`(, uint256 d) = …`).
        ast::StmtKind::DeclMulti(vars, _) => {
            for slot in vars.iter() {
                let Some(var) = slot.as_ref().unspan() else { continue };
                if let Some(name) = var.name {
                    out.push((
                        name.as_str().to_string(),
                        crate::batbelt::evm::parser::evm_file_parser::type_to_string(sess, &var.ty),
                    ));
                }
            }
        }
        ast::StmtKind::Block(b) | ast::StmtKind::UncheckedBlock(b) => {
            for s in b.stmts.iter() {
                collect_local_types_from_stmt(sess, &s.kind, out);
            }
        }
        ast::StmtKind::If(_, t, e) => {
            collect_local_types_from_stmt(sess, &t.kind, out);
            if let Some(e) = e {
                collect_local_types_from_stmt(sess, &e.kind, out);
            }
        }
        ast::StmtKind::For { init, body, .. } => {
            if let Some(i) = init {
                collect_local_types_from_stmt(sess, &i.kind, out);
            }
            collect_local_types_from_stmt(sess, &body.kind, out);
        }
        ast::StmtKind::While(_, b) | ast::StmtKind::DoWhile(b, _) => {
            collect_local_types_from_stmt(sess, &b.kind, out);
        }
        ast::StmtKind::Try(t) => {
            for clause in t.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_local_types_from_stmt(sess, &s.kind, out);
                }
            }
        }
        _ => {}
    }
}

/// Extract the EXTERNAL call targets of a function body as `(receiver, method)`
/// pairs — a call `x.foo(…)` on some receiver expression. The receiver is rendered
/// as a path (`positionManager`, `$.borrowerOps`, `arr[]`), so a caller can try to
/// resolve it to a concrete contract (a state var's interface type, a struct field's
/// type, …). Captures member/index chains that the deploy-time call-site extractor
/// drops, which is what lets `$.borrowerOps.adjustPosition` be surfaced at all.
/// Builtins (`require`, `.call`, …) are skipped. Dual-wrapped for the same
/// whole-function-vs-bare-statements `body_source` inconsistency as the write scan.
pub fn extract_external_call_targets(source: &str) -> Vec<(String, String)> {
    analyze_body(source, &[], &[]).call_targets
}

fn collect_targets_from_stmt(kind: &ast::StmtKind<'_>, out: &mut Targets) {
    match kind {
        ast::StmtKind::Expr(e) => collect_targets_from_expr(&e.kind, out),
        ast::StmtKind::Return(Some(e)) => collect_targets_from_expr(&e.kind, out),
        ast::StmtKind::Block(b) | ast::StmtKind::UncheckedBlock(b) => {
            for s in b.stmts.iter() {
                collect_targets_from_stmt(&s.kind, out);
            }
        }
        ast::StmtKind::If(cond, t, e) => {
            collect_targets_from_expr(&cond.kind, out);
            collect_targets_from_stmt(&t.kind, out);
            if let Some(e) = e {
                collect_targets_from_stmt(&e.kind, out);
            }
        }
        ast::StmtKind::For {
            init,
            cond,
            next,
            body,
        } => {
            if let Some(i) = init {
                collect_targets_from_stmt(&i.kind, out);
            }
            if let Some(c) = cond {
                collect_targets_from_expr(&c.kind, out);
            }
            if let Some(n) = next {
                collect_targets_from_expr(&n.kind, out);
            }
            collect_targets_from_stmt(&body.kind, out);
        }
        ast::StmtKind::While(c, b) => {
            collect_targets_from_expr(&c.kind, out);
            collect_targets_from_stmt(&b.kind, out);
        }
        ast::StmtKind::DoWhile(b, c) => {
            collect_targets_from_stmt(&b.kind, out);
            collect_targets_from_expr(&c.kind, out);
        }
        ast::StmtKind::DeclSingle(var) => {
            if let Some(init) = &var.initializer {
                collect_targets_from_expr(&init.kind, out);
            }
        }
        ast::StmtKind::DeclMulti(_, e) => collect_targets_from_expr(&e.kind, out),
        // Same reason as the call walkers: an external call can sit in an event's or an
        // error's arguments, and a boundary nobody walks is a boundary nobody marks.
        ast::StmtKind::Emit(_, args) | ast::StmtKind::Revert(_, args) => {
            for arg in args.exprs() {
                collect_targets_from_expr(&arg.kind, out);
            }
        }
        ast::StmtKind::Try(t) => {
            collect_targets_from_expr(&t.expr.kind, out);
            for clause in t.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_targets_from_stmt(&s.kind, out);
                }
            }
        }
        _ => {}
    }
}

/// What a call's argument is called, as the source writes it.
///
/// Casts are unwrapped — `address(treasury)` is `treasury`, `payable(feeSink)` is `feeSink` —
/// because the reader knows the name, not the cast. Anything this does not model is `…`: an
/// honest blank beats a rendering nobody can match to a line of code.
fn arg_text(kind: &ast::ExprKind<'_>) -> String {
    match kind {
        ast::ExprKind::Ident(id) => id.as_str().to_string(),
        ast::ExprKind::Member(obj, field) => format!("{}.{}", arg_text(&obj.kind), field.as_str()),
        ast::ExprKind::Index(obj, _) => format!("{}[]", arg_text(&obj.kind)),
        // A literal is a party too: `address(0)` is the burn address, and a map that drops it
        // says a transfer went nowhere in particular.
        ast::ExprKind::Lit(lit, _) => lit.symbol.as_str().to_string(),
        // `payable(x)` has a node of its own in the grammar, so it never reached the cast
        // branch below: every `payable(feeSink)` was an unreadable blank.
        ast::ExprKind::Payable(args) => cast_text("payable", args),
        // One of two branches runs, and WHICH one is a question about the transaction, not
        // about the code. Both are written down, separated — the map draws a party per branch,
        // so a ternary recipient is two destinations rather than none.
        ast::ExprKind::Ternary(_, yes, no) => {
            format!("{}{TERNARY_MARK}{}", arg_text(&yes.kind), arg_text(&no.kind))
        }
        // `(x)` is `x`. Anything longer is several values and not one party.
        ast::ExprKind::Tuple(elems) => match elems.as_ref() {
            [solar_parse::interface::SpannedOption::Some(single)] => arg_text(&single.kind),
            _ => "…".to_string(),
        },
        ast::ExprKind::Call(callee, args) => match &callee.kind {
            ast::ExprKind::Ident(id) => {
                let name = id.as_str();
                let inner: Vec<&ast::Expr<'_>> = args.exprs().collect();
                // A cast wraps exactly one thing and adds nothing to the reader. Beyond the
                // elementary ones, a single-argument call on a CAPITALISED name is a cast to a
                // user-defined type — `IERC20(token)`, `StakingFeeConverter(addr)` — which is
                // how most arguments that name a contract are written. Solidity capitalises
                // types and nothing else; a wrapper that is genuinely a function of one
                // argument reads the same way to a human too, and its name is kept below.
                let is_cast = matches!(name, "address" | "payable" | "uint160" | "uint256")
                    || name.starts_with(char::is_uppercase);
                if inner.len() == 1 && is_cast {
                    return arg_text(&inner[0].kind);
                }
                // A struct is built the same way a function is called, and it is where an
                // address hides from a signature: `exactInput(ExactInputParams(path, this,
                // …))` declares no address and hands one over. Written out, so whoever reads
                // the declaration can line the fields up with what was passed.
                format!("{name}({})", arg_list(&inner))
            }
            ast::ExprKind::Member(_, method) => {
                format!("{}({})", method.as_str(), arg_list(&args.exprs().collect::<Vec<_>>()))
            }
            // `address(x)`, `uint160(x)`, `IERC20(x)`: the callee is a TYPE, which is its own
            // node — not an identifier — so none of these were ever read. Most arguments that
            // name a contract are written this way, and every one of them came out as `…`.
            ast::ExprKind::Type(ty) => cast_text(&type_text(ty), args),
            _ => "…".to_string(),
        },
        _ => "…".to_string(),
    }
}

/// The contract a `new X(...)` expression deploys, when that is what this expression is.
fn created_type(kind: &ast::ExprKind<'_>) -> Option<String> {
    let new_type = match kind {
        ast::ExprKind::New(ty) => Some(ty),
        // `new X{salt: s}(args)` and the ordinary `new X(args)`: a call whose callee is the
        // `new`, with the option block in between in the salted form.
        ast::ExprKind::Call(callee, _) => match &callee.kind {
            ast::ExprKind::New(ty) => Some(ty),
            ast::ExprKind::CallOptions(inner, _) => match &inner.kind {
                ast::ExprKind::New(ty) => Some(ty),
                _ => None,
            },
            _ => None,
        },
        _ => None,
    }?;
    match type_text(new_type).as_str() {
        "…" => None,
        name => Some(name.to_string()),
    }
}

/// The arguments of a nested call, as the source writes them.
fn arg_list(inner: &[&ast::Expr<'_>]) -> String {
    inner.iter().map(|expr| arg_text(&expr.kind)).collect::<Vec<_>>().join(", ")
}

/// Two parties in one argument, from a ternary: `recipient|burnSink`.
pub const TERNARY_MARK: char = '|';

/// A cast, written as the reader reads it: the name inside, unless that name is a literal, in
/// which case the type is what gives it meaning (`address(0)`, not `0`).
fn cast_text(type_name: &str, args: &ast::CallArgs<'_>) -> String {
    let inner: Vec<&ast::Expr<'_>> = args.exprs().collect();
    let [single] = inner.as_slice() else {
        return "…".to_string();
    };
    let text = arg_text(&single.kind);
    if matches!(single.kind, ast::ExprKind::Lit(..)) {
        return format!("{type_name}({text})");
    }
    text
}

/// A type as its name, for a cast: `address`, `IERC20`, `uint160`.
fn type_text(ty: &ast::Type<'_>) -> String {
    match &ty.kind {
        ast::TypeKind::Elementary(elementary) => elementary.to_string(),
        ast::TypeKind::Custom(path) => path
            .segments()
            .last()
            .map(|segment| segment.as_str().to_string())
            .unwrap_or_else(|| "…".to_string()),
        _ => "…".to_string(),
    }
}

/// Where a walk of the body collects what it finds.
#[derive(Default)]
pub struct Targets {
    pub pairs: Vec<(String, String)>,
    pub arguments: Vec<(String, String, Vec<String>)>,
    /// Methods this body calls on `msg.sender` — the calls that go back to whoever called it.
    pub callbacks: Vec<String>,
    /// Calls that carry native value: `(receiver, method, the type it was cast to)` of every
    /// `x.f{value: v}(…)`. The cast is what names the thing being paid — `IKinetiqStaking(
    /// stakingManager)` — and unwrapping it would leave a bare address.
    pub value_calls: Vec<(String, String, String)>,
    /// `new X(…)` assigned to something: `(the variable it is stored in, the contract)`.
    pub creations: Vec<(String, String)>,
}

fn collect_targets_from_expr(kind: &ast::ExprKind<'_>, out: &mut Targets) {
    match kind {
        ast::ExprKind::Call(callee, args) => {
            // Past the option block of `recv.method{value: x}(args)`, which otherwise hides
            // the receiver and the method from this walk entirely.
            let mut sends_value = false;
            let callee = match &callee.kind {
                ast::ExprKind::CallOptions(inner, options) => {
                    for option in options.iter() {
                        sends_value |= option.name.as_str() == "value";
                        collect_targets_from_expr(&option.value.kind, out);
                    }
                    inner
                }
                _ => callee,
            };
            if let ast::ExprKind::Member(recv, method) = &callee.kind {
                let m = method.as_str().to_string();
                // Native value leaving on an ordinary call: `distributor.fund{value: bought}(…)`
                // moves HYPE, and nothing in `fund(address,uint256)` says so. The receiver is
                // written as the source writes it, casts unwrapped, so a cast of a state
                // address resolves like any other party.
                if sends_value {
                    // The cast, when there is one: `IKinetiqStaking(stakingManager)` pays a
                    // staking manager, and the unwrapped name alone is just an address.
                    let cast = match &recv.kind {
                        ast::ExprKind::Call(callee, _) => match &callee.kind {
                            ast::ExprKind::Ident(id)
                                if id.as_str().starts_with(char::is_uppercase) =>
                            {
                                id.as_str().to_string()
                            }
                            ast::ExprKind::Type(ty) => type_text(ty),
                            _ => String::new(),
                        },
                        _ => String::new(),
                    };
                    let call = (arg_text(&recv.kind), m.clone(), cast);
                    if !out.value_calls.contains(&call) {
                        out.value_calls.push(call);
                    }
                }
                // A call ON `msg.sender` goes back to whoever called this function — the
                // language says so, no name heuristic involved. `Vault.lock` ends in
                // `ILockCallback(msg.sender).lockAcquired(data)`, and that is the arrow that
                // tells an audit control came back here.
                if arg_text(&recv.kind) == "msg.sender" && !out.callbacks.contains(&m) {
                    out.callbacks.push(m.clone());
                }
                // `abi.encodeCall(I.f, (a, b))` names an interface and a method in full, and
                // the bytes it builds are handed straight to a low-level `call` or `staticcall`
                // — in an `assembly` block, where nothing else can see them. That is how the
                // vault reads its rate (`IStakingAccountant.kHYPEToHYPE`) and how OpenZeppelin's
                // `safeTransfer` reaches `token.transfer`: the call IS written down, just not
                // where a call usually is. Recorded as the call it describes.
                if m == "encodeCall" {
                    let mut parts = args.exprs();
                    if let Some(reference) = parts.next() {
                        if let ast::ExprKind::Member(target, called) = &reference.kind {
                            if let Some((_, target_path)) = lvalue_path(&target.kind) {
                                let mut arguments: Vec<String> = Vec::new();
                                if let Some(tuple) = parts.next() {
                                    if let ast::ExprKind::Tuple(elems) = &tuple.kind {
                                        for elem in elems.iter() {
                                            if let solar_parse::interface::SpannedOption::Some(e) =
                                                elem
                                            {
                                                arguments.push(arg_text(&e.kind));
                                            }
                                        }
                                    } else {
                                        arguments.push(arg_text(&tuple.kind));
                                    }
                                }
                                let called = called.as_str().to_string();
                                out.arguments.push((
                                    target_path.clone(),
                                    called.clone(),
                                    arguments,
                                ));
                                out.pairs.push((target_path, called));
                            }
                        }
                    }
                }
                if !is_builtin(&m) {
                    if let Some((_, recv_path)) = lvalue_path(&recv.kind) {
                        out.arguments.push((
                            recv_path.clone(),
                            m.clone(),
                            args.exprs().map(|arg| arg_text(&arg.kind)).collect(),
                        ));
                        out.pairs.push((recv_path, m));
                    }
                }
            }
            collect_targets_from_expr(&callee.kind, out);
            for arg in args.exprs() {
                collect_targets_from_expr(&arg.kind, out);
            }
        }
        ast::ExprKind::Assign(l, _, r) => {
            // WHICH variable holds what was just deployed. `opexTreasury = new Treasury(owner_)`
            // five times over is five addresses, and a map that knows only the type draws one
            // arrow to one box for all of them.
            if let Some(created) = created_type(&r.kind) {
                if let Some((_, held_by)) = lvalue_path(&l.kind) {
                    let creation = (held_by, created);
                    if !out.creations.contains(&creation) {
                        out.creations.push(creation);
                    }
                }
            }
            collect_targets_from_expr(&l.kind, out);
            collect_targets_from_expr(&r.kind, out);
        }
        ast::ExprKind::Binary(l, _, r) => {
            collect_targets_from_expr(&l.kind, out);
            collect_targets_from_expr(&r.kind, out);
        }
        ast::ExprKind::Unary(_, e) | ast::ExprKind::Delete(e) => {
            collect_targets_from_expr(&e.kind, out)
        }
        ast::ExprKind::Ternary(c, a, b) => {
            collect_targets_from_expr(&c.kind, out);
            collect_targets_from_expr(&a.kind, out);
            collect_targets_from_expr(&b.kind, out);
        }
        ast::ExprKind::Index(e, index) => {
            collect_targets_from_expr(&e.kind, out);
            match index {
                ast::IndexKind::Index(Some(at)) => collect_targets_from_expr(&at.kind, out),
                ast::IndexKind::Range(from, to) => {
                    for part in [from, to].into_iter().flatten() {
                        collect_targets_from_expr(&part.kind, out);
                    }
                }
                _ => {}
            }
        }
        ast::ExprKind::Member(e, _) => collect_targets_from_expr(&e.kind, out),
        ast::ExprKind::Tuple(elems) => {
            for el in elems.iter() {
                if let solar_parse::interface::SpannedOption::Some(e) = el {
                    collect_targets_from_expr(&e.kind, out);
                }
            }
        }
        _ => {}
    }
}

/// The leftmost base of an lvalue: a plain identifier, or a call (a storage
/// accessor like `_s()` in the diamond-storage pattern — `_s().field = …`).
enum LvalueRoot {
    Ident(String),
    Call,
}

/// The `(root, rendered_path)` of an lvalue expression, walking through
/// member/index access to the leftmost base. `None` for a non-lvalue base
/// (a literal, an unsupported form, …).
fn lvalue_path(kind: &ast::ExprKind<'_>) -> Option<(LvalueRoot, String)> {
    match kind {
        ast::ExprKind::Ident(id) => {
            let n = id.as_str().to_string();
            Some((LvalueRoot::Ident(n.clone()), n))
        }
        ast::ExprKind::Member(obj, field) => {
            let (root, path) = lvalue_path(&obj.kind)?;
            Some((root, format!("{}.{}", path, field.as_str())))
        }
        ast::ExprKind::Index(obj, _) => {
            let (root, path) = lvalue_path(&obj.kind)?;
            Some((root, format!("{}[]", path)))
        }
        // A call at the base of an lvalue — you can only meaningfully assign
        // THROUGH a call result if it returns a storage reference, so a
        // `getter().field = …` is a storage write. Render the callee for context.
        ast::ExprKind::Call(callee, _) => {
            let name = match &callee.kind {
                ast::ExprKind::Ident(id) => id.as_str().to_string(),
                ast::ExprKind::Member(_, m) => m.as_str().to_string(),
                _ => "call".to_string(),
            };
            Some((LvalueRoot::Call, format!("{name}()")))
        }
        _ => None,
    }
}

/// Record `lval` as a storage write iff its root is a state variable, a
/// storage-pointer local, or a storage accessor call.
fn record_write(
    lval: &ast::ExprKind<'_>,
    span: solar_parse::interface::Span,
    state_vars: &[String],
    storage_locals: &[String],
    out: &mut Vec<(String, solar_parse::interface::Span)>,
) {
    if let Some((root, path)) = lvalue_path(lval) {
        let is_storage = match root {
            LvalueRoot::Call => true,
            LvalueRoot::Ident(name) => {
                state_vars.iter().any(|s| *s == name) || storage_locals.iter().any(|s| *s == name)
            }
        };
        if is_storage {
            out.push((path, span));
        }
    }
}

/// AST walk: record storage writes from a statement.
fn collect_writes_from_stmt(
    kind: &ast::StmtKind<'_>,
    sv: &[String],
    sl: &[String],
    out: &mut Vec<(String, solar_parse::interface::Span)>,
) {
    match kind {
        ast::StmtKind::Expr(expr) => collect_writes_from_expr(&expr.kind, sv, sl, out),
        ast::StmtKind::Block(b) | ast::StmtKind::UncheckedBlock(b) => {
            for s in b.stmts.iter() {
                collect_writes_from_stmt(&s.kind, sv, sl, out);
            }
        }
        ast::StmtKind::If(cond, t, e) => {
            collect_writes_from_expr(&cond.kind, sv, sl, out);
            collect_writes_from_stmt(&t.kind, sv, sl, out);
            if let Some(e) = e {
                collect_writes_from_stmt(&e.kind, sv, sl, out);
            }
        }
        ast::StmtKind::For {
            init,
            cond,
            next,
            body,
        } => {
            if let Some(i) = init {
                collect_writes_from_stmt(&i.kind, sv, sl, out);
            }
            if let Some(c) = cond {
                collect_writes_from_expr(&c.kind, sv, sl, out);
            }
            if let Some(n) = next {
                collect_writes_from_expr(&n.kind, sv, sl, out);
            }
            collect_writes_from_stmt(&body.kind, sv, sl, out);
        }
        ast::StmtKind::While(cond, body) => {
            collect_writes_from_expr(&cond.kind, sv, sl, out);
            collect_writes_from_stmt(&body.kind, sv, sl, out);
        }
        ast::StmtKind::DoWhile(body, cond) => {
            collect_writes_from_stmt(&body.kind, sv, sl, out);
            collect_writes_from_expr(&cond.kind, sv, sl, out);
        }
        ast::StmtKind::DeclSingle(var) => {
            if let Some(init) = &var.initializer {
                collect_writes_from_expr(&init.kind, sv, sl, out);
            }
        }
        ast::StmtKind::DeclMulti(_, expr) => collect_writes_from_expr(&expr.kind, sv, sl, out),
        ast::StmtKind::Try(t) => {
            collect_writes_from_expr(&t.expr.kind, sv, sl, out);
            for clause in t.clauses.iter() {
                for s in clause.block.stmts.iter() {
                    collect_writes_from_stmt(&s.kind, sv, sl, out);
                }
            }
        }
        _ => {}
    }
}

/// AST walk: record storage writes from an expression.
fn collect_writes_from_expr(
    kind: &ast::ExprKind<'_>,
    sv: &[String],
    sl: &[String],
    out: &mut Vec<(String, solar_parse::interface::Span)>,
) {
    match kind {
        ast::ExprKind::Assign(left, _op, right) => {
            record_write(&left.kind, left.span, sv, sl, out);
            collect_writes_from_expr(&left.kind, sv, sl, out);
            collect_writes_from_expr(&right.kind, sv, sl, out);
        }
        ast::ExprKind::Delete(inner) => {
            record_write(&inner.kind, inner.span, sv, sl, out);
            collect_writes_from_expr(&inner.kind, sv, sl, out);
        }
        ast::ExprKind::Unary(op, inner) => {
            if matches!(
                op.kind,
                ast::UnOpKind::PreInc
                    | ast::UnOpKind::PreDec
                    | ast::UnOpKind::PostInc
                    | ast::UnOpKind::PostDec
            ) {
                record_write(&inner.kind, inner.span, sv, sl, out);
            }
            collect_writes_from_expr(&inner.kind, sv, sl, out);
        }
        ast::ExprKind::Call(callee, args) => {
            // `.push(...)` / `.pop()` mutate the storage array they target.
            if let ast::ExprKind::Member(base, method) = &callee.kind {
                let m = method.as_str();
                if m == "push" || m == "pop" {
                    record_write(&base.kind, base.span, sv, sl, out);
                }
            }
            collect_writes_from_expr(&callee.kind, sv, sl, out);
            for arg in args.exprs() {
                collect_writes_from_expr(&arg.kind, sv, sl, out);
            }
        }
        ast::ExprKind::Binary(l, _, r) => {
            collect_writes_from_expr(&l.kind, sv, sl, out);
            collect_writes_from_expr(&r.kind, sv, sl, out);
        }
        ast::ExprKind::Ternary(c, a, b) => {
            collect_writes_from_expr(&c.kind, sv, sl, out);
            collect_writes_from_expr(&a.kind, sv, sl, out);
            collect_writes_from_expr(&b.kind, sv, sl, out);
        }
        ast::ExprKind::Index(e, _) => collect_writes_from_expr(&e.kind, sv, sl, out),
        ast::ExprKind::Member(e, _) => collect_writes_from_expr(&e.kind, sv, sl, out),
        ast::ExprKind::Tuple(elems) => {
            for el in elems.iter() {
                if let solar_parse::interface::SpannedOption::Some(e) = el {
                    collect_writes_from_expr(&e.kind, sv, sl, out);
                }
            }
        }
        _ => {}
    }
}

#[cfg(test)]
mod storage_write_test {


    use super::{
        extract_external_call_targets, extract_local_types, extract_storage_writes_from_source,
    };

    #[test]
    fn local_types_capture_storage_pointer_and_casts() {
        let src = "\
            CollateralRebalancerStorage storage $ = _s();\n\
            IERC20 collToken = IERC20(address($.collVault));\n\
            uint256 x = 1;\n";
        let t = extract_local_types(src);
        assert!(t.contains(&("$".to_string(), "CollateralRebalancerStorage".to_string())), "{t:?}");
        assert!(t.iter().any(|(n, _)| n == "collToken"), "{t:?}");
    }

    #[test]
    fn external_call_targets_capture_member_chains() {
        let src = "\
            $.borrowerOps.adjustPosition(a, b);\n\
            positionManager.setX(1);\n\
            uint256 x = oracle.price();\n\
            require(ok, \"e\");\n\
            token.safeTransfer(to, amt);\n";
        let t = extract_external_call_targets(src);
        assert!(t.contains(&("$.borrowerOps".to_string(), "adjustPosition".to_string())), "{t:?}");
        assert!(t.contains(&("positionManager".to_string(), "setX".to_string())), "{t:?}");
        assert!(t.contains(&("oracle".to_string(), "price".to_string())), "{t:?}");
        assert!(t.contains(&("token".to_string(), "safeTransfer".to_string())), "{t:?}");
        // `require` is a builtin, never a call target.
        assert!(!t.iter().any(|(_, m)| m == "require"), "{t:?}");
    }

    #[test]
    fn detects_direct_pointer_index_and_ignores_locals() {
        let src = r#"
            Layout storage $ = _layout();
            uint256 memTotal = 0;      // local, NOT storage
            $.reserveStable += r0;     // storage pointer write
            totalSupply = memTotal;    // direct state-var write
            balances[user] = 5;        // mapping/index write
            paused = true;             // direct state-var write
            memTotal = 7;              // local write — must be ignored
            counter++;                 // unary state-var write
            delete lastActor;          // delete state-var
            _s().kappa = 3;            // storage accessor call write
        "#;
        let state_vars = vec![
            "totalSupply".to_string(),
            "balances".to_string(),
            "paused".to_string(),
            "counter".to_string(),
            "lastActor".to_string(),
        ];
        let w = extract_storage_writes_from_source(src, &state_vars, &[]);
        assert!(w.contains(&"$.reserveStable".to_string()), "pointer: {w:?}");
        assert!(w.contains(&"totalSupply".to_string()), "direct: {w:?}");
        assert!(w.contains(&"balances[]".to_string()), "index: {w:?}");
        assert!(w.contains(&"paused".to_string()), "bool: {w:?}");
        assert!(w.contains(&"counter".to_string()), "unary: {w:?}");
        assert!(w.contains(&"lastActor".to_string()), "delete: {w:?}");
        assert!(w.contains(&"_s().kappa".to_string()), "accessor call: {w:?}");
        assert!(!w.contains(&"memTotal".to_string()), "local leaked: {w:?}");
    }

    #[test]
    fn detects_write_through_a_storage_parameter() {
        // A library that takes the storage struct as a `storage` PARAMETER and
        // writes through it — the CvammSwapLib.execute shape. Whole-function body.
        let src = r#"
            function execute(CvammStore.CvammStorage storage $, SwapVars memory v)
                internal
                returns (uint256)
            {
                uint256 k = $.kappa;
                $.reserveStable += v.amountInUsed;
                return k;
            }
        "#;
        let w = extract_storage_writes_from_source(src, &[], &[]);
        assert!(w.contains(&"$.reserveStable".to_string()), "got {w:?}");
    }

    #[test]
    fn detects_accessor_write_in_a_real_setter_shape() {
        // The exact shape of a real `setPaused`: guard branches, then an
        // accessor-call storage write. No state_vars needed (call-rooted).
        let src = r#"
            if (p) {
                if (msg.sender != owner() && msg.sender != guardian()) revert NotPauser();
            } else if (msg.sender != owner()) {
                revert NotOwner();
            }
            _s().paused = p;
        "#;
        let w = extract_storage_writes_from_source(src, &[], &[]);
        assert_eq!(w, vec!["_s().paused".to_string()], "got {w:?}");
    }
}

#[cfg(test)]
mod call_site_test {
    use super::*;

    #[test]
    fn test_extract_call_sites_self_call_in_try() {
        // `try this.x()` is an external call to the contract itself; it must be drawn.
        let source = r#"
try this.cross(base, quote) returns (uint256 p, uint48 t) {
    return (true, p, t);
} catch {
    return (false, 0, 0);
}
"#;
        let sites = extract_call_sites_from_source(source);
        assert!(sites.iter().any(|s| s.name == "this.cross" && s.symbol == "cross"), "{sites:?}");
    }

    #[test]
    fn test_extract_call_sites_keeps_line_and_repeats() {
        let source = r#"
uint256 fee = FeeLib.feeOf(assets, entryFeeBps);
uint256 net = assets - fee;
mintedShares = ShareLib.toShares(net, totalAssets(), totalShares);
_mint(receiver, mintedShares);
"#;
        let sites = extract_call_sites_from_source(source);

        let find = |name: &str| -> Vec<usize> {
            sites
                .iter()
                .filter(|s| s.name == name)
                .map(|s| s.line)
                .collect()
        };

        assert_eq!(find("FeeLib.feeOf"), vec![2]);
        assert_eq!(find("ShareLib.toShares"), vec![4]);
        assert_eq!(find("totalAssets"), vec![4]);
        assert_eq!(find("_mint"), vec![5]);
    }

    /// Two calls on one line must be told apart by column, which is what lets
    /// the deployment draw one connector per token instead of stacking both on
    /// the same spot.
    #[test]
    fn test_two_calls_on_the_same_line_have_distinct_columns() {
        let source = "return MathLib.wadMul(amount, price(asset));";
        let sites = extract_call_sites_from_source(source);

        let wad = sites
            .iter()
            .find(|s| s.name == "MathLib.wadMul")
            .expect("wadMul missing");
        let price = sites
            .iter()
            .find(|s| s.name == "price")
            .expect("price missing");

        assert_eq!(wad.line, price.line, "both are on the same line");
        assert_ne!(wad.column, price.column);
        assert!(wad.column < price.column, "wadMul comes first");

        // The column must point at the method, not at the receiver.
        assert_eq!(wad.symbol, "wadMul");
        assert_eq!(&source[wad.column..wad.column + wad.symbol.len()], "wadMul");
        assert_eq!(
            &source[price.column..price.column + price.symbol.len()],
            "price"
        );
    }

    #[test]
    fn test_same_callee_on_several_lines() {
        let source = r#"
uint256 a = mulDiv(x, y, z);
uint256 b = other(1);
uint256 c = mulDiv(p, q, r);
"#;
        let lines: Vec<usize> = extract_call_sites_from_source(source)
            .into_iter()
            .filter(|s| s.name == "mulDiv")
            .map(|s| s.line)
            .collect();
        assert_eq!(lines, vec![2, 4], "one call site per occurrence");
    }
}


/// Blank out everything that is not the function's body, keeping line and
/// column positions intact.
///
/// Call sites are extracted by wrapping the source in
/// `contract _C { function _f() { … } }`. A slice that still carries its own
/// signature puts a function declaration inside a function body, which is not
/// valid Solidity: the parse fails, the regex fallback runs, and it reads the
/// signature `foo(uint256 a)` as a call to `foo` — every function ended up
/// calling itself.
///
/// Blanking rather than trimming means a reported line is still the line in
/// `slice`, and a reported column is still the column in that line.
pub fn body_only(slice: &[String]) -> Vec<String> {
    let Some(open_index) = slice.iter().position(|line| line.contains('{')) else {
        return slice.to_vec();
    };

    let mut body: Vec<String> = slice.to_vec();
    for line in body.iter_mut().take(open_index) {
        *line = " ".repeat(line.len());
    }
    // Keep whatever follows the brace on its own line, blanking the rest.
    if let Some(brace) = body[open_index].find('{') {
        let tail = body[open_index][brace + 1..].to_string();
        body[open_index] = format!("{}{}", " ".repeat(brace + 1), tail);
    }

    // The closing brace would end the wrapper's function early.
    if let Some(close_index) = body.iter().rposition(|line| line.contains('}')) {
        if close_index > open_index {
            if let Some(brace) = body[close_index].rfind('}') {
                let head = body[close_index][..brace].to_string();
                body[close_index] = format!("{}{}", head, " ".repeat(body[close_index].len() - brace));
            }
        }
    }
    body
}

#[cfg(test)]
mod body_only_test {
    use super::*;

    fn lines(text: &str) -> Vec<String> {
        text.lines().map(|line| line.to_string()).collect()
    }

    /// The signature must not be read as a call to the function itself.
    #[test]
    fn test_signature_is_not_a_call_site() {
        let slice = lines(
            "function depositWithReferral(uint256 assets, address receiver)\n             external\n             returns (uint256 mintedShares)\n             {\n             mintedShares = _deposit(msg.sender, receiver, assets);\n             }",
        );
        let sites = extract_call_sites_from_source(&body_only(&slice).join("\n"));
        let names: Vec<&str> = sites.iter().map(|site| site.name.as_str()).collect();

        assert!(!names.contains(&"depositWithReferral"), "got {names:?}");
        assert!(names.contains(&"_deposit"), "got {names:?}");
    }

    /// Blanking keeps every position, so anchors still land on the right token.
    #[test]
    fn test_positions_survive_blanking() {
        let slice = lines(
            "function quote(address asset) external view returns (uint256) {\n             return MathLib.wadMul(amount, price(asset));\n             }",
        );
        let sites = extract_call_sites_from_source(&body_only(&slice).join("\n"));

        let wad = sites.iter().find(|s| s.name == "MathLib.wadMul").unwrap();
        assert_eq!(wad.line, 2, "line numbers must still refer to the slice");
        assert_eq!(&slice[1][wad.column..wad.column + wad.symbol.len()], "wadMul");
    }

    /// A one-line body keeps its call.
    #[test]
    fn test_body_on_the_signature_line() {
        let slice = lines("function f() internal { return g(); }");
        let sites = extract_call_sites_from_source(&body_only(&slice).join("\n"));
        let names: Vec<&str> = sites.iter().map(|site| site.name.as_str()).collect();
        assert!(names.contains(&"g"), "got {names:?}");
        assert!(!names.contains(&"f"), "got {names:?}");
    }
}

#[cfg(test)]
mod ast_path_test {
    use super::*;

    /// Source where the two extractors must disagree.
    ///
    /// The regex fallback scans text, so it reads `ghost()` inside a block
    /// comment and `phantom()` inside a string literal as calls, and it takes
    /// `returns (` for a call to `returns`. The AST sees none of them. Asserting
    /// their absence is therefore an assertion that the AST path ran.
    const TRICKY: &str = r#"{
        /* ghost(); */
        string memory note = "phantom()";
        uint256 value = real(1);
    }"#;

    #[test]
    fn test_dependency_scan_uses_the_ast() {
        let names = extract_calls_from_source(TRICKY);
        assert!(names.contains(&"real".to_string()), "got {names:?}");
        assert!(!names.contains(&"ghost".to_string()), "comment read as a call: {names:?}");
        assert!(!names.contains(&"phantom".to_string()), "string read as a call: {names:?}");
    }

    #[test]
    fn test_call_sites_use_the_ast() {
        let sites = extract_call_sites_from_source(TRICKY);
        let names: Vec<&str> = sites.iter().map(|site| site.name.as_str()).collect();
        assert!(names.contains(&"real"), "got {names:?}");
        assert!(!names.contains(&"ghost"), "comment read as a call: {names:?}");
        assert!(!names.contains(&"phantom"), "string read as a call: {names:?}");
    }

    /// What the fallback would produce, kept as the contrast: if this ever
    /// matches the AST output, the probe above has stopped proving anything.
    #[test]
    fn test_the_fallback_really_is_fooled() {
        let names = extract_calls_regex(TRICKY);
        assert!(
            names.contains(&"ghost".to_string()) || names.contains(&"phantom".to_string()),
            "the regex fallback no longer differs from the AST: {names:?}"
        );
    }
}

#[cfg(test)]
mod cast_call_test {
    use super::{extract_call_sites_from_source, extract_calls_from_source};

    #[test]
    fn a_cast_receiver_keeps_the_call_inside_its_argument() {
        // BeaconProxy._implementation: the internal call lives INSIDE the cast, and both
        // extractors used to drop it — the scan dropped the whole statement.
        let body = "return IBeacon(_getBeacon()).implementation();";

        let names = extract_calls_from_source(body);
        assert!(names.contains(&"_getBeacon".to_string()), "{names:?}");
        assert!(names.contains(&"IBeacon().implementation".to_string()), "{names:?}");

        let sites = extract_call_sites_from_source(body);
        let site_names: Vec<&str> = sites.iter().map(|s| s.name.as_str()).collect();
        assert!(site_names.contains(&"_getBeacon"), "{site_names:?}");
        assert!(site_names.contains(&"IBeacon().implementation"), "{site_names:?}");
        // Walked once, not twice: a duplicate site would draw a duplicate arrow.
        assert_eq!(site_names.iter().filter(|n| **n == "_getBeacon").count(), 1);
    }

    #[test]
    fn a_nested_chain_is_still_walked() {
        let sites = extract_call_sites_from_source("foo(bar()).baz(qux());");
        let names: Vec<&str> = sites.iter().map(|s| s.name.as_str()).collect();
        for expected in ["bar", "qux"] {
            assert_eq!(names.iter().filter(|n| **n == expected).count(), 1, "{names:?}");
        }
    }
}

#[cfg(test)]
mod contract_creation_test {
    use super::{extract_call_sites_from_source, extract_calls_from_source};

    #[test]
    fn new_contract_is_a_call_to_its_constructor() {
        // MorphoAccountDeployer.deploy: the only thing it does is create the account, and the
        // creation was invisible — no callee, no arrow, one lonely screenshot.
        let body = "return address(new MorphoBlueAccount(router, pool, f(x), loanAsset, MORPHO));";

        let names = extract_calls_from_source(body);
        assert!(names.contains(&"MorphoBlueAccount.constructor".to_string()), "{names:?}");
        assert!(names.contains(&"f".to_string()), "arguments are still walked: {names:?}");

        let sites = extract_call_sites_from_source(body);
        let site = sites
            .iter()
            .find(|s| s.name == "MorphoBlueAccount.constructor")
            .expect("creation site");
        assert_eq!(site.symbol, "MorphoBlueAccount");
        assert_eq!(site.arg_count, 5);
    }

    #[test]
    fn salted_creation_and_memory_allocation() {
        let sites = extract_call_sites_from_source("Pool p = new Pool{salt: s}(a); bytes memory b = new bytes(32);");
        let names: Vec<&str> = sites.iter().map(|s| s.name.as_str()).collect();
        assert!(names.contains(&"Pool.constructor"), "{names:?}");
        assert!(!names.iter().any(|n| n.contains("bytes")), "{names:?}");
    }
}

#[cfg(test)]
mod using_for_call_test {
    use super::analyze_body;

    #[test]
    fn a_library_call_on_an_address_is_a_callee() {
        // Verbatim from OpenZeppelin v4's SafeERC20, which is where a token transfer
        // disappeared from a diagram: the call goes through `using Address for address`.
        let body = r#"function _callOptionalReturn(IERC20 token, bytes memory data) private {
        // We need to perform a low level call here, to bypass Solidity's return data size checking mechanism, since
        // we're implementing it ourselves. We use {Address-functionCall} to perform this call, which verifies that
        // the target address contains contract code and also asserts for success in the low-level call.

        bytes memory returndata = address(token).functionCall(data, "SafeERC20: low-level call failed");
        if (returndata.length > 0) {
            // Return data is optional
            require(abi.decode(returndata, (bool)), "SafeERC20: ERC20 operation did not succeed");
        }
    }"#;
        let analysis = analyze_body(body, &[], &[]);
        assert!(
            analysis.call_names.iter().any(|name| name.contains("functionCall")),
            "no callee recorded; got {:?}",
            analysis.call_names
        );
    }
}

#[cfg(test)]
mod nested_call_test {
    use super::{extract_call_sites_from_source, extract_calls_from_source};

    fn names(source: &str) -> Vec<String> {
        extract_call_sites_from_source(source)
            .into_iter()
            .map(|site| site.name)
            .collect()
    }

    /// A call that only ever appears inside an `emit` is still a call.
    #[test]
    fn a_call_in_an_emit_argument_is_a_call() {
        let source = "emit Funded(currentEpoch(), token, amount);";
        assert!(names(source).contains(&"currentEpoch".to_string()));
        assert!(extract_calls_from_source(source).contains(&"currentEpoch".to_string()));
    }

    /// And so is one inside a `revert`'s arguments.
    #[test]
    fn a_call_in_a_revert_argument_is_a_call() {
        let source = "revert TooLate(deadline());";
        assert!(names(source).contains(&"deadline".to_string()));
    }

    /// The INDEX of a mapping is an expression like any other.
    #[test]
    fn a_call_in_a_mapping_index_is_a_call() {
        let source = "budget[currentEpoch()][token] += received;";
        assert!(names(source).contains(&"currentEpoch".to_string()));
        assert!(extract_calls_from_source(source).contains(&"currentEpoch".to_string()));
    }

    /// The statement form still works — the control the report used.
    #[test]
    fn a_call_as_a_declaration_is_still_a_call() {
        let source = "uint256 into = currentEpoch();";
        assert!(names(source).contains(&"currentEpoch".to_string()));
    }
}

#[cfg(test)]
mod call_options_test {
    use super::{extract_call_sites_from_source, extract_calls_from_source};

    fn names(source: &str) -> Vec<String> {
        extract_call_sites_from_source(source)
            .into_iter()
            .map(|site| site.name)
            .collect()
    }

    /// `recv.method{value: x}(args)` is a call to `recv.method`, options or no options. It
    /// moves native value, so losing it loses the one call an auditor cannot afford to miss.
    #[test]
    fn a_call_with_a_value_option_is_still_a_call() {
        let source = "distributor.fund{value: bought}(address(0), bought);";
        assert!(names(source).contains(&"distributor.fund".to_string()));
        assert!(extract_calls_from_source(source).contains(&"distributor.fund".to_string()));
    }

    /// The same through an interface cast, which is how a boundary usually looks.
    #[test]
    fn a_cast_receiver_with_options_is_still_a_call() {
        let source = "IKinetiqStaking(stakingManager).stake{value: bought}();";
        assert!(names(source).iter().any(|name| name.ends_with(".stake")));
    }

    /// `gas:` too, and both at once.
    #[test]
    fn gas_and_value_together_are_still_a_call() {
        let source = "target.run{value: v, gas: g}(data);";
        assert!(names(source).contains(&"target.run".to_string()));
    }

    /// What the options themselves call is a call as well: `{value: price()}` calls `price`.
    #[test]
    fn a_call_inside_an_option_is_a_call() {
        let source = "target.run{value: price()}(data);";
        assert!(names(source).contains(&"price".to_string()));
    }
}

#[cfg(test)]
mod constructor_targets_test {
    use super::analyze_body;

    /// A constructor's body is walked like any other: `AscendLaunchpad`'s declares two locals
    /// from a cast and calls four methods on them, and none of those calls reached the map.
    #[test]
    fn a_constructors_calls_are_collected() {
        let src = r#"
constructor(address[2] memory feeConverters_, address owner_) TwoStepOwned(owner_) {
    StakingFeeConverter stakeConverter = StakingFeeConverter(payable(feeConverters_[0]));
    if (stakeConverter.extraYield()) revert BadQuoteConfig();
    treasury = stakeConverter.treasury();
}
"#;
        let analysis = analyze_body(src, &["treasury".to_string()], &[]);
        assert!(
            analysis.local_types.iter().any(|(n, t)| n == "stakeConverter" && t == "StakingFeeConverter"),
            "the local's type: {:?}",
            analysis.local_types
        );
        assert!(
            analysis.call_targets.contains(&("stakeConverter".to_string(), "extraYield".to_string())),
            "the call in the `if`: {:?}",
            analysis.call_targets
        );
        assert!(
            analysis.call_targets.contains(&("stakeConverter".to_string(), "treasury".to_string())),
            "the call in the assignment: {:?}",
            analysis.call_targets
        );
    }
}

#[cfg(test)]
mod argument_text_test {
    use super::analyze_body;

    fn arguments_of(src: &str, method: &str) -> Vec<String> {
        analyze_body(src, &[], &[])
            .call_arguments
            .into_iter()
            .find(|(_, called, _)| called == method)
            .map(|(_, _, args)| args)
            .unwrap_or_default()
    }

    /// `address(x)` is a cast to a TYPE, which the grammar gives a node of its own — not an
    /// identifier — so every argument written that way (most of the ones that name a contract)
    /// came out as the blank `…` and the map could not even ask about it.
    #[test]
    fn a_cast_to_a_type_is_the_name_inside_it() {
        let src = r#"
function pay(uint256 amount) external {
    launchpad.takeListing(poolId, address(listingTreasury), amount);
    vault.take(currency, address(this), amount);
    token.transfer(IERC20(other), amount);
}
"#;
        assert_eq!(arguments_of(src, "takeListing")[1], "listingTreasury");
        assert_eq!(arguments_of(src, "take")[1], "this");
        assert_eq!(arguments_of(src, "transfer")[0], "other");
    }

    /// A cast of a LITERAL keeps its type, because that is what gives it meaning: `address(0)`
    /// is the burn address and `0` is a number.
    #[test]
    fn a_cast_of_a_literal_keeps_the_type() {
        let src = r#"
function stop() external {
    launchpad.setListing(poolId, address(0), 0);
    sink.send(payable(0));
}
"#;
        assert_eq!(arguments_of(src, "setListing")[1], "address(0)");
        assert_eq!(arguments_of(src, "send")[0], "payable(0)");
    }

    /// `payable(x)` has its own node in the grammar and never reached the cast branch either.
    #[test]
    fn payable_is_the_name_inside_it() {
        let src = r#"
function pay() external {
    token.transfer(payable(feeSink), 1);
}
"#;
        assert_eq!(arguments_of(src, "transfer")[0], "feeSink");
    }

    /// A ternary is two possible parties, and which one runs is a question about the
    /// transaction. Both are recorded, so the map can draw a box for each.
    #[test]
    fn a_ternary_records_both_branches() {
        let src = r#"
function pay(bool burn) external {
    token.transfer(burn ? burnSink : treasury, 1);
}
"#;
        assert_eq!(arguments_of(src, "transfer")[0], "burnSink|treasury");
    }
}

#[cfg(test)]
mod creation_test {
    use super::analyze_body;

    /// WHICH variable holds what was deployed. `AscendLaunchpad`'s constructor runs
    /// `new Treasury(owner_)` five times into five variables — five addresses — and a map that
    /// knows only the type draws one arrow to one box for all of them.
    #[test]
    fn a_creation_remembers_the_variable_it_lands_in() {
        let src = r#"
constructor(address owner_) {
    opexTreasury = new Treasury(owner_);
    kntqTreasury = new Treasury(owner_);
    listingDesk = new ListingDesk{salt: bytes32(0)}(owner_, verifier_, 500);
    uint256 x = 1;
}
"#;
        let creations = analyze_body(src, &[], &[]).creations;
        assert!(creations.contains(&("opexTreasury".to_string(), "Treasury".to_string())));
        assert!(creations.contains(&("kntqTreasury".to_string(), "Treasury".to_string())));
        assert!(
            creations.contains(&("listingDesk".to_string(), "ListingDesk".to_string())),
            "a salted `new` is the same deployment: {creations:?}"
        );
    }
}

#[cfg(test)]
mod callback_test {
    use super::analyze_body;

    /// A call ON `msg.sender` goes back to whoever called — the language says so, and no name
    /// heuristic is involved. `Vault.lock` is the reason: `lock(bytes)` shows nothing, and the
    /// body hands the thread straight back to the caller.
    #[test]
    fn a_call_on_msg_sender_is_a_callback() {
        let src = r#"
function lock(bytes calldata data) external returns (bytes memory result) {
    SettlementGuard.setLocker(msg.sender);
    result = ILockCallback(msg.sender).lockAcquired(data);
}
"#;
        assert_eq!(analyze_body(src, &[], &[]).callbacks, vec!["lockAcquired".to_string()]);

        // Written without the cast, it is the same fact.
        let bare = r#"
function pull() external {
    msg.sender.onPulled(1);
}
"#;
        assert_eq!(analyze_body(bare, &[], &[]).callbacks, vec!["onPulled".to_string()]);

        // A call on anything else is not one.
        let ordinary = r#"
function pay() external {
    token.transfer(msg.sender, 1);
}
"#;
        assert!(analyze_body(ordinary, &[], &[]).callbacks.is_empty());
    }
}

#[cfg(test)]
mod encode_call_test {
    use super::analyze_body;

    /// `abi.encodeCall(I.f, (…))` IS the call, written where no walker looks for one: the bytes
    /// go straight into a low-level `call` or `staticcall`, often inside `assembly`. The vault
    /// reads its kHYPE rate that way and OpenZeppelin's `safeTransfer` reaches `token.transfer`
    /// that way, and neither reached the map.
    #[test]
    fn an_encoded_call_is_a_call() {
        let src = r#"
function _readRate() internal view returns (uint256) {
    bytes memory data = abi.encodeCall(IStakingAccountant.kHYPEToHYPE, (WAD));
    (bool ok, bytes memory out) = accountant.staticcall(data);
    return abi.decode(out, (uint256));
}
"#;
        let analysis = analyze_body(src, &[], &[]);
        assert!(
            analysis
                .call_targets
                .contains(&("IStakingAccountant".to_string(), "kHYPEToHYPE".to_string())),
            "the interface and the method are written down: {:?}",
            analysis.call_targets
        );
    }

    /// And its arguments come with it, so the parties of that call can be read.
    #[test]
    fn an_encoded_calls_arguments_are_kept() {
        let src = r#"
function _move(address to, uint256 value) internal {
    _callOptionalReturn(token, abi.encodeCall(token.transfer, (to, value)));
}
"#;
        let analysis = analyze_body(src, &[], &[]);
        let found = analysis
            .call_arguments
            .iter()
            .find(|(receiver, method, _)| receiver == "token" && method == "transfer")
            .map(|(_, _, args)| args.clone());
        assert_eq!(found, Some(vec!["to".to_string(), "value".to_string()]));
    }
}