use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use crate::query::traversal::{adj_ast, adj_cfg, pred_cfg, pred_lazy, EdgeKind};
use crate::{CodeGraph, NodeId};
#[must_use]
pub fn is_connected_to_cfg<G: CodeGraph>(graph: &G, n_id: NodeId) -> bool {
!pred_cfg(graph, n_id, None).is_empty() || !adj_cfg(graph, n_id, None, &[]).is_empty()
}
#[must_use]
pub fn lookup_first_cfg_parent<G: CodeGraph>(graph: &G, n_id: NodeId) -> NodeId {
core::iter::once(n_id)
.chain(pred_lazy(graph, n_id, Some(-1), EdgeKind::Ast))
.find(|p_id| is_connected_to_cfg(graph, *p_id))
.unwrap_or(n_id)
}
#[must_use]
pub fn match_ast<G: CodeGraph>(
graph: &G,
n_id: NodeId,
label_types: &[&str],
depth: Option<i64>,
) -> BTreeMap<String, Option<NodeId>> {
let mut index: usize = 0;
let mut nodes: BTreeMap<String, Option<NodeId>> = label_types
.iter()
.map(|label| (label.to_string(), None))
.collect();
for c_id in adj_ast(graph, n_id, depth, &[]) {
let c_type = graph.label_type(c_id).unwrap_or_default();
let filled = match nodes.get_mut(c_type) {
Some(slot) if slot.is_none() => {
*slot = Some(c_id);
true
}
_ => false,
};
if !filled {
nodes.insert(format!("__{index}__"), Some(c_id));
index = index.saturating_add(1);
}
}
nodes
}
#[must_use]
pub fn match_ast_d<G: CodeGraph>(
graph: &G,
n_id: NodeId,
node_type: &str,
depth: Option<i64>,
) -> Option<NodeId> {
match_ast(graph, n_id, &[node_type], depth)
.get(node_type)
.copied()
.flatten()
}
#[must_use]
pub fn match_ast_group<G: CodeGraph>(
graph: &G,
n_id: NodeId,
label_types: &[&str],
depth: Option<i64>,
) -> BTreeMap<String, Vec<NodeId>> {
let mut index: usize = 0;
let mut nodes: BTreeMap<String, Vec<NodeId>> = label_types
.iter()
.map(|label| (label.to_string(), Vec::new()))
.collect();
for c_id in adj_ast(graph, n_id, depth, &[]) {
let c_type = graph.label_type(c_id).unwrap_or_default();
if let Some(group) = nodes.get_mut(c_type) {
group.push(c_id);
} else {
nodes.insert(format!("__{index}__"), vec![c_id]);
index = index.saturating_add(1);
}
}
nodes
}
#[must_use]
pub fn match_ast_group_d<G: CodeGraph>(
graph: &G,
n_id: NodeId,
node_type: &str,
depth: Option<i64>,
) -> Vec<NodeId> {
match_ast_group(graph, n_id, &[node_type], depth)
.remove(node_type)
.unwrap_or_default()
}
#[cfg(test)]
mod cfg_parent_tests {
use super::{is_connected_to_cfg, lookup_first_cfg_parent};
use crate::syntax::{SyntaxGraph, SyntaxNode};
use crate::NodeId;
#[test]
fn climbs_ast_parents_to_the_first_cfg_connected_node() {
let mut graph = SyntaxGraph::new();
graph.add_node(NodeId(1), SyntaxNode::ExecutionBlock);
graph.add_node(NodeId(2), SyntaxNode::ExecutionBlock);
graph.add_node(
NodeId(3),
SyntaxNode::SymbolLookup {
symbol: alloc::string::String::from("x"),
symbol_scope: None,
value: None,
},
);
graph.add_cfg_edge(NodeId(1), NodeId(2));
graph.add_ast_edge(NodeId(2), NodeId(3));
assert!(!is_connected_to_cfg(&graph, NodeId(3)));
assert!(is_connected_to_cfg(&graph, NodeId(2)));
assert_eq!(lookup_first_cfg_parent(&graph, NodeId(3)), NodeId(2));
assert_eq!(lookup_first_cfg_parent(&graph, NodeId(2)), NodeId(2));
}
#[test]
fn falls_back_to_the_node_itself_without_any_cfg_ancestor() {
let mut graph = SyntaxGraph::new();
graph.add_node(NodeId(1), SyntaxNode::ExecutionBlock);
graph.add_node(NodeId(2), SyntaxNode::ExecutionBlock);
graph.add_ast_edge(NodeId(1), NodeId(2));
assert_eq!(lookup_first_cfg_parent(&graph, NodeId(2)), NodeId(2));
}
}
#[cfg(test)]
mod tests {
use super::{match_ast, match_ast_d, match_ast_group, match_ast_group_d};
use crate::ast::{AstGraph, AstNode};
use crate::NodeId;
use alloc::borrow::ToOwned;
use alloc::vec;
fn sample_ast() -> AstGraph {
let mut ast = AstGraph::new();
ast.add_node(NodeId(1), AstNode::new(1, 1, "stream".to_owned()));
ast.add_node(NodeId(2), AstNode::new(1, 1, "document".to_owned()));
ast.add_node(NodeId(3), AstNode::new(2, 1, "document".to_owned()));
ast.add_node(NodeId(4), AstNode::new(3, 1, "block_node".to_owned()));
ast.add_edge(NodeId(1), NodeId(2), 0);
ast.add_edge(NodeId(1), NodeId(3), 1);
ast.add_edge(NodeId(1), NodeId(4), 2);
ast
}
#[test]
fn match_ast_fills_first_occurrence_and_indexes_the_rest() {
let ast = sample_ast();
let matched = match_ast(&ast, NodeId(1), &["document"], None);
assert_eq!(matched.get("document"), Some(&Some(NodeId(2))));
assert_eq!(matched.get("__0__"), Some(&Some(NodeId(3))));
assert_eq!(matched.get("__1__"), Some(&Some(NodeId(4))));
}
#[test]
fn match_ast_keeps_unmatched_labels_as_none() {
let ast = sample_ast();
let matched = match_ast(&ast, NodeId(1), &["mapping"], None);
assert_eq!(matched.get("mapping"), Some(&None));
}
#[test]
fn match_ast_d_returns_first_child_of_kind() {
let ast = sample_ast();
assert_eq!(
match_ast_d(&ast, NodeId(1), "document", None),
Some(NodeId(2))
);
assert_eq!(match_ast_d(&ast, NodeId(1), "mapping", None), None);
}
#[test]
fn match_ast_group_collects_by_kind_and_indexes_the_rest() {
let ast = sample_ast();
let groups = match_ast_group(&ast, NodeId(1), &["document"], None);
assert_eq!(groups.get("document"), Some(&vec![NodeId(2), NodeId(3)]));
assert_eq!(groups.get("__0__"), Some(&vec![NodeId(4)]));
}
#[test]
fn match_ast_group_d_returns_all_children_of_kind() {
let ast = sample_ast();
assert_eq!(
match_ast_group_d(&ast, NodeId(1), "document", None),
[NodeId(2), NodeId(3)]
);
assert_eq!(match_ast_group_d(&ast, NodeId(1), "mapping", None), []);
}
}