use crate::graph::GraphStore;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::time::Instant;
use tracing::{debug, info};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct FunctionData {
pub qualified_name: String,
pub file_path: String,
pub line_start: u32,
pub line_end: u32,
pub complexity: i32,
pub loc: i32,
pub parameters: Vec<String>,
pub return_type: Option<String>,
pub is_async: bool,
pub decorators: Vec<String>,
pub docstring: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ClassData {
pub qualified_name: String,
pub file_path: String,
pub line_start: u32,
pub line_end: u32,
pub complexity: i32,
pub method_count: i32,
pub methods: Vec<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct FileData {
pub file_path: String,
pub loc: i64,
pub language: String,
}
#[derive(Debug, Clone, Default)]
pub struct QueryCache {
pub functions: HashMap<String, FunctionData>,
pub classes: HashMap<String, ClassData>,
pub files: HashMap<String, FileData>,
pub calls: HashMap<String, HashSet<String>>,
pub imports: HashMap<String, HashSet<String>>,
pub inheritance: HashMap<String, HashSet<String>>,
pub functions_by_file: HashMap<String, Vec<String>>,
pub classes_by_file: HashMap<String, Vec<String>>,
pub callers: HashMap<String, HashSet<String>>,
}
impl QueryCache {
pub fn new() -> Self {
Self::default()
}
pub fn prefetch(&mut self, graph: &GraphStore) -> Result<()> {
let start = Instant::now();
info!("Prefetching graph data for detectors...");
self.prefetch_functions(graph)?;
self.prefetch_classes(graph)?;
self.prefetch_files(graph)?;
self.prefetch_calls(graph)?;
self.prefetch_imports(graph)?;
self.prefetch_inheritance(graph)?;
let elapsed = start.elapsed();
info!(
"Prefetched {} functions, {} classes, {} files in {:?}",
self.functions.len(),
self.classes.len(),
self.files.len(),
elapsed
);
Ok(())
}
fn prefetch_functions(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching functions...");
for func in graph.get_functions() {
let qn = func.qualified_name.clone();
self.functions.insert(
qn.clone(),
FunctionData {
qualified_name: qn.clone(),
file_path: func.file_path.clone(),
line_start: func.line_start,
line_end: func.line_end,
complexity: func.complexity().unwrap_or(1) as i32,
loc: func.loc() as i32,
parameters: vec![],
return_type: func.get_str("return_type").map(String::from),
is_async: func.get_bool("is_async").unwrap_or(false),
decorators: vec![],
docstring: func.get_str("docstring").map(String::from),
},
);
self.functions_by_file
.entry(func.file_path.clone())
.or_default()
.push(qn);
}
Ok(())
}
fn prefetch_classes(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching classes...");
for class in graph.get_classes() {
let qn = class.qualified_name.clone();
self.classes.insert(
qn.clone(),
ClassData {
qualified_name: qn.clone(),
file_path: class.file_path.clone(),
line_start: class.line_start,
line_end: class.line_end,
complexity: class.complexity().unwrap_or(1) as i32,
method_count: class.get_i64("methodCount").unwrap_or(0) as i32,
methods: vec![],
},
);
self.classes_by_file
.entry(class.file_path.clone())
.or_default()
.push(qn);
}
Ok(())
}
fn prefetch_files(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching files...");
for file in graph.get_files() {
self.files.insert(
file.file_path.clone(),
FileData {
file_path: file.file_path.clone(),
loc: file.get_i64("loc").unwrap_or(0),
language: file.language.clone().unwrap_or_default(),
},
);
}
Ok(())
}
fn prefetch_calls(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching call edges...");
for (caller, callee) in graph.get_calls() {
self.calls
.entry(caller.clone())
.or_default()
.insert(callee.clone());
self.callers
.entry(callee)
.or_default()
.insert(caller);
}
Ok(())
}
fn prefetch_imports(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching import edges...");
for (importer, imported) in graph.get_imports() {
self.imports
.entry(importer)
.or_default()
.insert(imported);
}
Ok(())
}
fn prefetch_inheritance(&mut self, graph: &GraphStore) -> Result<()> {
debug!("Prefetching inheritance...");
for (child, parent) in graph.get_inheritance() {
self.inheritance
.entry(child)
.or_default()
.insert(parent);
}
Ok(())
}
pub fn get_fan_in(&self, qn: &str) -> usize {
self.callers.get(qn).map(|s| s.len()).unwrap_or(0)
}
pub fn get_fan_out(&self, qn: &str) -> usize {
self.calls.get(qn).map(|s| s.len()).unwrap_or(0)
}
pub fn get_functions_in_file(&self, file_path: &str) -> Vec<&FunctionData> {
self.functions_by_file
.get(file_path)
.map(|qns| {
qns.iter()
.filter_map(|qn| self.functions.get(qn))
.collect()
})
.unwrap_or_default()
}
pub fn get_classes_in_file(&self, file_path: &str) -> Vec<&ClassData> {
self.classes_by_file
.get(file_path)
.map(|qns| {
qns.iter()
.filter_map(|qn| self.classes.get(qn))
.collect()
})
.unwrap_or_default()
}
}