pyrograph 0.1.0

use syn::visit::Visit;
use syn::{
    Block, Expr, ExprMacro, File, ImplItem, Item, ItemConst, ItemFn, ItemImpl, ItemStatic, Local,
    Pat, Stmt,
};

use super::{EdgeKind, EvalResult, FunctionInfo, NodeId, NodeKind, Scope, TaintGraph};

pub(super) struct RustParser {
    pub(super) graph: TaintGraph,
    current_file: std::path::PathBuf,
    pub(super) scopes: Vec<Scope>,
    pub(super) functions: std::collections::HashMap<String, FunctionInfo>,
    pub(super) closures: std::collections::HashMap<NodeId, Vec<NodeId>>,
    pub(super) current_function: Option<NodeId>,
    /// Node that receives taint from `break expr;` inside the current `loop`.
    pub(super) loop_result_node: Option<NodeId>,
    /// Maps imported short names to fully-qualified paths, e.g. `get_env` -> `std::env::var`.
    pub(super) use_aliases: std::collections::HashMap<String, String>,
    /// Cache of node IDs created for imported items so repeated references resolve consistently.
    pub(super) import_nodes: std::collections::HashMap<String, NodeId>,
}

impl RustParser {
    pub(super) fn new(current_file: std::path::PathBuf) -> Self {
        Self {
            graph: TaintGraph::new(),
            current_file,
            scopes: Vec::new(),
            functions: std::collections::HashMap::new(),
            closures: std::collections::HashMap::new(),
            current_function: None,
            loop_result_node: None,
            use_aliases: std::collections::HashMap::new(),
            import_nodes: std::collections::HashMap::new(),
        }
    }

    pub(super) fn into_graph(self) -> TaintGraph {
        self.graph
    }

    pub(super) fn push_scope(&mut self) {
        self.scopes.push(Scope::default());
    }

    pub(super) fn pop_scope(&mut self) {
        self.scopes.pop();
    }

    pub(super) fn define(&mut self, name: String, node: NodeId) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.define(name, node);
        }
    }

    pub(super) fn resolve(&self, name: &str) -> Option<NodeId> {
        self.scopes.iter().rev().find_map(|scope| scope.resolve(name))
    }

    pub(super) fn new_node(&mut self, kind: NodeKind, name: impl Into<String>) -> NodeId {
        self.graph.add_node(kind, name.into(), None)
    }

    pub(super) fn literal_result(&mut self, value: String) -> EvalResult {
        let node = self.new_node(NodeKind::Literal, value.clone());
        EvalResult {
            node,
            descriptor: None,
            literal: Some(value),
        }
    }

    pub(super) fn flow(&mut self, from: NodeId, to: NodeId, kind: EdgeKind) {
        self.graph.add_edge(from, to, kind);
    }

    pub(super) fn bind_pat(&mut self, pat: &Pat, src: Option<NodeId>) -> Vec<NodeId> {
        let mut bound = Vec::new();
        match pat {
            Pat::Ident(ident) => {
                let name = ident.ident.to_string();
                let node = self.new_node(NodeKind::Variable, name.clone());
                self.define(name, node);
                if let Some(src) = src {
                    self.flow(src, node, EdgeKind::Assignment);
                }
                bound.push(node);
            }
            Pat::Tuple(tuple) => {
                for elem in &tuple.elems {
                    bound.extend(self.bind_pat(elem, src));
                }
            }
            Pat::Struct(strukt) => {
                for field in &strukt.fields {
                    bound.extend(self.bind_pat(&field.pat, src));
                }
            }
            Pat::TupleStruct(tuple) => {
                for elem in &tuple.elems {
                    bound.extend(self.bind_pat(elem, src));
                }
            }
            Pat::Reference(reference) => {
                bound.extend(self.bind_pat(&reference.pat, src));
            }
            Pat::Type(typed) => {
                bound.extend(self.bind_pat(&typed.pat, src));
            }
            Pat::Slice(slice) => {
                for elem in &slice.elems {
                    bound.extend(self.bind_pat(elem, src));
                }
            }
            Pat::Or(or_pat) => {
                for case in &or_pat.cases {
                    bound.extend(self.bind_pat(case, src));
                }
            }
            Pat::Paren(paren) => {
                bound.extend(self.bind_pat(&paren.pat, src));
            }
            Pat::Macro(mac) => {
                let name = mac.mac.path.segments.last().map(|s| s.ident.to_string()).unwrap_or_else(|| "macro_pat".into());
                let node = self.new_node(NodeKind::Variable, name);
                if let Some(src) = src {
                    self.flow(src, node, EdgeKind::Assignment);
                }
                bound.push(node);
            }
            Pat::Wild(_) => {}
            Pat::Lit(_) => {}
            Pat::Const(_) => {}
            Pat::Range(_) => {}
            Pat::Rest(_) => {}
            _ => {}
        }
        bound
    }

    pub(super) fn path_to_string(path: &syn::Path) -> String {
        path.segments
            .iter()
            .map(|segment| segment.ident.to_string())
            .collect::<Vec<_>>()
            .join("::")
    }

    /// Expand the first segment of a path if it matches a `use` alias.
    /// e.g. `get_env` -> `std::env::var`, `env::var` -> `std::env::var`.
    pub(super) fn expand_use_alias(&self, name: &str) -> Option<String> {
        if let Some((first, rest)) = name.split_once("::") {
            self.use_aliases.get(first).map(|prefix| format!("{}::{}", prefix, rest))
        } else {
            self.use_aliases.get(name).cloned()
        }
    }

    fn collect_use_tree(&mut self, tree: &syn::UseTree, prefix: &str) {
        match tree {
            syn::UseTree::Path(path) => {
                let new_prefix = if prefix.is_empty() {
                    path.ident.to_string()
                } else {
                    format!("{}::{}", prefix, path.ident)
                };
                self.collect_use_tree(&path.tree, &new_prefix);
            }
            syn::UseTree::Name(name) => {
                let full = if prefix.is_empty() {
                    name.ident.to_string()
                } else {
                    format!("{}::{}", prefix, name.ident)
                };
                self.use_aliases.insert(name.ident.to_string(), full);
            }
            syn::UseTree::Rename(rename) => {
                let full = if prefix.is_empty() {
                    rename.ident.to_string()
                } else {
                    format!("{}::{}", prefix, rename.ident)
                };
                self.use_aliases.insert(rename.rename.to_string(), full);
            }
            syn::UseTree::Glob(_) => {
                // Expand known dangerous modules so taint analysis can
                // resolve common names after a glob import.
                let expanded: &[&str] = match prefix {
                    "std::env" => &["var", "vars", "set_var", "remove_var", "current_dir", "home_dir", "args", "args_os"],
                    "std::process" => &["Command", "Child", "Stdio", "Output", "ExitStatus", "id"],
                    "std::fs" => &["File", "OpenOptions", "read", "write", "read_to_string", "copy", "remove_file", "remove_dir", "create_dir", "metadata"],
                    "std::net" => &["TcpStream", "TcpListener", "UdpSocket", "SocketAddr", "IpAddr", "Ipv4Addr", "Ipv6Addr"],
                    _ => &[],
                };
                for name in expanded {
                    self.use_aliases.insert(name.to_string(), format!("{}::{}", prefix, name));
                }
            }
            syn::UseTree::Group(group) => {
                for item in &group.items {
                    self.collect_use_tree(item, prefix);
                }
            }
        }
    }

    pub(super) fn infer_descriptor(name: &str) -> Option<String> {
        for candidate in [
            "Command",
            "Client",
            "TcpStream",
            "File",
            "OpenOptions",
            "Path",
            "PathBuf",
        ] {
            if name == candidate || name.ends_with(&format!("::{candidate}")) {
                return Some(candidate.to_string());
            }
        }
        if name.ends_with("::home_dir") || name.ends_with("::config_dir") {
            return Some("Path".into());
        }
        if let Some(prefix) = name.strip_suffix("::new") {
            return Self::infer_descriptor(prefix);
        }
        None
    }

    pub(super) fn visit_local_stmt(&mut self, local: &Local) {
        let init = local.init.as_ref().map(|init| self.eval_expr(&init.expr));
        let src = init.as_ref().map(|value| value.node);
        self.bind_pat(&local.pat, src);
    }

    pub(super) fn visit_const_like(&mut self, pat: &Pat, expr: &Expr) {
        let value = self.eval_expr(expr);
        self.bind_pat(pat, Some(value.node));
    }

    pub(super) fn visit_block_scoped(&mut self, block: &Block) {
        self.push_scope();
        for stmt in &block.stmts {
            self.visit_stmt(stmt);
        }
        self.pop_scope();
    }
}

impl Visit<'_> for RustParser {
    fn visit_file(&mut self, node: &File) {
        let _ = &self.current_file;
        self.push_scope();
        for item in &node.items {
            self.visit_item(item);
        }
        self.pop_scope();
    }

    fn visit_item(&mut self, node: &Item) {
        match node {
            Item::Const(ItemConst { ident, expr, .. }) => {
                self.visit_const_like(
                    &Pat::Ident(syn::PatIdent {
                        attrs: Vec::new(),
                        by_ref: None,
                        mutability: None,
                        ident: ident.clone(),
                        subpat: None,
                    }),
                    expr,
                );
            }
            Item::Fn(item_fn) => self.visit_item_fn(item_fn),
            Item::Impl(item_impl) => self.visit_item_impl(item_impl),
            Item::Macro(mac) => {
                let expr = Expr::Macro(ExprMacro {
                    attrs: mac.attrs.clone(),
                    mac: mac.mac.clone(),
                });
                self.eval_expr(&expr);
            }
            Item::Mod(module) => {
                if let Some((_, items)) = &module.content {
                    for item in items {
                        self.visit_item(item);
                    }
                }
            }
            Item::Use(item_use) => {
                self.collect_use_tree(&item_use.tree, "");
            }
            Item::Static(ItemStatic { ident, expr, .. }) => {
                self.visit_const_like(
                    &Pat::Ident(syn::PatIdent {
                        attrs: Vec::new(),
                        by_ref: None,
                        mutability: None,
                        ident: ident.clone(),
                        subpat: None,
                    }),
                    expr,
                );
            }
            _ => {}
        }
    }

    fn visit_item_fn(&mut self, node: &ItemFn) {
        let name = node.sig.ident.to_string();
        let fn_node = self.new_node(NodeKind::Variable, name.clone());
        self.define(name.clone(), fn_node);
        let prev_fn = self.current_function.replace(fn_node);
        self.push_scope();
        let mut params = Vec::new();
        for input in &node.sig.inputs {
            if let syn::FnArg::Typed(param) = input {
                params.extend(self.bind_pat(&param.pat, None));
            }
        }
        self.functions.insert(name, FunctionInfo { node: fn_node, params });
        let mut last_result = None;
        for stmt in &node.block.stmts {
            self.visit_stmt(stmt);
            // Track the last expression statement (Rust implicit return)
            if let Stmt::Expr(expr, None) = stmt {
                last_result = Some(self.eval_expr(expr));
            }
        }
        // Wire implicit return: last expression without semicolon → fn_node
        if let Some(result) = last_result {
            self.flow(result.node, fn_node, EdgeKind::Return);
        }
        self.pop_scope();
        self.current_function = prev_fn;
    }

    fn visit_item_impl(&mut self, node: &ItemImpl) {
        for item in &node.items {
            match item {
                ImplItem::Fn(method) => {
                    let name = method.sig.ident.to_string();
                    let fn_node = self.new_node(NodeKind::Variable, name.clone());
                    self.define(name.clone(), fn_node);
                    let prev_fn = self.current_function.replace(fn_node);
                    self.push_scope();
                    let mut params = Vec::new();
                    for input in &method.sig.inputs {
                        if let syn::FnArg::Typed(param) = input {
                            params.extend(self.bind_pat(&param.pat, None));
                        }
                    }
                    self.functions.insert(name, FunctionInfo { node: fn_node, params });
                    let mut last_result = None;
                    for stmt in &method.block.stmts {
                        self.visit_stmt(stmt);
                        if let Stmt::Expr(expr, None) = stmt {
                            last_result = Some(self.eval_expr(expr));
                        }
                    }
                    if let Some(result) = last_result {
                        self.flow(result.node, fn_node, EdgeKind::Return);
                    }
                    self.pop_scope();
                    self.current_function = prev_fn;
                }
                ImplItem::Const(item_const) => {
                    self.visit_const_like(
                        &Pat::Ident(syn::PatIdent {
                            attrs: Vec::new(),
                            by_ref: None,
                            mutability: None,
                            ident: item_const.ident.clone(),
                            subpat: None,
                        }),
                        &item_const.expr,
                    );
                }
                ImplItem::Type(_) => {}
                ImplItem::Macro(mac) => {
                    let expr = Expr::Macro(ExprMacro {
                        attrs: mac.attrs.clone(),
                        mac: mac.mac.clone(),
                    });
                    self.eval_expr(&expr);
                }
                _ => {}
            }
        }
    }

    fn visit_stmt(&mut self, node: &Stmt) {
        match node {
            Stmt::Expr(expr, _) => {
                self.eval_expr(expr);
            }
            Stmt::Item(item) => self.visit_item(item),
            Stmt::Local(local) => self.visit_local_stmt(local),
            Stmt::Macro(mac) => {
                let expr = Expr::Macro(ExprMacro {
                    attrs: mac.attrs.clone(),
                    mac: mac.mac.clone(),
                });
                self.eval_expr(&expr);
            }
        }
    }
}