fluidattacks_blends/

ast.rs

//! Walk a tree-sitter parse tree into a `domain::ast::AstGraph`.

use std::collections::HashMap;

use blends_domain::ast::{AstGraph, AstNode};
use blends_domain::NodeId;
use tree_sitter::Node;

use crate::content::Content;
use crate::language::Language;
use crate::parse::parse;

const fn needs_field_resolution(language: Language) -> bool {
    matches!(language, Language::Swift)
}

/// Whether a node with children must NOT be descended into (its subtree is
/// collapsed to a single node)
fn is_node_terminal(language: Language, kind: &str) -> bool {
    matches!(
        (language, kind),
        (Language::Swift, "bang")
            | (
                Language::Yaml,
                "single_quote_scalar" | "double_quote_scalar"
            )
    )
}

/// Whether a non-leaf node should still capture its source text, ported from
/// the python `_has_content_node` table (dart arrives with that language).
fn has_content_node(language: Language, kind: &str) -> bool {
    matches!(
        (language, kind),
        (
            Language::CSharp,
            "character_literal" | "predefined_type" | "string_literal" | "verbatim_string_literal"
        ) | (Language::Go, "interpreted_string_literal")
            | (
                Language::JavaScript | Language::TypeScript,
                "template_string" | "regex"
            )
            | (Language::Kotlin, "string_literal")
            | (Language::Swift, "bang")
            | (
                Language::Yaml,
                "single_quote_scalar" | "double_quote_scalar"
            )
    )
}

fn to_u32(value: usize) -> u32 {
    u32::try_from(value).unwrap_or(u32::MAX)
}

fn decode_latin1(bytes: &[u8]) -> String {
    bytes.iter().map(|&byte| char::from(byte)).collect()
}

fn build_node(
    content: &Content,
    graph: &mut AstGraph,
    node: Node<'_>,
    counter: &mut u64,
) -> NodeId {
    *counter = counter.saturating_add(1);
    let id = NodeId(*counter);

    let start = node.start_position();
    let mut ast_node = AstNode::new(
        to_u32(start.row.saturating_add(1)),
        to_u32(start.column.saturating_add(1)),
        node.kind().to_owned(),
    );

    let child_count = node.child_count();
    if child_count == 0 || has_content_node(content.language, node.kind()) {
        ast_node.text = Some(
            content
                .bytes
                .get(node.start_byte()..node.end_byte())
                .map(decode_latin1)
                .unwrap_or_default(),
        );
    }
    graph.add_node(id, ast_node);

    if child_count > 0 && !is_node_terminal(content.language, node.kind()) {
        build_children(content, graph, node, id, counter);
    }

    id
}

type NodeKey = (usize, usize, &'static str);

fn node_key(node: &Node<'_>) -> NodeKey {
    (node.start_byte(), node.end_byte(), node.kind())
}

fn field_by_child(parent: Node<'_>) -> HashMap<NodeKey, &'static str> {
    let language = parent.language();
    let mut field_ids: Vec<u16> = (1..=language.field_count())
        .filter_map(|id| u16::try_from(id).ok())
        .collect();
    field_ids.sort_by_key(|&id| language.field_name_for_id(id));

    let mut map = HashMap::new();
    for id in field_ids {
        if let (Some(name), Some(child)) =
            (language.field_name_for_id(id), parent.child_by_field_id(id))
        {
            map.insert(node_key(&child), name);
        }
    }
    map
}

fn build_children(
    content: &Content,
    graph: &mut AstGraph,
    parent: Node<'_>,
    parent_id: NodeId,
    counter: &mut u64,
) {
    let overrides = needs_field_resolution(content.language).then(|| field_by_child(parent));

    let mut cursor = parent.walk();
    if !cursor.goto_first_child() {
        return;
    }

    let mut index: u32 = 0;
    loop {
        let child = cursor.node();
        let field = overrides.as_ref().map_or_else(
            || cursor.field_name(),
            |fields| fields.get(&node_key(&child)).copied(),
        );
        let child_id = build_node(content, graph, child, counter);
        graph.add_edge(parent_id, child_id, index);

        if let Some(name) = field {
            graph.set_field(parent_id, format!("{name}_id"), child_id);
        }

        index = index.saturating_add(1);
        if !cursor.goto_next_sibling() {
            break;
        }
    }
}

pub fn get_ast_graph(content: &Content) -> Option<AstGraph> {
    let Ok(tree) = parse(content) else {
        tracing::warn!(path = %content.path.display(), "Unable to parse possibly malformed file");
        return None;
    };

    let mut graph = AstGraph::new();
    let mut counter: u64 = 0;
    build_node(content, &mut graph, tree.root_node(), &mut counter);

    Some(graph)
}

#[cfg(test)]
mod tests {
    use super::get_ast_graph;
    use crate::content::Content;
    use blends_domain::ast::AstGraph;
    use blends_domain::NodeId;
    use std::fs;

    fn build_graph(source: &[u8]) -> AstGraph {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("snippet.java");
        fs::write(&path, source).unwrap();
        let content = Content::from_path(&path, None).unwrap();
        get_ast_graph(&content).unwrap()
    }

    #[test]
    fn root_node_is_program_at_one_one() {
        let graph = build_graph(b"class A {}");
        let root = graph.nodes.get(&NodeId(1)).unwrap();

        assert_eq!(root.kind, "program");
        assert_eq!(root.line, 1);
        assert_eq!(root.col, 1);
    }

    #[test]
    fn leaf_nodes_carry_their_text() {
        let graph = build_graph(b"class A {}");

        assert!(graph
            .nodes
            .values()
            .any(|node| node.text.as_deref() == Some("class")));
    }

    #[test]
    fn non_leaf_nodes_have_no_text() {
        let graph = build_graph(b"class A {}");
        let root = graph.nodes.get(&NodeId(1)).unwrap();

        assert!(root.text.is_none());
    }

    #[test]
    fn first_child_edge_is_indexed_zero() {
        let graph = build_graph(b"class A {}");
        let from_root = graph.edges.get(&NodeId(1)).unwrap();

        assert!(from_root.values().any(|edge| edge.index == 0));
    }

    #[test]
    fn empty_file_still_builds_a_root() {
        let graph = build_graph(b"");

        assert!(graph.nodes.contains_key(&NodeId(1)));
    }
}