use crate::error::Result;
use crate::graph::{CodeGraph, Direction, Node, NodeId, NodeType, PropertyValue};
type FilterFn = Box<dyn Fn(&Node) -> bool>;
pub struct QueryBuilder<'a> {
graph: &'a CodeGraph,
filters: Vec<FilterFn>,
limit_value: Option<usize>,
in_file_filter: Option<String>,
}
impl<'a> QueryBuilder<'a> {
pub fn new(graph: &'a CodeGraph) -> Self {
Self {
graph,
filters: Vec::new(),
limit_value: None,
in_file_filter: None,
}
}
pub fn node_type(mut self, node_type: NodeType) -> Self {
self.filters
.push(Box::new(move |node| node.node_type == node_type));
self
}
pub fn in_file(mut self, file_path: &str) -> Self {
self.in_file_filter = Some(file_path.to_string());
self
}
pub fn file_pattern(mut self, pattern: &str) -> Self {
let pattern = pattern.to_string();
self.filters.push(Box::new(move |node| {
if let Some(path) = node.properties.get_string("path") {
glob_match(&pattern, path)
} else {
false
}
}));
self
}
pub fn property<V: Into<PropertyValue>>(mut self, key: &str, value: V) -> Self {
let key = key.to_string();
let value = value.into();
self.filters.push(Box::new(move |node| {
if let Some(prop_value) = node.properties.get(&key) {
match (&value, prop_value) {
(PropertyValue::String(v1), PropertyValue::String(v2)) => v1 == v2,
(PropertyValue::Int(v1), PropertyValue::Int(v2)) => v1 == v2,
(PropertyValue::Float(v1), PropertyValue::Float(v2)) => {
(v1 - v2).abs() < f64::EPSILON
}
(PropertyValue::Bool(v1), PropertyValue::Bool(v2)) => v1 == v2,
_ => false,
}
} else {
false
}
}));
self
}
pub fn property_exists(mut self, key: &str) -> Self {
let key = key.to_string();
self.filters
.push(Box::new(move |node| node.properties.contains_key(&key)));
self
}
pub fn name_contains(mut self, substring: &str) -> Self {
let substring = substring.to_lowercase();
self.filters.push(Box::new(move |node| {
if let Some(name) = node.properties.get_string("name") {
name.to_lowercase().contains(&substring)
} else {
false
}
}));
self
}
pub fn name_matches(mut self, pattern: &str) -> Self {
let pattern = pattern.to_string();
self.filters.push(Box::new(move |node| {
if let Some(name) = node.properties.get_string("name") {
regex_match(&pattern, name)
} else {
false
}
}));
self
}
pub fn custom<F>(mut self, predicate: F) -> Self
where
F: Fn(&Node) -> bool + 'static,
{
self.filters.push(Box::new(predicate));
self
}
pub fn limit(mut self, n: usize) -> Self {
self.limit_value = Some(n);
self
}
pub fn execute(&self) -> Result<Vec<NodeId>> {
let mut results = Vec::new();
let limit = self.limit_value.unwrap_or(usize::MAX);
let search_nodes: Vec<NodeId> = if let Some(file_path) = &self.in_file_filter {
self.get_nodes_in_file(file_path)?
} else {
(0..self.graph.node_count() as u64).collect()
};
for node_id in search_nodes {
if results.len() >= limit {
break;
}
if let Ok(node) = self.graph.get_node(node_id) {
if self.matches_filters(node) {
results.push(node_id);
}
}
}
Ok(results)
}
pub fn count(&self) -> Result<usize> {
let mut count = 0;
let search_nodes: Vec<NodeId> = if let Some(file_path) = &self.in_file_filter {
self.get_nodes_in_file(file_path)?
} else {
(0..self.graph.node_count() as u64).collect()
};
for node_id in search_nodes {
if let Ok(node) = self.graph.get_node(node_id) {
if self.matches_filters(node) {
count += 1;
}
}
}
Ok(count)
}
pub fn exists(&self) -> Result<bool> {
let search_nodes: Vec<NodeId> = if let Some(file_path) = &self.in_file_filter {
self.get_nodes_in_file(file_path)?
} else {
(0..self.graph.node_count() as u64).collect()
};
for node_id in search_nodes {
if let Ok(node) = self.graph.get_node(node_id) {
if self.matches_filters(node) {
return Ok(true);
}
}
}
Ok(false)
}
fn get_nodes_in_file(&self, file_path: &str) -> Result<Vec<NodeId>> {
for node_id in 0..self.graph.node_count() as u64 {
if let Ok(node) = self.graph.get_node(node_id) {
if node.node_type == NodeType::CodeFile {
if let Some(path) = node.properties.get_string("path") {
if path == file_path {
return self.graph.get_neighbors(node_id, Direction::Outgoing);
}
}
}
}
}
Ok(Vec::new())
}
fn matches_filters(&self, node: &Node) -> bool {
self.filters.iter().all(|filter| filter(node))
}
}
fn glob_match(pattern: &str, path: &str) -> bool {
if pattern.contains("**") {
let parts: Vec<&str> = pattern.split("**").collect();
if parts.len() == 2 {
let prefix = parts[0];
let suffix = parts[1].trim_start_matches('/');
if !prefix.is_empty() && !path.starts_with(prefix) {
return false;
}
if suffix.contains('*') {
if let Some(last_slash) = path.rfind('/') {
let filename = &path[last_slash + 1..];
return glob_match(suffix, filename);
} else {
return glob_match(suffix, path);
}
}
if !suffix.is_empty() && !path.ends_with(suffix) {
return false;
}
return true;
}
}
let pattern_parts: Vec<&str> = pattern.split('*').collect();
if pattern_parts.len() == 1 {
return pattern == path;
}
let mut pos = 0;
for (i, part) in pattern_parts.iter().enumerate() {
if part.is_empty() {
continue;
}
if i == 0 {
if !path[pos..].starts_with(part) {
return false;
}
pos += part.len();
} else if i == pattern_parts.len() - 1 {
return path[pos..].ends_with(part);
} else {
if let Some(index) = path[pos..].find(part) {
pos += index + part.len();
} else {
return false;
}
}
}
true
}
fn regex_match(pattern: &str, text: &str) -> bool {
let starts_with = pattern.starts_with('^');
let ends_with = pattern.ends_with('$');
let pattern = pattern.trim_start_matches('^').trim_end_matches('$');
if starts_with && ends_with {
text == pattern
} else if starts_with {
text.starts_with(pattern)
} else if ends_with {
text.ends_with(pattern)
} else {
text.contains(pattern)
}
}