use crate::tag::{Tag, TagKind};
use crate::weight::compute_edge_weight;
use petgraph::graph::DiGraph;
use std::collections::{HashMap, HashSet};
#[derive(Debug, Default)]
pub struct TagIndex {
pub defines: HashMap<String, HashSet<String>>,
pub references: HashMap<String, Vec<String>>,
pub definitions: HashMap<(String, String), HashSet<Tag>>,
}
impl TagIndex {
pub fn from_tags(tags: impl Iterator<Item = Tag>) -> Self {
let mut index = TagIndex::default();
for tag in tags {
match tag.kind {
TagKind::Def => {
index
.defines
.entry(tag.name.clone())
.or_default()
.insert(tag.rel_fname.clone());
index
.definitions
.entry((tag.rel_fname.clone(), tag.name.clone()))
.or_default()
.insert(tag);
}
TagKind::Ref => {
index
.references
.entry(tag.name.clone())
.or_default()
.push(tag.rel_fname.clone());
}
}
}
index
}
pub fn apply_no_reference_fallback(&mut self) {
if self.references.is_empty() {
for (name, fnames) in &self.defines {
self.references
.insert(name.clone(), fnames.iter().cloned().collect());
}
}
}
}
#[derive(Debug, Clone)]
pub struct Edge {
pub ident: String,
pub weight: f64,
}
pub type Graph = DiGraph<String, Edge>;
pub fn build_graph(
index: &TagIndex,
mentioned_idents: &HashSet<String>,
chat_rel_fnames: &HashSet<String>,
self_edge_weight: f64,
) -> Graph {
let mut graph = Graph::new();
let mut node_indices: HashMap<String, petgraph::graph::NodeIndex> = HashMap::new();
let mut get_node = |graph: &mut Graph, fname: &str| -> petgraph::graph::NodeIndex {
if let Some(&idx) = node_indices.get(fname) {
idx
} else {
let idx = graph.add_node(fname.to_string());
node_indices.insert(fname.to_string(), idx);
idx
}
};
let idents_with_refs: HashSet<&String> = index.references.keys().collect();
for ident in index.defines.keys() {
if !index.references.contains_key(ident) {
continue;
}
let definers = &index.defines[ident];
let num_definers = definers.len();
let mut ref_counts: HashMap<&str, usize> = HashMap::new();
for referencer in &index.references[ident] {
*ref_counts.entry(referencer.as_str()).or_default() += 1;
}
for (referencer, num_refs) in ref_counts {
for definer in definers {
let weight = compute_edge_weight(
ident,
referencer,
num_refs,
mentioned_idents,
chat_rel_fnames,
num_definers,
);
let src = get_node(&mut graph, referencer);
let dst = get_node(&mut graph, definer);
graph.add_edge(
src,
dst,
Edge {
ident: ident.clone(),
weight,
},
);
}
}
}
for (ident, definers) in &index.defines {
if idents_with_refs.contains(ident) {
continue;
}
for definer in definers {
let node = get_node(&mut graph, definer);
graph.add_edge(
node,
node,
Edge {
ident: ident.clone(),
weight: self_edge_weight,
},
);
}
}
graph
}
pub fn total_out_weight(graph: &Graph, node: petgraph::graph::NodeIndex) -> f64 {
graph.edges(node).map(|e| e.weight().weight).sum()
}
pub fn graph_nodes(graph: &Graph) -> Vec<&str> {
graph.node_weights().map(|s| s.as_str()).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tag::Tag;
#[test]
fn tag_index_from_tags() {
let tags = vec![
Tag::def("a.rs", "/a.rs", 1, "Foo"),
Tag::def("a.rs", "/a.rs", 5, "Bar"),
Tag::reference("b.rs", "/b.rs", 10, "Foo"),
Tag::reference("b.rs", "/b.rs", 15, "Foo"),
];
let index = TagIndex::from_tags(tags.into_iter());
assert!(index.defines["Foo"].contains("a.rs"));
assert!(index.defines["Bar"].contains("a.rs"));
assert_eq!(index.references["Foo"].len(), 2); assert!(!index.references.contains_key("Bar")); }
#[test]
fn no_reference_fallback() {
let tags = vec![
Tag::def("a.rs", "/a.rs", 1, "Foo"),
Tag::def("b.rs", "/b.rs", 1, "Bar"),
];
let mut index = TagIndex::from_tags(tags.into_iter());
assert!(index.references.is_empty());
index.apply_no_reference_fallback();
assert!(index.references.contains_key("Foo"));
assert!(index.references.contains_key("Bar"));
}
#[test]
fn build_graph_basic() {
let tags = vec![
Tag::def("a.rs", "/a.rs", 1, "Foo"),
Tag::reference("b.rs", "/b.rs", 10, "Foo"),
];
let index = TagIndex::from_tags(tags.into_iter());
let mentioned = HashSet::new();
let chat_files = HashSet::new();
let graph = build_graph(&index, &mentioned, &chat_files, 0.1);
assert_eq!(graph.node_count(), 2);
assert_eq!(graph.edge_count(), 1);
}
#[test]
fn build_graph_self_edges() {
let tags = vec![
Tag::def("a.rs", "/a.rs", 1, "Foo"),
];
let mut index = TagIndex::from_tags(tags.into_iter());
let _ = &mut index; let mentioned = HashSet::new();
let chat_files = HashSet::new();
let graph = build_graph(&index, &mentioned, &chat_files, 0.5);
assert_eq!(graph.node_count(), 1);
assert_eq!(graph.edge_count(), 1);
let edge = graph.edge_weights().next().unwrap();
assert!((edge.weight - 0.5).abs() < 0.001);
}
#[test]
fn total_out_weight_calculation() {
let tags = vec![
Tag::def("a.rs", "/a.rs", 1, "Foo"),
Tag::def("a.rs", "/a.rs", 2, "Bar"),
Tag::reference("b.rs", "/b.rs", 10, "Foo"),
Tag::reference("b.rs", "/b.rs", 20, "Bar"),
];
let index = TagIndex::from_tags(tags.into_iter());
let graph = build_graph(&index, &HashSet::new(), &HashSet::new(), 0.1);
let b_node = graph.node_indices().find(|&n| graph[n] == "b.rs").unwrap();
let total = total_out_weight(&graph, b_node);
assert!((total - 2.0).abs() < 0.001);
}
}