pyrograph 0.1.0

use syn::parse::Parser;
use syn::visit::Visit;
use syn::{BinOp, Expr, ExprBinary, ExprCall, ExprMacro, ExprMethodCall, ExprPath, Lit, Stmt};

use super::{EdgeKind, EvalResult, NodeKind, RustParser};

impl RustParser {
    pub(super) fn expr_path_name(path: &ExprPath) -> String {
        Self::path_to_string(&path.path)
    }

    pub(super) fn eval_path(&mut self, path: &ExprPath) -> EvalResult {
        let name = Self::expr_path_name(path);
        // Try local resolution first (shadowed bindings take precedence).
        if let Some(node) = self.resolve(&name) {
            return EvalResult {
                node,
                descriptor: Self::infer_descriptor(&name),
                literal: None,
            };
        }
        if let Some(node) = path.path.get_ident().and_then(|ident| self.resolve(&ident.to_string())) {
            return EvalResult {
                node,
                descriptor: Self::infer_descriptor(&name),
                literal: None,
            };
        }
        // Expand `use` aliases so that `use std::env::var as v; v(...)` matches `env::var`.
        let expanded = self.expand_use_alias(&name);
        if let Some(ref exp) = expanded {
            if let Some(&node) = self.import_nodes.get(exp) {
                return EvalResult {
                    node,
                    descriptor: Self::infer_descriptor(exp),
                    literal: None,
                };
            }
        }
        let final_name = expanded.clone().unwrap_or_else(|| name.clone());
        let node = self.new_node(NodeKind::Variable, final_name.clone());
        if let Some(exp) = expanded {
            self.import_nodes.insert(exp, node);
        }
        EvalResult {
            node,
            descriptor: Self::infer_descriptor(&final_name),
            literal: None,
        }
    }

    pub(super) fn eval_macro(&mut self, mac: &ExprMacro) -> EvalResult {
        let name = Self::path_to_string(&mac.mac.path);
        let call_name = if name == "include_str" {
            "include_str!".to_string()
        } else {
            format!("{name}!")
        };
        let node = self.new_node(NodeKind::Call, call_name);

        // Heuristic: format-like macros have comma-separated expressions.
        // Parse them manually so taint from arguments flows to the macro result.
        if matches!(name.as_str(), "format" | "format_args" | "println" | "eprintln" | "write" | "writeln") {
            let parser = syn::punctuated::Punctuated::<Expr, syn::Token![,]>::parse_terminated;
            if let Ok(args) = parser.parse2(mac.mac.tokens.clone()) {
                let mut descriptor = None;
                let mut literal = None;
                for arg in &args {
                    let value = self.eval_expr(arg);
                    self.flow(value.node, node, EdgeKind::Argument);
                    if descriptor.is_none() {
                        descriptor = value.descriptor;
                    }
                    if literal.is_none() {
                        literal = value.literal.clone();
                    }
                }
                return EvalResult { node, descriptor, literal };
            }
        }

        if let Ok(expr) = syn::parse2::<Expr>(mac.mac.tokens.clone()) {
            let arg = self.eval_expr(&expr);
            self.flow(arg.node, node, EdgeKind::Argument);
            return EvalResult {
                node,
                descriptor: arg.descriptor,
                literal: arg.literal,
            };
        }
        let token_text = mac.mac.tokens.to_string();
        let lit = self.literal_result(token_text);
        self.flow(lit.node, node, EdgeKind::Argument);
        EvalResult {
            node,
            descriptor: None,
            literal: lit.literal,
        }
    }

    pub(super) fn eval_call(&mut self, call: &ExprCall) -> EvalResult {
        let callee = self.eval_expr(&call.func);
        let call_name = self
            .graph
            .node(callee.node)
            .map(|node| node.name.clone())
            .unwrap_or_else(|| "call".into());
        let node = self.new_node(NodeKind::Call, call_name.clone());
        self.flow(callee.node, node, EdgeKind::Assignment);
        for arg in &call.args {
            let value = self.eval_expr(arg);
            self.flow(value.node, node, EdgeKind::Argument);
        }
        // Interprocedural taint: wire call arguments to function parameters.
        // 1. Direct name lookup (normal calls)
        // 2. Follow INCOMING Assignment edges to callee (fn pointer indirection:
        //    `let f = execute; f(arg)` — edge is execute→f, so we search for
        //    any node that assigned TO callee.node)
        let resolved_fn = self.functions.get(&call_name).cloned().or_else(|| {
            // Scan for nodes with an outgoing Assignment edge TO callee.node
            (0..self.graph.node_count() as u32).find_map(|id| {
                let has_edge = self.graph.edges_from(id).any(|(target, kind)| {
                    target == callee.node && kind == EdgeKind::Assignment
                });
                if has_edge {
                    let name = self.graph.node(id).map(|n| n.name.clone())?;
                    self.functions.get(&name).cloned()
                } else {
                    None
                }
            })
        });
        if let Some(function) = resolved_fn {
            for (arg, param) in call.args.iter().zip(&function.params) {
                let value = self.eval_expr(arg);
                self.flow(value.node, *param, EdgeKind::Assignment);
            }
            self.flow(function.node, node, EdgeKind::Return);
        }
        // Closure argument interprocedural taint: wire call arguments to closure parameters.
        let closure_node = self.closures.get(&callee.node).map(|_| callee.node).or_else(|| {
            self.graph
                .edges_from(callee.node)
                .find_map(|(target, kind)| {
                    if kind == EdgeKind::Assignment && self.closures.contains_key(&target) {
                        Some(target)
                    } else {
                        None
                    }
                })
        });
        if let Some(closure_node) = closure_node {
            if let Some(params) = self.closures.get(&closure_node).cloned() {
                for (arg, param) in call.args.iter().zip(&params) {
                    let value = self.eval_expr(arg);
                    self.flow(value.node, *param, EdgeKind::Assignment);
                }
                // The closure body already wired its result to the closure node via Return.
                // Flow closure node to call result so taint propagates.
                self.flow(closure_node, node, EdgeKind::Return);
            }
        }
        EvalResult {
            node,
            descriptor: Self::infer_descriptor(&call_name).or(callee.descriptor),
            literal: None,
        }
    }

    pub(super) fn eval_method_call(&mut self, method: &ExprMethodCall) -> EvalResult {
        let receiver = self.eval_expr(&method.receiver);
        let base = receiver
            .descriptor
            .clone()
            .map(|desc| if desc.contains("Path") { "Path".to_string() } else { desc })
            .unwrap_or_else(|| "value".into());
        let method_name = method.method.to_string();
        let node = self.new_node(NodeKind::Call, format!(".{method_name}"));
        self.flow(receiver.node, node, EdgeKind::Assignment);
        for arg in &method.args {
            let value = self.eval_expr(arg);
            self.flow(value.node, node, EdgeKind::Argument);
            // &mut arguments receive data FROM the call (e.g., read_to_string(&mut buf)
            // fills buf with the receiver's data). Flow taint back to the argument.
            if let Expr::Reference(ref_expr) = arg {
                if ref_expr.mutability.is_some() {
                    self.flow(node, value.node, EdgeKind::Assignment);
                }
            }
        }
        let descriptor = if method_name == "file_name" {
            Some("Path".into())
        } else if matches!(
            method_name.as_str(),
            "new" | "arg" | "args" | "env" | "stdout" | "stderr" | "current_dir" | "spawn"
                | "output" | "status" | "post" | "body" | "header" | "bearer_auth" | "send"
                | "join" | "push" | "unwrap" | "expect" | "open"
        ) {
            Some(base)
        } else {
            receiver.descriptor
        };
        EvalResult {
            node,
            descriptor,
            literal: None,
        }
    }

    pub(super) fn eval_binary(&mut self, binary: &ExprBinary) -> EvalResult {
        let lhs = self.eval_expr(&binary.left);
        let rhs = self.eval_expr(&binary.right);
        let node = self.new_node(NodeKind::Call, "binary".to_string());
        self.flow(lhs.node, node, EdgeKind::Argument);
        self.flow(rhs.node, node, EdgeKind::Argument);
        if matches!(binary.op, BinOp::Add(_)) {
            if let (Some(left), Some(right)) = (lhs.literal, rhs.literal) {
                return self.literal_result(format!("{left}{right}"));
            }
        }
        EvalResult {
            node,
            descriptor: None,
            literal: None,
        }
    }

    pub(super) fn eval_expr(&mut self, expr: &Expr) -> EvalResult {
        match expr {
            Expr::Array(arr) => {
                let node = self.new_node(NodeKind::Call, "array".to_string());
                for elem in &arr.elems {
                    let v = self.eval_expr(elem);
                    self.flow(v.node, node, EdgeKind::Argument);
                }
                EvalResult {
                    node,
                    descriptor: None,
                    literal: None,
                }
            }
            Expr::Assign(assign) => {
                let value = self.eval_expr(&assign.right);
                if let Expr::Path(path) = &*assign.left {
                    let target = self.eval_path(path);
                    self.flow(value.node, target.node, EdgeKind::Assignment);
                    target
                } else if let Expr::Field(field) = &*assign.left {
                    let base = self.eval_expr(&field.base);
                    self.flow(value.node, base.node, EdgeKind::Assignment);
                    EvalResult {
                        node: base.node,
                        descriptor: base.descriptor,
                        literal: None,
                    }
                } else if let Expr::Index(index) = &*assign.left {
                    let base = self.eval_expr(&index.expr);
                    self.flow(value.node, base.node, EdgeKind::Assignment);
                    EvalResult {
                        node: base.node,
                        descriptor: base.descriptor,
                        literal: None,
                    }
                } else {
                    value
                }
            }
            Expr::Async(block) => {
                let result = self.eval_block(&block.block);
                let node = self.new_node(NodeKind::Call, "async".to_string());
                self.flow(result.node, node, EdgeKind::Return);
                EvalResult {
                    node,
                    descriptor: result.descriptor,
                    literal: result.literal,
                }
            }
            Expr::Await(await_expr) => {
                let base = self.eval_expr(&await_expr.base);
                EvalResult {
                    node: base.node,
                    descriptor: base.descriptor,
                    literal: base.literal,
                }
            }
            Expr::Binary(binary) => self.eval_binary(binary),
            Expr::Block(block) => self.eval_block(&block.block),
            Expr::Break(brk) => {
                let value = brk
                    .expr
                    .as_ref()
                    .map(|e| self.eval_expr(e))
                    .unwrap_or_else(|| self.literal_result("break".into()));
                if let Some(loop_node) = self.loop_result_node {
                    self.flow(value.node, loop_node, EdgeKind::Assignment);
                }
                value
            }
            Expr::Call(call) => self.eval_call(call),
            Expr::Cast(cast) => self.eval_expr(&cast.expr),
            Expr::Closure(closure) => {
                let node = self.new_node(NodeKind::Call, "closure".to_string());
                self.push_scope();
                let mut params = Vec::new();
                for input in &closure.inputs {
                    params.extend(self.bind_pat(input, None));
                }
                let body = self.eval_expr(&closure.body);
                self.flow(body.node, node, EdgeKind::Return);
                self.pop_scope();
                self.closures.insert(node, params);
                EvalResult {
                    node,
                    descriptor: body.descriptor,
                    literal: body.literal,
                }
            }
            Expr::Const(block) => {
                let result = self.eval_block(&block.block);
                let node = self.new_node(NodeKind::Call, "const_block".to_string());
                self.flow(result.node, node, EdgeKind::Return);
                EvalResult {
                    node,
                    descriptor: result.descriptor,
                    literal: result.literal,
                }
            }
            Expr::Continue(_) => self.literal_result("continue".into()),
            Expr::Field(field) => {
                let base_result = self.eval_expr(&field.base);
                let member_name = match &field.member {
                    syn::Member::Named(ident) => ident.to_string(),
                    syn::Member::Unnamed(index) => index.index.to_string(),
                };
                let name = format!(
                    "{}.{}",
                    self.graph
                        .node(base_result.node)
                        .map(|n| n.name.as_str())
                        .unwrap_or("value"),
                    member_name
                );
                let node = self.new_node(NodeKind::Variable, name);
                self.flow(base_result.node, node, EdgeKind::Assignment);
                EvalResult {
                    node,
                    descriptor: base_result.descriptor,
                    literal: None,
                }
            }
            Expr::ForLoop(for_loop) => {
                let iter = self.eval_expr(&for_loop.expr);
                self.push_scope();
                self.bind_pat(&for_loop.pat, Some(iter.node));
                for stmt in &for_loop.body.stmts {
                    self.visit_stmt(stmt);
                }
                self.pop_scope();
                iter
            }
            Expr::Group(group) => self.eval_expr(&group.expr),
            Expr::If(if_expr) => {
                self.push_scope();
                let cond_value = self.eval_expr(&if_expr.cond);
                let mut then_result = self.literal_result("block".into());
                for (i, stmt) in if_expr.then_branch.stmts.iter().enumerate() {
                    if i == if_expr.then_branch.stmts.len() - 1 {
                        match stmt {
                            Stmt::Expr(expr, _) => then_result = self.eval_expr(expr),
                            Stmt::Local(local) => self.visit_local_stmt(local),
                            Stmt::Item(item) => self.visit_item(item),
                            Stmt::Macro(mac) => {
                                let expr = Expr::Macro(ExprMacro {
                                    attrs: mac.attrs.clone(),
                                    mac: mac.mac.clone(),
                                });
                                self.eval_expr(&expr);
                            }
                        }
                    } else {
                        self.visit_stmt(stmt);
                    }
                }
                self.pop_scope();
                let else_result = if_expr.else_branch.as_ref().map(|(_, e)| self.eval_expr(e));
                let node = self.new_node(NodeKind::Call, "if".to_string());
                self.flow(cond_value.node, node, EdgeKind::Argument);
                self.flow(then_result.node, node, EdgeKind::Assignment);
                if let Some(else_res) = &else_result {
                    self.flow(else_res.node, node, EdgeKind::Assignment);
                }
                EvalResult {
                    node,
                    descriptor: then_result.descriptor.or_else(|| else_result.as_ref().and_then(|r| r.descriptor.clone())),
                    literal: then_result.literal.or_else(|| else_result.as_ref().and_then(|r| r.literal.clone())),
                }
            }
            Expr::Index(index) => {
                let base = self.eval_expr(&index.expr);
                let idx = self.eval_expr(&index.index);
                let node = self.new_node(NodeKind::Variable, "index".to_string());
                self.flow(base.node, node, EdgeKind::Assignment);
                self.flow(idx.node, node, EdgeKind::Argument);
                EvalResult {
                    node,
                    descriptor: base.descriptor,
                    literal: None,
                }
            }
            Expr::Infer(_) => self.literal_result("infer".into()),
            Expr::Let(let_expr) => {
                let value = self.eval_expr(&let_expr.expr);
                self.bind_pat(&let_expr.pat, Some(value.node));
                EvalResult {
                    node: value.node,
                    descriptor: value.descriptor,
                    literal: value.literal,
                }
            }
            Expr::Lit(lit) => match &lit.lit {
                Lit::Str(value) => self.literal_result(value.value()),
                Lit::ByteStr(value) => {
                    self.literal_result(String::from_utf8_lossy(&value.value()).into_owned())
                }
                Lit::Int(value) => self.literal_result(value.base10_digits().to_string()),
                Lit::Bool(value) => self.literal_result(value.value.to_string()),
                _ => self.literal_result("literal".into()),
            },
            Expr::Loop(loop_expr) => {
                let node = self.new_node(NodeKind::Call, "loop".to_string());
                let prev_loop = self.loop_result_node.replace(node);
                self.visit_block_scoped(&loop_expr.body);
                self.loop_result_node = prev_loop;
                EvalResult {
                    node,
                    descriptor: None,
                    literal: None,
                }
            }
            Expr::Macro(mac) => self.eval_macro(mac),
            Expr::Match(match_expr) => {
                let source = self.eval_expr(&match_expr.expr);
                let node = self.new_node(NodeKind::Call, "match".to_string());
                self.flow(source.node, node, EdgeKind::Argument);
                let mut descriptor = None;
                let mut literal = None;
                for arm in &match_expr.arms {
                    self.push_scope();
                    self.bind_pat(&arm.pat, Some(source.node));
                    if let Some((_, guard)) = &arm.guard {
                        self.eval_expr(guard);
                    }
                    let arm_result = self.eval_expr(&arm.body);
                    self.flow(arm_result.node, node, EdgeKind::Assignment);
                    if descriptor.is_none() {
                        descriptor = arm_result.descriptor;
                    }
                    if literal.is_none() {
                        literal = arm_result.literal.clone();
                    }
                    self.pop_scope();
                }
                EvalResult {
                    node,
                    descriptor,
                    literal,
                }
            }
            Expr::MethodCall(method) => self.eval_method_call(method),
            Expr::Paren(paren) => self.eval_expr(&paren.expr),
            Expr::Path(path) => self.eval_path(path),
            Expr::Range(range) => {
                let node = self.new_node(NodeKind::Call, "range".to_string());
                if let Some(start) = &range.start {
                    let v = self.eval_expr(start);
                    self.flow(v.node, node, EdgeKind::Argument);
                }
                if let Some(end) = &range.end {
                    let v = self.eval_expr(end);
                    self.flow(v.node, node, EdgeKind::Argument);
                }
                EvalResult {
                    node,
                    descriptor: None,
                    literal: None,
                }
            }
            Expr::RawAddr(raw) => self.eval_expr(&raw.expr),
            Expr::Reference(reference) => self.eval_expr(&reference.expr),
            Expr::Repeat(repeat) => {
                let value = self.eval_expr(&repeat.expr);
                let node = self.new_node(NodeKind::Call, "repeat".to_string());
                self.flow(value.node, node, EdgeKind::Argument);
                EvalResult {
                    node,
                    descriptor: value.descriptor,
                    literal: value.literal,
                }
            }
            Expr::Return(ret) => {
                let value = ret.expr.as_ref().map(|inner| self.eval_expr(inner));
                if let (Some(value), Some(current_fn)) = (value, self.current_function) {
                    self.flow(value.node, current_fn, EdgeKind::Return);
                    value
                } else {
                    self.literal_result("return".into())
                }
            }
            Expr::Struct(strukt) => {
                let node = self.new_node(NodeKind::Call, "struct".to_string());
                let mut descriptor = None;
                let mut literal = None;
                for field in &strukt.fields {
                    let val = self.eval_expr(&field.expr);
                    self.flow(val.node, node, EdgeKind::Argument);
                    if descriptor.is_none() {
                        descriptor = val.descriptor;
                    }
                    if literal.is_none() {
                        literal = val.literal.clone();
                    }
                }
                if let Some(rest) = &strukt.rest {
                    let val = self.eval_expr(rest);
                    self.flow(val.node, node, EdgeKind::Argument);
                }
                EvalResult {
                    node,
                    descriptor,
                    literal,
                }
            }
            Expr::Try(try_expr) => {
                let base = self.eval_expr(&try_expr.expr);
                let node = self.new_node(NodeKind::Call, "try".to_string());
                self.flow(base.node, node, EdgeKind::Assignment);
                EvalResult {
                    node,
                    descriptor: base.descriptor,
                    literal: base.literal,
                }
            }
            Expr::TryBlock(block) => {
                let result = self.eval_block(&block.block);
                let node = self.new_node(NodeKind::Call, "try_block".to_string());
                self.flow(result.node, node, EdgeKind::Return);
                EvalResult {
                    node,
                    descriptor: result.descriptor,
                    literal: result.literal,
                }
            }
            Expr::Tuple(tuple) => {
                let node = self.new_node(NodeKind::Call, "tuple".to_string());
                let mut descriptor = None;
                let mut literal = None;
                for elem in &tuple.elems {
                    let e = self.eval_expr(elem);
                    self.flow(e.node, node, EdgeKind::Argument);
                    if descriptor.is_none() {
                        descriptor = e.descriptor;
                    }
                    if literal.is_none() {
                        literal = e.literal.clone();
                    }
                }
                EvalResult {
                    node,
                    descriptor,
                    literal,
                }
            }
            Expr::Unary(unary) => self.eval_expr(&unary.expr),
            Expr::Unsafe(block) => {
                let result = self.eval_block(&block.block);
                let node = self.new_node(NodeKind::Call, "unsafe".to_string());
                self.flow(result.node, node, EdgeKind::Return);
                EvalResult {
                    node,
                    descriptor: result.descriptor,
                    literal: result.literal,
                }
            }
            Expr::Verbatim(_) => self.literal_result("verbatim".into()),
            Expr::While(while_expr) => {
                self.push_scope();
                self.eval_expr(&while_expr.cond);
                for stmt in &while_expr.body.stmts {
                    self.visit_stmt(stmt);
                }
                self.pop_scope();
                self.literal_result("while".into())
            }
            Expr::Yield(yield_expr) => yield_expr
                .expr
                .as_ref()
                .map(|e| self.eval_expr(e))
                .unwrap_or_else(|| self.literal_result("yield".into())),
            _ => self.literal_result("expr".into()),
        }
    }

    pub(super) fn eval_block(&mut self, block: &syn::Block) -> EvalResult {
        self.push_scope();
        let mut result = self.literal_result("block".into());
        for (i, stmt) in block.stmts.iter().enumerate() {
            if i == block.stmts.len() - 1 {
                match stmt {
                    Stmt::Expr(expr, _) => result = self.eval_expr(expr),
                    Stmt::Local(local) => self.visit_local_stmt(local),
                    Stmt::Item(item) => self.visit_item(item),
                    Stmt::Macro(mac) => {
                        let expr = Expr::Macro(ExprMacro {
                            attrs: mac.attrs.clone(),
                            mac: mac.mac.clone(),
                        });
                        self.eval_expr(&expr);
                    }
                }
            } else {
                self.visit_stmt(stmt);
            }
        }
        self.pop_scope();
        result
    }
}