fluidattacks-blends-domain 0.34.0

use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;

use crate::stack::module_path::normalize_module_path;
use crate::stack::ImportBinding;
use crate::syntax::{SyntaxGraph, SyntaxNode};
use crate::{Language, NodeId};

/// One `import` node resolved to what it binds.
#[derive(Clone, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct ImportSite {
    /// Id of the AST node holding the bound name.
    pub node: NodeId,
    /// Id of the scope the binding is visible in.
    pub scope: NodeId,
    /// The local name the import binds.
    pub symbol: String,
    /// What module (and symbol, if any) that name is bound to.
    pub binding: ImportBinding,
}

/// Every import site in `syntax`, for the languages this resolves imports for.
///
/// Returns an empty list for a `language` this has no import reader for.
#[must_use]
pub fn import_sites(syntax: &SyntaxGraph, language: Language) -> Vec<ImportSite> {
    let read_binding = match language {
        Language::Python => python_binding,
        Language::JavaScript | Language::TypeScript => script_binding,
        _ => return Vec::new(),
    };
    let bound_names: BTreeMap<NodeId, (NodeId, &str)> = syntax
        .symbol_index
        .iter()
        .flat_map(|(symbol, scopes)| {
            scopes.iter().flat_map(move |(scope, nodes)| {
                nodes
                    .iter()
                    .map(move |node| (*node, (*scope, symbol.as_str())))
            })
        })
        .collect();
    bound_names
        .into_iter()
        .filter_map(|(node, (scope, symbol))| {
            let binding = read_binding(syntax.nodes.get(&node)?, symbol)?;
            Some(ImportSite {
                node,
                scope,
                symbol: String::from(symbol),
                binding,
            })
        })
        .collect()
}

fn python_binding(node: &SyntaxNode, symbol: &str) -> Option<ImportBinding> {
    let SyntaxNode::Import {
        expression: Some(expression),
        alias,
        module_name,
        ..
    } = node
    else {
        return None;
    };
    let Some(module_name) = module_name else {
        let module = if alias.is_some() {
            expression.as_str()
        } else {
            symbol
        };
        return Some(module_binding(module, 0));
    };
    let module = module_name.trim_start_matches('.');
    let imported = expression.rsplit('.').next().unwrap_or(expression);
    Some(symbol_binding(
        module,
        level_of(module_name.len().saturating_sub(module.len())),
        imported,
        symbol,
    ))
}

fn script_binding(node: &SyntaxNode, symbol: &str) -> Option<ImportBinding> {
    match node {
        SyntaxNode::Import {
            alias: Some(_),
            module_name: Some(specifier),
            import_type,
            ..
        } => {
            let (module, relative_level) = script_module(specifier);
            match import_type.as_deref() {
                None => Some(symbol_binding(&module, relative_level, "default", symbol)),
                Some("namespace_import" | "import_require_clause") => {
                    Some(module_binding(&module, relative_level))
                }
                Some(_) => None,
            }
        }
        SyntaxNode::ModuleImport {
            expression,
            module_name: Some(specifier),
            ..
        } => {
            let imported = expression
                .strip_prefix(specifier.as_str())?
                .strip_prefix('.')?;
            let (module, relative_level) = script_module(specifier);
            Some(symbol_binding(&module, relative_level, imported, symbol))
        }
        _ => None,
    }
}

fn script_module(specifier: &str) -> (String, u32) {
    let mut segments = specifier.split('/').peekable();
    let mut relative = false;
    let mut parents = 0_usize;
    while let Some(segment) = segments.next_if(|segment| matches!(*segment, "." | "..")) {
        relative = true;
        parents = parents.saturating_add(usize::from(segment == ".."));
    }
    let module = normalize_module_path(&segments.collect::<Vec<_>>().join("/"));
    let relative_level = if relative {
        level_of(parents.saturating_add(1))
    } else {
        0
    };
    (module, relative_level)
}

fn module_binding(module: &str, relative_level: u32) -> ImportBinding {
    ImportBinding {
        module_stem: String::from(module),
        source_symbol: None,
        relative_level,
        binds_module: true,
    }
}

fn symbol_binding(
    module: &str,
    relative_level: u32,
    imported: &str,
    symbol: &str,
) -> ImportBinding {
    ImportBinding {
        module_stem: String::from(module),
        source_symbol: (imported != symbol).then(|| String::from(imported)),
        relative_level,
        binds_module: false,
    }
}

fn level_of(depth: usize) -> u32 {
    u32::try_from(depth).unwrap_or(u32::MAX)
}