use std::collections::HashSet;
use tree_sitter::StreamingIterator;
use tree_sitter::{Node, Query, QueryCursor, Tree};
#[derive(Debug, Clone, Copy)]
pub struct QuerySpec {
pub name: &'static str,
pub capture: &'static str,
pub source: &'static str,
}
impl QuerySpec {
pub fn run<'tree>(
&self,
tree: &'tree Tree,
content: &[u8],
) -> Result<Vec<Node<'tree>>, String> {
let language = tree.language();
let query = Query::new(&language, self.source)
.map_err(|err| format!("failed to compile query `{}`: {err}", self.name))?;
let capture_index = query
.capture_names()
.iter()
.position(|name| *name == self.capture)
.ok_or_else(|| format!("query `{}` missing capture `{}`", self.name, self.capture))?;
let mut cursor = QueryCursor::new();
let mut matches = cursor.matches(&query, tree.root_node(), content);
let mut seen_offsets = HashSet::new();
let mut nodes = Vec::new();
while let Some(m) = matches.next() {
for capture in m.captures {
if capture.index as usize == capture_index
&& seen_offsets.insert(capture.node.start_byte())
{
nodes.push(capture.node);
}
}
}
Ok(nodes)
}
}
#[derive(Debug)]
pub struct ImportQuerySet {
pub statements: &'static [QuerySpec],
}
#[derive(Debug)]
pub struct ExportQuerySet {
pub statements: &'static [QuerySpec],
}
#[derive(Debug)]
pub struct CallQuerySet {
pub roots: &'static [QuerySpec],
}
pub mod typescript {
use super::{CallQuerySet, ExportQuerySet, ImportQuerySet, QuerySpec};
pub const IMPORT_QUERIES: ImportQuerySet = ImportQuerySet {
statements: &[QuerySpec {
name: "typescript_import_statements",
capture: "import.statement",
source: include_str!("typescript/imports.scm"),
}],
};
pub const EXPORT_QUERIES: ExportQuerySet = ExportQuerySet {
statements: &[QuerySpec {
name: "typescript_export_statements",
capture: "export.statement",
source: include_str!("typescript/exports.scm"),
}],
};
pub const CALL_QUERIES: CallQuerySet = CallQuerySet {
roots: &[QuerySpec {
name: "typescript_call_roots",
capture: "call.root",
source: include_str!("typescript/calls.scm"),
}],
};
}
pub mod javascript {
use super::{CallQuerySet, ExportQuerySet, ImportQuerySet, QuerySpec};
pub const IMPORT_QUERIES: ImportQuerySet = ImportQuerySet {
statements: &[QuerySpec {
name: "javascript_import_statements",
capture: "import.statement",
source: include_str!("javascript/imports.scm"),
}],
};
pub const EXPORT_QUERIES: ExportQuerySet = ExportQuerySet {
statements: &[QuerySpec {
name: "javascript_export_statements",
capture: "export.statement",
source: include_str!("javascript/exports.scm"),
}],
};
pub const CALL_QUERIES: CallQuerySet = CallQuerySet {
roots: &[QuerySpec {
name: "javascript_call_roots",
capture: "call.root",
source: include_str!("javascript/calls.scm"),
}],
};
}
pub mod csharp {
use super::{CallQuerySet, ImportQuerySet, QuerySpec};
pub const IMPORT_QUERIES: ImportQuerySet = ImportQuerySet {
statements: &[QuerySpec {
name: "csharp_using_directives",
capture: "import.directive",
source: include_str!("csharp/imports.scm"),
}],
};
pub const CALL_QUERIES: CallQuerySet = CallQuerySet {
roots: &[QuerySpec {
name: "csharp_call_roots",
capture: "call.root",
source: include_str!("csharp/calls.scm"),
}],
};
}
pub mod go {
use super::{CallQuerySet, ExportQuerySet, ImportQuerySet, QuerySpec};
pub const IMPORT_QUERIES: ImportQuerySet = ImportQuerySet {
statements: &[QuerySpec {
name: "go_import_declarations",
capture: "import.declaration",
source: include_str!("go/imports.scm"),
}],
};
pub const CALL_QUERIES: CallQuerySet = CallQuerySet {
roots: &[QuerySpec {
name: "go_call_roots",
capture: "call.root",
source: include_str!("go/calls.scm"),
}],
};
pub const EXPORT_QUERIES: ExportQuerySet = ExportQuerySet {
statements: &[QuerySpec {
name: "go_export_declarations",
capture: "export.declaration",
source: include_str!("go/exports.scm"),
}],
};
}