use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use harn_lexer::{Lexer, Span, Token, TokenKind};
use harn_parser::{
lexical::{resolved_identifier_bindings, BindingId},
visit::walk_program_interpolated,
Node,
};
use super::{DefKind, DefSite, ModuleGraph, ParsedModuleSource};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RefSite {
pub file: PathBuf,
pub span: Span,
pub name: String,
}
#[derive(Debug, Clone, Default)]
pub struct ReferenceIndex {
pub files: Vec<PathBuf>,
pub has_unsaved_buffers: bool,
by_def: HashMap<DefKey, Vec<RefSite>>,
definitions: HashMap<DefKey, DefSite>,
site_defs: Vec<(RefSite, DefKey)>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct DefKey {
file: PathBuf,
name: String,
start: usize,
end: usize,
}
impl DefKey {
fn from_def(def: &DefSite) -> Self {
Self {
file: def.file.clone(),
name: def.name.clone(),
start: def.span.start,
end: def.span.end,
}
}
}
impl ReferenceIndex {
pub fn definition_at(
&self,
file: &std::path::Path,
name: &str,
offset: usize,
) -> Option<DefSite> {
let file = super::canonical_path(file);
self.site_defs
.iter()
.filter(|(site, _)| {
site.file == file
&& site.name == name
&& offset >= site.span.start
&& offset <= site.span.end
})
.min_by_key(|(site, _)| site.span.end.saturating_sub(site.span.start))
.and_then(|(_, key)| self.definitions.get(key))
.cloned()
}
pub fn references_to(&self, def: &DefSite) -> Vec<RefSite> {
self.by_def
.get(&DefKey::from_def(def))
.cloned()
.unwrap_or_default()
}
pub fn edges(&self) -> Vec<ReferenceEdge> {
let mut edges = Vec::new();
for (key, refs) in &self.by_def {
for site in refs {
edges.push(ReferenceEdge {
from: site.clone(),
to_file: key.file.clone(),
to_name: key.name.clone(),
});
}
}
edges.sort_by(|left, right| {
left.from
.file
.cmp(&right.from.file)
.then_with(|| left.from.span.start.cmp(&right.from.span.start))
.then_with(|| left.to_file.cmp(&right.to_file))
.then_with(|| left.to_name.cmp(&right.to_name))
});
edges
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReferenceEdge {
pub from: RefSite,
pub to_file: PathBuf,
pub to_name: String,
}
pub fn index_references(
graph: &ModuleGraph,
sources: &HashMap<PathBuf, ParsedModuleSource>,
unsaved: &HashSet<PathBuf>,
) -> ReferenceIndex {
let mut files: Vec<PathBuf> = sources.keys().cloned().collect();
files.sort();
let has_unsaved_buffers = files.iter().any(|file| unsaved.contains(file));
let mut index = ReferenceIndex {
files,
has_unsaved_buffers,
by_def: HashMap::new(),
definitions: HashMap::new(),
site_defs: Vec::new(),
};
for (file, parsed) in sources {
let lexical = resolved_identifier_bindings(&[], &parsed.program);
let tokens = Lexer::new(&parsed.source)
.tokenize()
.expect("a retained parsed source must still lex");
let mut lexical_defs: HashMap<BindingId, DefSite> = HashMap::new();
for binding in lexical.values() {
lexical_defs.entry(binding.clone()).or_insert_with(|| {
graph
.definition_of(file, &binding.name)
.filter(|def| {
def.file == *file
&& def.span.start == binding.declaration_start
&& def.span.end == binding.declaration_end
})
.unwrap_or_else(|| DefSite {
name: binding.name.clone(),
file: file.clone(),
kind: DefKind::Variable,
span: Span::with_offsets(
binding.declaration_start,
binding.declaration_end,
1,
1,
),
})
});
}
for (binding, def) in &lexical_defs {
if let Some(span) = identifier_span(
&tokens,
&binding.name,
binding.declaration_start,
binding.declaration_end,
) {
insert_reference(
&mut index,
def.clone(),
RefSite {
file: file.clone(),
span,
name: binding.name.clone(),
},
);
}
}
walk_program_interpolated(&parsed.source, &parsed.program, &mut |node| {
for (name, broad_span) in name_uses(node) {
let def = lexical
.get(&(broad_span.start, broad_span.end))
.and_then(|binding| lexical_defs.get(binding).cloned())
.or_else(|| graph.definition_of(file, name));
if let (Some(def), Some(span)) = (
def,
identifier_span(&tokens, name, broad_span.start, broad_span.end),
) {
insert_reference(
&mut index,
def,
RefSite {
file: file.clone(),
span,
name: name.to_string(),
},
);
}
}
});
}
for refs in index.by_def.values_mut() {
refs.sort_by(|left, right| {
left.file
.cmp(&right.file)
.then_with(|| left.span.start.cmp(&right.span.start))
});
refs.dedup_by(|left, right| {
left.file == right.file
&& left.span.start == right.span.start
&& left.name == right.name
});
}
index.site_defs.sort_by(|(left, _), (right, _)| {
left.file
.cmp(&right.file)
.then_with(|| left.span.start.cmp(&right.span.start))
.then_with(|| left.name.cmp(&right.name))
});
index
.site_defs
.dedup_by(|(left, left_key), (right, right_key)| {
left.file == right.file
&& left.span == right.span
&& left.name == right.name
&& left_key == right_key
});
index
}
fn insert_reference(index: &mut ReferenceIndex, def: DefSite, site: RefSite) {
let key = DefKey::from_def(&def);
index.definitions.entry(key.clone()).or_insert(def);
index
.by_def
.entry(key.clone())
.or_default()
.push(site.clone());
index.site_defs.push((site, key));
}
fn identifier_span(tokens: &[Token], name: &str, start: usize, end: usize) -> Option<Span> {
tokens.iter().find_map(|token| match &token.kind {
TokenKind::Identifier(found)
if found == name && token.span.start >= start && token.span.end <= end =>
{
Some(token.span)
}
_ => None,
})
}
fn name_uses(node: &harn_parser::SNode) -> Vec<(&str, Span)> {
match &node.node {
Node::Identifier(name) => vec![(name.as_str(), node.span)],
Node::FunctionCall { name, .. } => vec![(name.as_str(), node.span)],
Node::FnDecl { name, .. }
| Node::Pipeline { name, .. }
| Node::ToolDecl { name, .. }
| Node::SkillDecl { name, .. }
| Node::StructDecl { name, .. }
| Node::EnumDecl { name, .. }
| Node::InterfaceDecl { name, .. }
| Node::TypeDecl { name, .. }
| Node::OverrideDecl { name, .. } => vec![(name.as_str(), node.span)],
Node::EvalPackDecl { binding_name, .. } => vec![(binding_name.as_str(), node.span)],
_ => Vec::new(),
}
}