use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;
use crate::query::adj_ast;
use crate::syntax::{SyntaxGraph, SyntaxKind, SyntaxNode};
use crate::{CodeGraph, NodeId};
#[must_use]
pub const fn nodes_by_type(graph: &SyntaxGraph) -> &BTreeMap<SyntaxKind, Vec<NodeId>> {
&graph.nodes_by_type
}
#[must_use]
pub fn get_nodes_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> BTreeSet<NodeId> {
let Some((first, rest)) = path.split_first() else {
return BTreeSet::new();
};
let matches = adj_ast(graph, n_id, Some(1), &[*first]);
if rest.is_empty() {
matches.into_iter().collect()
} else {
matches
.into_iter()
.flat_map(|child| get_nodes_by_path(graph, child, rest))
.collect()
}
}
#[must_use]
pub fn get_node_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> Option<NodeId> {
get_nodes_by_path(graph, n_id, path).into_iter().next()
}
#[must_use]
pub fn get_args_n_ids(graph: &SyntaxGraph, n_id: NodeId) -> Vec<NodeId> {
let Some(args_parent_n_id) = graph.nodes.get(&n_id).and_then(SyntaxNode::arguments_id) else {
return Vec::new();
};
adj_ast(graph, args_parent_n_id, Some(1), &[])
.into_iter()
.filter(|arg_id| {
graph
.nodes
.get(arg_id)
.is_some_and(|node| !matches!(node, SyntaxNode::Comment { .. }))
})
.collect()
}
#[must_use]
pub fn get_n_arg(graph: &SyntaxGraph, n_id: NodeId, arg_idx: usize) -> Option<NodeId> {
get_args_n_ids(graph, n_id).into_iter().nth(arg_idx)
}
#[cfg(test)]
mod tests {
use super::{get_args_n_ids, get_n_arg, get_node_by_path, get_nodes_by_path, nodes_by_type};
use crate::ast::{AstGraph, AstNode};
use crate::syntax::{SyntaxGraph, SyntaxKind, SyntaxNode};
use crate::NodeId;
use alloc::borrow::ToOwned;
use alloc::collections::BTreeSet;
use alloc::vec;
use alloc::vec::Vec;
fn sample() -> AstGraph {
let mut ast = AstGraph::new();
ast.add_node(
NodeId(1),
AstNode::new(1, 1, "class_declaration".to_owned()),
);
ast.add_node(NodeId(2), AstNode::new(1, 1, "base_list".to_owned()));
ast.add_node(NodeId(3), AstNode::new(1, 1, "identifier".to_owned()));
ast.add_node(NodeId(4), AstNode::new(1, 1, "modifier".to_owned()));
ast.add_node(NodeId(5), AstNode::new(1, 2, "identifier".to_owned()));
ast.add_edge(NodeId(1), NodeId(2), 0);
ast.add_edge(NodeId(1), NodeId(4), 1);
ast.add_edge(NodeId(2), NodeId(3), 0);
ast.add_edge(NodeId(2), NodeId(5), 1);
ast
}
#[test]
fn nodes_by_type_exposes_the_index_maintained_by_add_node() {
let mut graph = SyntaxGraph::new();
graph.add_node(NodeId(9), SyntaxNode::Break);
graph.add_node(NodeId(1), SyntaxNode::File);
graph.add_node(NodeId(2), SyntaxNode::Break);
assert_eq!(
nodes_by_type(&graph).get(&SyntaxKind::Break),
Some(&vec![NodeId(9), NodeId(2)])
);
assert_eq!(
nodes_by_type(&graph).get(&SyntaxKind::File),
Some(&vec![NodeId(1)])
);
assert_eq!(nodes_by_type(&graph).get(&SyntaxKind::If), None);
assert!(nodes_by_type(&SyntaxGraph::new()).is_empty());
}
#[test]
fn nodes_by_path_collect_every_match_of_the_last_step() {
let ast = sample();
assert_eq!(
get_nodes_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
BTreeSet::from([NodeId(3), NodeId(5)])
);
assert_eq!(
get_nodes_by_path(&ast, NodeId(1), &["base_list"]),
BTreeSet::from([NodeId(2)])
);
}
#[test]
fn nodes_by_path_are_empty_when_the_path_breaks_or_is_empty() {
let ast = sample();
assert_eq!(
get_nodes_by_path(&ast, NodeId(1), &["base_list", "block"]),
BTreeSet::new()
);
assert_eq!(get_nodes_by_path(&ast, NodeId(1), &[]), BTreeSet::new());
}
#[test]
fn follows_a_multi_step_label_path() {
let ast = sample();
assert_eq!(
get_node_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
Some(NodeId(3))
);
}
#[test]
fn returns_the_direct_child_for_a_single_step_path() {
let ast = sample();
assert_eq!(
get_node_by_path(&ast, NodeId(1), &["base_list"]),
Some(NodeId(2))
);
}
#[test]
fn is_none_when_the_path_breaks_or_is_empty() {
let ast = sample();
assert_eq!(
get_node_by_path(&ast, NodeId(1), &["base_list", "block"]),
None
);
assert_eq!(get_node_by_path(&ast, NodeId(1), &["missing"]), None);
assert_eq!(get_node_by_path(&ast, NodeId(1), &[]), None);
}
fn invocation(arguments_id: Option<NodeId>) -> SyntaxNode {
SyntaxNode::MethodInvocation {
expression: "run".to_owned(),
object: None,
symbol_scope: None,
expression_id: None,
arguments_id,
object_id: None,
block_id: None,
receiver_type_fqn: None,
}
}
fn literal(value: &str) -> SyntaxNode {
SyntaxNode::Literal {
value: value.to_owned(),
value_type: "string".to_owned(),
}
}
fn call_with_commented_args() -> SyntaxGraph {
let mut graph = SyntaxGraph::new();
graph.add_node(NodeId(1), invocation(Some(NodeId(2))));
graph.add_node(NodeId(2), SyntaxNode::ArgumentList);
graph.add_node(NodeId(3), literal("first"));
graph.add_node(
NodeId(4),
SyntaxNode::Comment {
comment: "// second arg".to_owned(),
},
);
graph.add_node(NodeId(5), literal("second"));
graph.add_ast_edge(NodeId(1), NodeId(2));
graph.add_ast_edge(NodeId(2), NodeId(3));
graph.add_ast_edge(NodeId(2), NodeId(4));
graph.add_ast_edge(NodeId(2), NodeId(5));
graph
}
#[test]
fn args_skip_comments_interleaved_in_the_argument_list() {
let graph = call_with_commented_args();
assert_eq!(get_args_n_ids(&graph, NodeId(1)), [NodeId(3), NodeId(5)]);
}
#[test]
fn args_are_empty_without_an_argument_list() {
let mut graph = SyntaxGraph::new();
graph.add_node(NodeId(1), invocation(None));
graph.add_node(NodeId(2), invocation(Some(NodeId(9))));
graph.add_node(NodeId(3), literal("plain"));
assert_eq!(get_args_n_ids(&graph, NodeId(1)), Vec::new());
assert_eq!(get_args_n_ids(&graph, NodeId(2)), Vec::new());
assert_eq!(get_args_n_ids(&graph, NodeId(3)), Vec::new());
assert_eq!(get_args_n_ids(&graph, NodeId(404)), Vec::new());
}
#[test]
fn nth_arg_indexes_the_comment_free_positions() {
let graph = call_with_commented_args();
assert_eq!(get_n_arg(&graph, NodeId(1), 0), Some(NodeId(3)));
assert_eq!(get_n_arg(&graph, NodeId(1), 1), Some(NodeId(5)));
assert_eq!(get_n_arg(&graph, NodeId(1), 2), None);
}
#[test]
fn args_ignore_edges_pointing_at_unknown_nodes() {
let mut graph = call_with_commented_args();
graph.add_ast_edge(NodeId(2), NodeId(0));
assert_eq!(get_args_n_ids(&graph, NodeId(1)), [NodeId(3), NodeId(5)]);
assert_eq!(get_n_arg(&graph, NodeId(1), 0), Some(NodeId(3)));
assert_eq!(get_n_arg(&graph, NodeId(1), 1), Some(NodeId(5)));
}
}