#![feature(portable_simd)]
pub mod ast;
pub mod callgraph;
pub mod cfg;
pub mod dfg;
pub mod embedding;
pub mod error;
pub mod metrics;
pub mod patterns;
pub mod pdg;
pub mod quality;
pub mod security;
pub mod semantic;
pub mod simd;
pub mod util;
pub mod lang {
pub mod c;
pub mod cpp;
pub mod go;
pub mod java;
pub mod python;
pub mod registry;
pub mod rust_lang;
pub mod traits;
pub mod typescript;
pub use registry::LanguageRegistry;
pub use traits::{BoxedLanguage, Language};
}
pub use error::{Result, BrrrError};
pub use ast::{
CallGraphInfo, ClassInfo, ClassSummary, CodeStructure, FieldInfo, FileTreeEntry, FunctionInfo,
FunctionSummary, ImportInfo, ModuleInfo,
};
pub use ast::AstExtractor;
pub use ast::{clear_parser_cache, clear_query_cache};
pub use ast::extract_imports;
pub use cfg::{BlockId, BlockType, CFGBlock, CFGEdge, CFGError, CFGInfo, EdgeType};
pub use cfg::{
to_ascii as cfg_to_ascii, to_dot as cfg_to_dot, to_json as cfg_to_json,
to_json_compact as cfg_to_json_compact, to_mermaid as cfg_to_mermaid,
};
pub use dfg::{DFGInfo, DataflowEdge, DataflowKind};
pub use pdg::{
BranchType, ControlDependence, PDGInfo, SliceCriteria, SliceMetrics, SliceResult,
};
pub use pdg::{backward_slice as pdg_backward_slice, forward_slice as pdg_forward_slice};
pub use metrics::{
analyze_complexity, analyze_file_complexity, ComplexityAnalysis, ComplexityStats,
CyclomaticComplexity, FunctionComplexity, RiskLevel,
};
pub use metrics::{
analyze_nesting, analyze_file_nesting, NestingMetrics, NestingAnalysis,
NestingStats, FunctionNesting, NestingDepthLevel, NestingConstruct,
DeepNesting, NestingAnalysisError,
};
pub use callgraph::{CallEdge, CallGraph, FunctionRef};
pub use callgraph::{FunctionDef, FunctionIndex, IndexStats};
pub use callgraph::scanner::{
ErrorHandling, FileMetadata, ProjectScanner, ScanConfig, ScanError, ScanErrorKind, ScanResult,
};
pub use callgraph::{
analyze_dead_code, analyze_dead_code_with_config, DeadCodeConfig, DeadCodeResult,
DeadCodeStats, DeadFunction, DeadReason,
};
pub use callgraph::{classify_entry_point, detect_entry_points_with_config, EntryPointKind};
pub use callgraph::{analyze_impact, CallerInfo, ImpactConfig, ImpactResult};
pub use callgraph::{analyze_architecture, ArchAnalysis, ArchStats, CycleDependency};
pub use callgraph::{
get_cache_dir, get_cache_file, get_or_build_graph_with_config,
invalidate_cache, warm_cache_with_config, CachedCallGraph, CachedEdge,
};
pub use lang::{BoxedLanguage, Language, LanguageRegistry};
pub use semantic::{
ChunkInfo, CodeComplexity, CodeLocation, ContentHashedIndex, EmbeddingUnit, SearchResult,
SemanticPattern, UnitKind, CHUNK_OVERLAP_TOKENS, MAX_CODE_PREVIEW_TOKENS, MAX_EMBEDDING_TOKENS,
SEMANTIC_PATTERNS,
};
pub use embedding::{
distances_to_scores, distances_to_scores_for_metric, is_normalized, normalize_vector,
IndexConfig, Metric, Quantization, VectorIndex,
};
pub use security::injection::command::{
CommandInjectionFinding, CommandSink, Confidence, InjectionKind,
Severity as CommandSeverity, SourceLocation, TaintSource, TaintSourceKind,
scan_command_injection, scan_file_command_injection,
};
pub use security::injection::sql::{
Location as SqlLocation, SQLInjectionFinding, ScanResult as SqlScanResult,
Severity as SqlSeverity, SqlInjectionDetector, SqlSinkType, UnsafePattern,
};
pub use security::injection::path_traversal::{
Confidence as PathTraversalConfidence, FileOperationType, FileSink,
PathTraversalFinding, ScanResult as PathTraversalScanResult,
Severity as PathTraversalSeverity, SourceLocation as PathTraversalLocation,
VulnerablePattern as PathTraversalPattern,
scan_path_traversal, scan_file_path_traversal, get_file_sinks,
};
pub use security::crypto::{
Algorithm as CryptoAlgorithm, Confidence as CryptoConfidence,
Location as CryptoLocation, ScanResult as CryptoScanResult,
Severity as CryptoSeverity, UsageContext as CryptoUsageContext,
WeakCryptoDetector, WeakCryptoFinding, WeakCryptoIssue,
scan_weak_crypto, scan_file_weak_crypto,
};
pub use security::{
scan_security, Confidence as UnifiedConfidence, InjectionType,
Location as UnifiedLocation, ScanSummary, SecurityCategory, SecurityConfig,
SecurityFinding, SecurityReport, Severity as UnifiedSeverity,
check_suppression, is_suppressed,
};
pub use security::sarif::SarifLog;
pub use quality::clones::{
detect_clones, format_clone_summary, Clone, CloneAnalysis, CloneConfig, CloneError,
CloneInstance, CloneStats, CloneType, TextualCloneDetector,
};
pub use patterns::{
detect_patterns, format_pattern_summary, DesignPattern, Location as PatternLocation,
PatternAnalysis, PatternCategory, PatternConfig, PatternDetector, PatternError,
PatternMatch, PatternStats,
};
use serde::{Deserialize, Serialize};
use std::path::Path;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionContext {
pub name: String,
pub file: String,
pub start_line: usize,
pub end_line: usize,
pub signature: Option<String>,
pub docstring: Option<String>,
pub source: String,
pub language: String,
}
impl FunctionContext {
pub fn from_function_info(
info: &FunctionInfo,
file_path: &str,
source: &str,
language: &str,
) -> Self {
let start = info.line_number.saturating_sub(1);
let end = info.end_line_number.unwrap_or(info.line_number);
let lines: Vec<&str> = source.lines().collect();
let func_source = lines
.get(start..end)
.map(|ls| ls.join("\n"))
.unwrap_or_default();
Self {
name: info.name.clone(),
file: file_path.to_string(),
start_line: info.line_number,
end_line: end,
signature: Some(info.signature()),
docstring: info.docstring.clone(),
source: func_source,
language: language.to_string(),
}
}
pub fn minimal(name: &str, file: &str, line: usize, language: &str) -> Self {
Self {
name: name.to_string(),
file: file.to_string(),
start_line: line,
end_line: line,
signature: None,
docstring: None,
source: String::new(),
language: language.to_string(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelevantContext {
pub entry: FunctionContext,
pub callees: Vec<FunctionContext>,
pub callers: Vec<FunctionContext>,
pub token_count: usize,
pub depth: usize,
}
impl RelevantContext {
pub fn new(entry: FunctionContext, depth: usize) -> Self {
Self {
entry,
callees: Vec::new(),
callers: Vec::new(),
token_count: 0,
depth,
}
}
pub fn with_callee(mut self, callee: FunctionContext) -> Self {
self.callees.push(callee);
self
}
pub fn with_caller(mut self, caller: FunctionContext) -> Self {
self.callers.push(caller);
self
}
pub fn with_token_count(mut self, count: usize) -> Self {
self.token_count = count;
self
}
pub fn estimate_tokens(&mut self) {
let mut total_chars = self.entry.source.len();
for callee in &self.callees {
total_chars += callee.source.len();
}
for caller in &self.callers {
total_chars += caller.source.len();
}
self.token_count = total_chars / 4;
}
pub fn function_count(&self) -> usize {
1 + self.callees.len() + self.callers.len()
}
pub fn to_llm_string(&self) -> String {
let mut output = String::new();
output.push_str(&format!(
"## Code Context: {} (depth={})\n\n",
self.entry.name, self.depth
));
let short_file = Path::new(&self.entry.file)
.file_name()
.map(|f| f.to_string_lossy().to_string())
.unwrap_or_else(|| self.entry.file.clone());
output.push_str(&format!(
"### Entry Point: {} ({}:{})\n",
self.entry.name, short_file, self.entry.start_line
));
if let Some(sig) = &self.entry.signature {
output.push_str(&format!("```\n{}\n```\n", sig));
}
if let Some(doc) = &self.entry.docstring {
let first_line = doc.lines().next().unwrap_or("");
output.push_str(&format!("> {}\n", first_line));
}
output.push('\n');
if !self.callees.is_empty() {
output.push_str("### Calls:\n");
for callee in &self.callees {
let short = Path::new(&callee.file)
.file_name()
.map(|f| f.to_string_lossy().to_string())
.unwrap_or_else(|| callee.file.clone());
output.push_str(&format!("- {} ({}:{})\n", callee.name, short, callee.start_line));
}
output.push('\n');
}
if !self.callers.is_empty() {
output.push_str("### Called By:\n");
for caller in &self.callers {
let short = Path::new(&caller.file)
.file_name()
.map(|f| f.to_string_lossy().to_string())
.unwrap_or_else(|| caller.file.clone());
output.push_str(&format!("- {} ({}:{})\n", caller.name, short, caller.start_line));
}
output.push('\n');
}
output.push_str(&format!(
"Token estimate: {} | Functions: {}\n",
self.token_count,
self.function_count()
));
output
}
}
#[derive(Debug, Clone)]
pub enum SourceInput<'a> {
Path(&'a str),
Source {
code: &'a str,
language: &'a str,
},
}
impl<'a> SourceInput<'a> {
pub fn resolve(&self) -> Result<(Vec<u8>, &'static dyn Language, Option<&'a str>)> {
let registry = LanguageRegistry::global();
match self {
SourceInput::Path(path) => {
let p = Path::new(path);
let lang = registry.detect_language(p).ok_or_else(|| {
BrrrError::UnsupportedLanguage(
p.extension()
.map(|e| e.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string()),
)
})?;
let source = std::fs::read(p)
.map_err(|e| BrrrError::io_with_path(e, p))?;
Ok((source, lang, Some(*path)))
}
SourceInput::Source { code, language } => {
let lang = registry
.get_by_name(language)
.ok_or_else(|| BrrrError::UnsupportedLanguage((*language).to_string()))?;
Ok((code.as_bytes().to_vec(), lang, None))
}
}
}
}
pub fn get_tree(
path: &str,
ext_filter: Option<&str>,
exclude_hidden: bool,
respect_ignore: bool,
) -> Result<FileTreeEntry> {
let ext_vec: Vec<String> = ext_filter
.map(|e| vec![e.to_string()])
.unwrap_or_default();
let show_hidden = !exclude_hidden;
let no_ignore = !respect_ignore;
ast::file_tree(path, &ext_vec, show_hidden, no_ignore, None)
}
#[inline]
pub fn get_tree_default(path: &str, ext_filter: Option<&str>) -> Result<FileTreeEntry> {
get_tree(path, ext_filter, true, true)
}
pub fn get_structure(
path: &str,
lang_filter: Option<&str>,
max_results: usize,
respect_ignore: bool,
) -> Result<CodeStructure> {
let no_ignore = !respect_ignore;
ast::code_structure(path, lang_filter, max_results, no_ignore)
}
#[inline]
pub fn get_structure_default(
path: &str,
lang_filter: Option<&str>,
max_results: usize,
) -> Result<CodeStructure> {
get_structure(path, lang_filter, max_results, true)
}
pub fn extract_file(file_path: &str, base_path: Option<&str>) -> Result<ModuleInfo> {
ast::extract_file(file_path, base_path)
}
#[inline]
pub fn extract_file_unchecked(file_path: &str) -> Result<ModuleInfo> {
ast::extract_file_unchecked(file_path)
}
pub fn extract_from_source(source: &str, language: &str) -> Result<ModuleInfo> {
ast::AstExtractor::extract_from_source(source, language)
}
pub fn get_imports(file_path: &str, language: Option<&str>) -> Result<Vec<ImportInfo>> {
use std::path::Path;
let path = Path::new(file_path);
let registry = LanguageRegistry::global();
if let Some(lang_name) = language {
if registry.get_by_name(lang_name).is_none() {
return Err(BrrrError::UnsupportedLanguage(lang_name.to_string()));
}
}
ast::extract_imports(path)
}
pub fn get_context(project: &str, entry_point: &str, depth: usize) -> Result<serde_json::Value> {
callgraph::get_context_with_lang(project, entry_point, depth, None)
}
pub fn query(
project: &str,
entry_point: &str,
depth: usize,
language: Option<&str>,
) -> Result<String> {
let result = callgraph::get_context_with_lang(project, entry_point, depth, language)?;
Ok(result
.get("llm_context")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string())
}
pub fn get_cfg(file: &str, function: &str, language: Option<&str>) -> Result<CFGInfo> {
cfg::extract_with_language(file, function, language)
}
#[inline]
pub fn get_cfg_auto(file: &str, function: &str) -> Result<CFGInfo> {
get_cfg(file, function, None)
}
pub fn get_cfg_from_source(source: &str, function: &str, language: &str) -> Result<CFGInfo> {
cfg::CfgBuilder::extract_from_source(source, language, function)
}
pub fn get_cfg_blocks(file: &str, function: &str) -> Result<Vec<CFGBlock>> {
let cfg = get_cfg(file, function, None)?;
Ok(cfg.blocks.into_values().collect())
}
pub fn get_cfg_edges(file: &str, function: &str) -> Result<Vec<CFGEdge>> {
let cfg = get_cfg(file, function, None)?;
Ok(cfg.edges)
}
pub fn get_cfg_mermaid(file: &str, function: &str) -> Result<String> {
let cfg = get_cfg(file, function, None)?;
Ok(cfg::to_mermaid(&cfg))
}
pub fn get_cfg_dot(file: &str, function: &str) -> Result<String> {
let cfg = get_cfg(file, function, None)?;
Ok(cfg::to_dot(&cfg))
}
pub fn get_cfg_ascii(file: &str, function: &str) -> Result<String> {
let cfg = get_cfg(file, function, None)?;
Ok(cfg::to_ascii(&cfg))
}
pub fn get_cfg_json(file: &str, function: &str) -> Result<serde_json::Value> {
let cfg = get_cfg(file, function, None)?;
Ok(serde_json::to_value(&cfg)?)
}
pub fn get_dfg(file: &str, function: &str, language: Option<&str>) -> Result<DFGInfo> {
dfg::extract_with_language(file, function, language)
}
#[inline]
pub fn get_dfg_auto(file: &str, function: &str) -> Result<DFGInfo> {
get_dfg(file, function, None)
}
pub fn get_dfg_from_source(source: &str, function: &str, language: &str) -> Result<DFGInfo> {
dfg::DfgBuilder::extract_from_source(source, language, function)
}
pub fn get_dfg_edges(file: &str, function: &str) -> Result<Vec<DataflowEdge>> {
let dfg = get_dfg(file, function, None)?;
Ok(dfg.edges)
}
pub fn get_dfg_variables(file: &str, function: &str) -> Result<Vec<String>> {
let dfg = get_dfg(file, function, None)?;
use std::collections::HashSet;
let vars: HashSet<_> = dfg.edges.iter().map(|e| e.variable.clone()).collect();
Ok(vars.into_iter().collect())
}
pub fn get_def_use_chains(file: &str, function: &str, variable: &str) -> Result<Vec<(usize, usize)>> {
let dfg = get_dfg(file, function, None)?;
let chains: Vec<_> = dfg
.edges
.iter()
.filter(|e| e.variable == variable)
.map(|e| (e.from_line, e.to_line))
.collect();
Ok(chains)
}
pub fn get_pdg(file: &str, function: &str, language: Option<&str>) -> Result<PDGInfo> {
pdg::build_pdg_with_language(file, function, language)
}
#[inline]
pub fn get_pdg_auto(file: &str, function: &str) -> Result<PDGInfo> {
get_pdg(file, function, None)
}
pub fn get_slice(
file: &str,
function: &str,
line: usize,
direction: Option<&str>,
variable: Option<&str>,
language: Option<&str>,
) -> Result<Vec<usize>> {
if line == 0 {
return Err(BrrrError::InvalidArgument(
"Line numbers are 1-indexed, got 0".to_string(),
));
}
let pdg_info = pdg::build_pdg_with_language(file, function, language)?;
let criteria = match variable {
Some(var) => pdg::SliceCriteria::at_line_variable(line, var),
None => pdg::SliceCriteria::at_line(line),
};
let dir = direction.unwrap_or("backward");
let result = match dir {
"backward" => pdg::backward_slice(&pdg_info, &criteria),
"forward" => pdg::forward_slice(&pdg_info, &criteria),
_ => {
return Err(BrrrError::InvalidArgument(format!(
"Invalid direction '{}', expected 'backward' or 'forward'",
dir
)))
}
};
Ok(result.lines)
}
pub fn get_slice_from_source(
source: &str,
function: &str,
line: usize,
direction: Option<&str>,
variable: Option<&str>,
language: &str,
) -> Result<Vec<usize>> {
if line == 0 {
return Err(BrrrError::InvalidArgument(
"Line numbers are 1-indexed, got 0".to_string(),
));
}
let dfg_info = dfg::DfgBuilder::extract_from_source(source, language, function)?;
let criteria = match variable {
Some(var) => dfg::SliceCriteria::at_line_variable(line, var),
None => dfg::SliceCriteria::at_line(line),
};
let dir = direction.unwrap_or("backward");
let result = match dir {
"backward" => dfg::backward_slice(&dfg_info, &criteria).lines,
"forward" => dfg::forward_slice(&dfg_info, &criteria).lines,
_ => {
return Err(BrrrError::InvalidArgument(format!(
"Invalid direction '{}', expected 'backward' or 'forward'",
dir
)))
}
};
Ok(result)
}
pub fn get_backward_slice(file: &str, function: &str, line: usize) -> Result<Vec<usize>> {
get_slice(file, function, line, Some("backward"), None, None)
}
pub fn get_slice_dfg_only(file: &str, function: &str, line: usize) -> Result<Vec<usize>> {
if line == 0 {
return Err(BrrrError::InvalidArgument(
"Line numbers are 1-indexed, got 0".to_string(),
));
}
dfg::get_slice(file, function, line)
}
pub fn get_forward_slice(file: &str, function: &str, line: usize) -> Result<Vec<usize>> {
if line == 0 {
return Err(BrrrError::InvalidArgument(
"Line numbers are 1-indexed, got 0".to_string(),
));
}
pdg::get_forward_slice(file, function, line)
}
pub fn get_pdg_slice(file: &str, function: &str, line: usize, direction: &str) -> Result<Vec<usize>> {
if line == 0 {
return Err(BrrrError::InvalidArgument(
"Line numbers are 1-indexed, got 0".to_string(),
));
}
match direction {
"backward" => pdg::get_slice(file, function, line),
"forward" => pdg::get_forward_slice(file, function, line),
_ => Err(BrrrError::InvalidArgument(format!(
"Invalid direction '{}', expected 'backward' or 'forward'",
direction
))),
}
}
pub fn build_callgraph(path: &str) -> Result<CallGraph> {
callgraph::build(path)
}
pub fn get_impact(path: &str, function: &str, depth: usize) -> Result<Vec<FunctionRef>> {
use callgraph::{cache, analyze_impact, ImpactConfig};
let project = std::path::Path::new(path);
let graph = cache::get_or_build_graph_with_config(project, None, false)?;
let config = ImpactConfig::new().with_depth(depth);
let result = analyze_impact(&graph, function, config);
Ok(result
.callers
.into_iter()
.map(|c| FunctionRef {
file: c.file,
name: c.name,
qualified_name: c.qualified_name,
})
.collect())
}
pub fn find_dead_code(path: &str) -> Result<Vec<FunctionRef>> {
use callgraph::{cache, dead, DeadCodeConfig};
let project = std::path::Path::new(path);
let mut graph = cache::get_or_build_graph_with_config(project, None, false)?;
graph.build_indexes();
let result = dead::analyze_dead_code_with_config(&graph, &DeadCodeConfig::default());
Ok(result
.dead_functions
.into_iter()
.map(|d| FunctionRef {
file: d.file,
name: d.name,
qualified_name: d.qualified_name,
})
.collect())
}
pub fn warm_callgraph(path: &str, langs: Option<&[String]>) -> Result<()> {
let project = std::path::Path::new(path);
let lang = langs.and_then(|l| l.first().map(|s| s.as_str()));
callgraph::warm_cache_with_config(project, lang, false)
}
pub fn extract_semantic_units(path: &str, lang: &str) -> Result<Vec<EmbeddingUnit>> {
semantic::extract_units(path, lang)
}
pub fn extract_semantic_units_with_callgraph(path: &str, lang: &str) -> Result<Vec<EmbeddingUnit>> {
semantic::extract_units_with_callgraph(path, lang)
}
pub fn extract_file_units(file_path: &str) -> Result<Vec<EmbeddingUnit>> {
semantic::extract_units_from_file(file_path)
}
pub fn build_embedding_text(unit: &EmbeddingUnit) -> String {
semantic::build_embedding_text(unit)
}
pub fn detect_semantic_patterns(code: &str) -> Vec<String> {
semantic::detect_semantic_patterns(code)
}
pub fn count_tokens(text: &str) -> usize {
semantic::count_tokens(text)
}
pub fn scan_project_files(
root: &str,
language: Option<&str>,
_respect_ignore: bool, ) -> Result<ScanResult> {
let scanner = ProjectScanner::new(root)?;
match language {
Some(lang) => scanner.scan_language_with_errors(lang),
None => scanner.scan_files_with_errors(),
}
}
pub fn scan_extensions(root: &str, extensions: &[&str]) -> Result<Vec<std::path::PathBuf>> {
let scanner = ProjectScanner::new(root)?;
scanner.scan_extensions(extensions)
}
pub fn get_project_metadata(root: &str, language: Option<&str>) -> Result<Vec<FileMetadata>> {
let scanner = ProjectScanner::new(root)?;
match language {
Some(lang) => scanner.scan_language_with_metadata(lang),
None => scanner.scan_with_metadata(),
}
}
pub fn scan_with_config(root: &str, config: &ScanConfig) -> Result<ScanResult> {
let scanner = ProjectScanner::new(root)?;
scanner.scan_with_config(config)
}
pub fn estimate_file_count(root: &str) -> Result<usize> {
let scanner = ProjectScanner::new(root)?;
scanner.estimate_file_count()
}
#[derive(Debug, Clone, Default)]
pub struct IndexingConfig {
pub language: Option<String>,
pub respect_ignore: bool,
pub parallel: bool,
pub include_tests: bool,
}
impl IndexingConfig {
pub fn new() -> Self {
Self {
language: None,
respect_ignore: true,
parallel: true,
include_tests: true,
}
}
pub fn with_language(mut self, lang: &str) -> Self {
self.language = Some(lang.to_string());
self
}
pub fn with_respect_ignore(mut self, respect: bool) -> Self {
self.respect_ignore = respect;
self
}
pub fn with_parallel(mut self, parallel: bool) -> Self {
self.parallel = parallel;
self
}
pub fn exclude_tests(mut self) -> Self {
self.include_tests = false;
self
}
pub fn include_tests(mut self) -> Self {
self.include_tests = true;
self
}
}
pub fn build_function_index(root: &str, language: Option<&str>) -> Result<FunctionIndex> {
use std::path::Path;
let scanner = ProjectScanner::new(root)?;
let project_root = Path::new(root);
let files = match language {
Some(lang) => scanner.scan_language(lang)?,
None => scanner.scan_files()?,
};
FunctionIndex::build_with_root(&files, Some(project_root))
}
pub fn build_function_index_with_config(root: &str, config: &IndexingConfig) -> Result<FunctionIndex> {
use std::path::Path;
let scanner = ProjectScanner::new(root)?;
let project_root = Path::new(root);
let mut scan_config = match &config.language {
Some(lang) => ScanConfig::for_language(lang),
None => ScanConfig::default(),
};
if !config.include_tests {
scan_config = scan_config.with_excludes(&[
"**/test/**",
"**/tests/**",
"**/*_test.*",
"**/*_spec.*",
"**/test_*.*",
]);
}
let result = scanner.scan_with_config(&scan_config)?;
FunctionIndex::build_with_root(&result.files, Some(project_root))
}
#[derive(Debug, Clone)]
pub struct ImporterInfo {
pub file: std::path::PathBuf,
pub import: ImportInfo,
}
pub fn get_importers(
root: &str,
module: &str,
language: Option<&str>,
) -> Result<Vec<ImporterInfo>> {
let scanner = ProjectScanner::new(root)?;
let files = match language {
Some(lang) => scanner.scan_language(lang)?,
None => scanner.scan_files()?,
};
let mut importers = Vec::new();
for file_path in files {
let imports = match ast::extract_imports(&file_path) {
Ok(i) => i,
Err(_) => continue, };
for import in imports {
if import_matches_module(&import, module) {
importers.push(ImporterInfo {
file: file_path.clone(),
import,
});
}
}
}
Ok(importers)
}
fn import_matches_module(import: &ImportInfo, module: &str) -> bool {
if import.module == module {
return true;
}
if import.module.ends_with(&format!(".{}", module)) {
return true;
}
if import.module.starts_with(&format!("{}.", module)) {
return true;
}
if import.module.contains(&format!(".{}.", module)) {
return true;
}
if import.names.iter().any(|name| name == module) {
return true;
}
false
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct IntraFileCall {
pub caller: String,
pub callee: String,
pub line: usize,
pub column: usize,
}
pub fn get_intra_file_calls(file_path: &str) -> Result<std::collections::HashMap<String, Vec<String>>> {
use std::collections::{HashMap, HashSet};
use std::path::Path;
let path = Path::new(file_path);
let registry = LanguageRegistry::global();
let lang = registry.detect_language(path).ok_or_else(|| {
BrrrError::UnsupportedLanguage(
path.extension()
.and_then(|e| e.to_str())
.unwrap_or("unknown")
.to_string(),
)
})?;
let source = std::fs::read(path)
.map_err(|e| BrrrError::io_with_path(e, path))?;
let mut parser = lang.parser_for_path(path)?;
let tree = parser.parse(&source, None).ok_or_else(|| BrrrError::Parse {
file: file_path.to_string(),
message: "Failed to parse file".to_string(),
})?;
let module_info = ast::AstExtractor::extract_file(path)?;
let mut defined_functions: HashSet<String> = HashSet::new();
for func in &module_info.functions {
defined_functions.insert(func.name.clone());
}
for class in &module_info.classes {
for method in &class.methods {
defined_functions.insert(method.name.clone());
}
}
struct FuncRange {
name: String,
start_line: usize,
end_line: usize,
}
let mut func_ranges: Vec<FuncRange> = Vec::new();
for func in &module_info.functions {
func_ranges.push(FuncRange {
name: func.name.clone(),
start_line: func.line_number,
end_line: func.end_line_number.unwrap_or(func.line_number),
});
}
for class in &module_info.classes {
for method in &class.methods {
func_ranges.push(FuncRange {
name: method.name.clone(),
start_line: method.line_number,
end_line: method.end_line_number.unwrap_or(method.line_number),
});
}
}
func_ranges.sort_by_key(|r| r.start_line);
let mut calls: HashMap<String, Vec<String>> = HashMap::new();
for func_name in &defined_functions {
calls.insert(func_name.clone(), Vec::new());
}
let query_str = lang.call_query();
let query = tree_sitter::Query::new(&tree.language(), query_str).map_err(|e| {
BrrrError::TreeSitter(format!("Failed to compile call query: {}", e))
})?;
let mut cursor = tree_sitter::QueryCursor::new();
let mut matches = cursor.matches(&query, tree.root_node(), source.as_slice());
use streaming_iterator::StreamingIterator;
while let Some(match_) = matches.next() {
let callee_capture = match_.captures.iter().find(|c| {
let name = &query.capture_names()[c.index as usize];
*name == "callee"
});
if let Some(capture) = callee_capture {
let callee_node = capture.node;
let callee_name = std::str::from_utf8(
&source[callee_node.start_byte()..callee_node.end_byte()],
)
.unwrap_or("")
.to_string();
if !defined_functions.contains(&callee_name) {
continue;
}
let call_line = callee_node.start_position().row + 1;
for func_range in &func_ranges {
if call_line >= func_range.start_line && call_line <= func_range.end_line {
calls
.entry(func_range.name.clone())
.or_default()
.push(callee_name.clone());
break;
}
}
}
}
for callees in calls.values_mut() {
callees.sort();
callees.dedup();
}
Ok(calls)
}
pub fn get_intra_file_calls_detailed(file_path: &str) -> Result<Vec<IntraFileCall>> {
use std::collections::HashSet;
use std::path::Path;
let path = Path::new(file_path);
let registry = LanguageRegistry::global();
let lang = registry.detect_language(path).ok_or_else(|| {
BrrrError::UnsupportedLanguage(
path.extension()
.and_then(|e| e.to_str())
.unwrap_or("unknown")
.to_string(),
)
})?;
let source = std::fs::read(path)
.map_err(|e| BrrrError::io_with_path(e, path))?;
let mut parser = lang.parser_for_path(path)?;
let tree = parser.parse(&source, None).ok_or_else(|| BrrrError::Parse {
file: file_path.to_string(),
message: "Failed to parse file".to_string(),
})?;
let module_info = ast::AstExtractor::extract_file(path)?;
let mut defined_functions: HashSet<String> = HashSet::new();
for func in &module_info.functions {
defined_functions.insert(func.name.clone());
}
for class in &module_info.classes {
for method in &class.methods {
defined_functions.insert(method.name.clone());
}
}
struct FuncRange {
name: String,
start_line: usize,
end_line: usize,
}
let mut func_ranges: Vec<FuncRange> = Vec::new();
for func in &module_info.functions {
func_ranges.push(FuncRange {
name: func.name.clone(),
start_line: func.line_number,
end_line: func.end_line_number.unwrap_or(func.line_number),
});
}
for class in &module_info.classes {
for method in &class.methods {
func_ranges.push(FuncRange {
name: method.name.clone(),
start_line: method.line_number,
end_line: method.end_line_number.unwrap_or(method.line_number),
});
}
}
func_ranges.sort_by_key(|r| r.start_line);
let mut detailed_calls: Vec<IntraFileCall> = Vec::new();
let query_str = lang.call_query();
let query = tree_sitter::Query::new(&tree.language(), query_str).map_err(|e| {
BrrrError::TreeSitter(format!("Failed to compile call query: {}", e))
})?;
let mut cursor = tree_sitter::QueryCursor::new();
let mut matches = cursor.matches(&query, tree.root_node(), source.as_slice());
use streaming_iterator::StreamingIterator;
while let Some(match_) = matches.next() {
let callee_capture = match_.captures.iter().find(|c| {
let name = &query.capture_names()[c.index as usize];
*name == "callee"
});
if let Some(capture) = callee_capture {
let callee_node = capture.node;
let callee_name = std::str::from_utf8(
&source[callee_node.start_byte()..callee_node.end_byte()],
)
.unwrap_or("")
.to_string();
if !defined_functions.contains(&callee_name) {
continue;
}
let position = callee_node.start_position();
let call_line = position.row + 1;
let call_column = position.column;
for func_range in &func_ranges {
if call_line >= func_range.start_line && call_line <= func_range.end_line {
detailed_calls.push(IntraFileCall {
caller: func_range.name.clone(),
callee: callee_name.clone(),
line: call_line,
column: call_column,
});
break;
}
}
}
}
detailed_calls.sort_by(|a, b| {
a.line.cmp(&b.line).then_with(|| a.column.cmp(&b.column))
});
Ok(detailed_calls)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_public_api_types_exported() {
fn _assert_types() {
let _: Option<FileTreeEntry> = None;
let _: Option<CodeStructure> = None;
let _: Option<ModuleInfo> = None;
let _: Option<SourceInput> = None;
let _: Option<FunctionInfo> = None;
let _: Option<ClassInfo> = None;
let _: Option<ImportInfo> = None;
let _: Option<CFGInfo> = None;
let _: Option<CFGBlock> = None;
let _: Option<CFGEdge> = None;
let _: Option<BlockId> = None;
let _: Option<DFGInfo> = None;
let _: Option<DataflowEdge> = None;
let _: Option<DataflowKind> = None;
let _: Option<PDGInfo> = None;
let _: Option<ControlDependence> = None;
let _: Option<BranchType> = None;
let _: Option<SliceCriteria> = None;
let _: Option<SliceResult> = None;
let _: Option<SliceMetrics> = None;
let _: Option<CallGraph> = None;
let _: Option<CallEdge> = None;
let _: Option<FunctionRef> = None;
let _: Option<BrrrError> = None;
let _: Option<FunctionIndex> = None;
let _: Option<FunctionDef> = None;
let _: Option<IndexStats> = None;
let _: Option<IndexingConfig> = None;
let _: Option<EmbeddingUnit> = None;
let _: Option<SearchResult> = None;
let _: Option<UnitKind> = None;
let _: Option<CodeComplexity> = None;
let _: Option<ChunkInfo> = None;
let _: Option<ProjectScanner> = None;
let _: Option<ScanConfig> = None;
let _: Option<ScanResult> = None;
let _: Option<ScanError> = None;
let _: Option<ScanErrorKind> = None;
let _: Option<FileMetadata> = None;
let _: Option<ErrorHandling> = None;
let _: Option<DeadCodeConfig> = None;
let _: Option<DeadCodeResult> = None;
let _: Option<DeadCodeStats> = None;
let _: Option<DeadFunction> = None;
let _: Option<DeadReason> = None;
let _: Option<EntryPointKind> = None;
let _: Option<ImpactConfig> = None;
let _: Option<ImpactResult> = None;
let _: Option<CallerInfo> = None;
let _: Option<ArchAnalysis> = None;
let _: Option<ArchStats> = None;
let _: Option<CycleDependency> = None;
let _: Option<CachedCallGraph> = None;
let _: Option<CachedEdge> = None;
let _: Option<ImporterInfo> = None;
let _: Option<IntraFileCall> = None;
let _: Option<FunctionContext> = None;
let _: Option<RelevantContext> = None;
}
}
#[test]
fn test_public_api_functions_exist() {
fn _assert_functions() {
let _ = get_tree as fn(&str, Option<&str>, bool, bool) -> Result<FileTreeEntry>;
let _ = get_structure as fn(&str, Option<&str>, usize, bool) -> Result<CodeStructure>;
let _ = extract_file as fn(&str, Option<&str>) -> Result<ModuleInfo>;
let _ = extract_file_unchecked as fn(&str) -> Result<ModuleInfo>;
let _ = extract_from_source as fn(&str, &str) -> Result<ModuleInfo>;
let _ = get_imports as fn(&str, Option<&str>) -> Result<Vec<ImportInfo>>;
let _ = get_context as fn(&str, &str, usize) -> Result<serde_json::Value>;
let _ = get_cfg as fn(&str, &str, Option<&str>) -> Result<CFGInfo>;
let _ = get_cfg_auto as fn(&str, &str) -> Result<CFGInfo>;
let _ = get_cfg_from_source as fn(&str, &str, &str) -> Result<CFGInfo>;
let _ = get_dfg as fn(&str, &str, Option<&str>) -> Result<DFGInfo>;
let _ = get_dfg_auto as fn(&str, &str) -> Result<DFGInfo>;
let _ = get_dfg_from_source as fn(&str, &str, &str) -> Result<DFGInfo>;
let _ = get_pdg as fn(&str, &str, Option<&str>) -> Result<PDGInfo>;
let _ = get_pdg_auto as fn(&str, &str) -> Result<PDGInfo>;
let _ = get_slice as fn(&str, &str, usize, Option<&str>, Option<&str>, Option<&str>) -> Result<Vec<usize>>;
let _ = get_slice_from_source as fn(&str, &str, usize, Option<&str>, Option<&str>, &str) -> Result<Vec<usize>>;
let _ = get_backward_slice as fn(&str, &str, usize) -> Result<Vec<usize>>;
let _ = get_slice_dfg_only as fn(&str, &str, usize) -> Result<Vec<usize>>;
let _ = get_forward_slice as fn(&str, &str, usize) -> Result<Vec<usize>>;
let _ = get_pdg_slice as fn(&str, &str, usize, &str) -> Result<Vec<usize>>;
let _ = pdg_backward_slice as fn(&PDGInfo, &SliceCriteria) -> SliceResult;
let _ = pdg_forward_slice as fn(&PDGInfo, &SliceCriteria) -> SliceResult;
let _ = build_callgraph as fn(&str) -> Result<CallGraph>;
let _ = get_impact as fn(&str, &str, usize) -> Result<Vec<FunctionRef>>;
let _ = find_dead_code as fn(&str) -> Result<Vec<FunctionRef>>;
let _ = warm_callgraph as fn(&str, Option<&[String]>) -> Result<()>;
let _ = extract_semantic_units as fn(&str, &str) -> Result<Vec<EmbeddingUnit>>;
let _ = extract_semantic_units_with_callgraph as fn(&str, &str) -> Result<Vec<EmbeddingUnit>>;
let _ = extract_file_units as fn(&str) -> Result<Vec<EmbeddingUnit>>;
let _ = build_embedding_text as fn(&EmbeddingUnit) -> String;
let _ = count_tokens as fn(&str) -> usize;
let _ = scan_project_files as fn(&str, Option<&str>, bool) -> Result<ScanResult>;
let _ = scan_extensions as fn(&str, &[&str]) -> Result<Vec<std::path::PathBuf>>;
let _ = get_project_metadata as fn(&str, Option<&str>) -> Result<Vec<FileMetadata>>;
let _ = scan_with_config as fn(&str, &ScanConfig) -> Result<ScanResult>;
let _ = estimate_file_count as fn(&str) -> Result<usize>;
let _ = build_function_index as fn(&str, Option<&str>) -> Result<FunctionIndex>;
let _ = build_function_index_with_config as fn(&str, &IndexingConfig) -> Result<FunctionIndex>;
let _ = get_importers as fn(&str, &str, Option<&str>) -> Result<Vec<ImporterInfo>>;
let _ = get_intra_file_calls as fn(&str) -> Result<std::collections::HashMap<String, Vec<String>>>;
let _ = get_intra_file_calls_detailed as fn(&str) -> Result<Vec<IntraFileCall>>;
}
}
#[test]
fn test_semantic_constants_exported() {
assert!(MAX_EMBEDDING_TOKENS > 0);
assert!(MAX_CODE_PREVIEW_TOKENS > 0);
assert!(CHUNK_OVERLAP_TOKENS > 0);
}
#[test]
fn test_semantic_pattern_exports() {
let pattern = &SEMANTIC_PATTERNS[0];
let _: &SemanticPattern = pattern;
assert!(!pattern.name.is_empty());
assert!(!pattern.pattern.is_empty());
let patterns = detect_semantic_patterns("def validate_user(): assert x");
assert!(patterns.contains(&"validation".to_string()));
assert!(!SEMANTIC_PATTERNS.is_empty());
}
#[test]
fn test_callgraph_advanced_type_exports() {
let dead_config = DeadCodeConfig::default();
assert!(!dead_config.include_public_api_patterns);
assert!(dead_config.min_confidence > 0.0);
let _: DeadCodeResult;
let _: DeadCodeStats;
let _: DeadFunction;
let _: DeadReason;
let kind = classify_entry_point("main");
assert_eq!(kind, Some(EntryPointKind::Main));
let test_kind = classify_entry_point("test_something");
assert_eq!(test_kind, Some(EntryPointKind::Test));
let impact_config = ImpactConfig::default();
assert_eq!(impact_config.max_depth, 0);
assert!(!impact_config.exclude_tests);
let _: ImpactResult;
let _: CallerInfo;
let _: ArchAnalysis;
let _: ArchStats;
let _: CycleDependency;
let _: CachedCallGraph;
let _: CachedEdge;
let _ = analyze_dead_code as fn(&CallGraph) -> DeadCodeResult;
let _ = analyze_dead_code_with_config as fn(&CallGraph, &DeadCodeConfig) -> DeadCodeResult;
let _ = analyze_impact as fn(&CallGraph, &str, ImpactConfig) -> ImpactResult;
let _ = detect_entry_points_with_config as fn(&CallGraph, &DeadCodeConfig) -> Vec<FunctionRef>;
let _ = analyze_architecture as fn(&CallGraph) -> ArchAnalysis;
}
#[test]
fn test_extract_from_source() {
let source = r#"
def greet(name: str) -> str:
"""Say hello."""
return f"Hello, {name}!"
class Greeter:
def __init__(self, prefix: str):
self.prefix = prefix
"#;
let module = extract_from_source(source, "python").unwrap();
assert_eq!(module.language, "python");
assert_eq!(module.path, "<string>");
assert_eq!(module.functions.len(), 1);
assert_eq!(module.functions[0].name, "greet");
assert_eq!(module.classes.len(), 1);
assert_eq!(module.classes[0].name, "Greeter");
}
#[test]
fn test_get_cfg_from_source() {
let source = r#"
def process(x):
if x > 0:
return x * 2
return 0
"#;
let cfg = get_cfg_from_source(source, "process", "python").unwrap();
assert_eq!(cfg.function_name, "process");
assert!(cfg.blocks.len() >= 2, "CFG should have multiple blocks");
}
#[test]
fn test_get_dfg_from_source() {
let source = r#"
def compute(x, y):
z = x + y
result = z * 2
return result
"#;
let dfg = get_dfg_from_source(source, "compute", "python").unwrap();
assert_eq!(dfg.function_name, "compute");
assert!(dfg.definitions.contains_key("z"), "Should track 'z' definition");
assert!(dfg.definitions.contains_key("result"), "Should track 'result' definition");
assert!(dfg.uses.contains_key("x"), "Should track 'x' use");
assert!(dfg.uses.contains_key("y"), "Should track 'y' use");
}
#[test]
fn test_get_slice_from_source() {
let source = r#"
def compute(x):
a = x + 1
b = a * 2
c = b + x
return c
"#;
let slice = get_slice_from_source(source, "compute", 5, None, None, "python").unwrap();
assert!(!slice.is_empty(), "Backward slice should not be empty");
let fwd_slice = get_slice_from_source(source, "compute", 3, Some("forward"), None, "python").unwrap();
assert!(!fwd_slice.is_empty(), "Forward slice should not be empty");
}
#[test]
fn test_source_input_enum() {
let path_input: SourceInput = SourceInput::Path("./test.py");
match &path_input {
SourceInput::Path(p) => assert_eq!(*p, "./test.py"),
_ => panic!("Expected Path variant"),
}
let source_input: SourceInput = SourceInput::Source {
code: "def foo(): pass",
language: "python",
};
match &source_input {
SourceInput::Source { code, language } => {
assert_eq!(*code, "def foo(): pass");
assert_eq!(*language, "python");
}
_ => panic!("Expected Source variant"),
}
let (bytes, lang, path) = source_input.resolve().unwrap();
assert_eq!(bytes, b"def foo(): pass");
assert_eq!(lang.name(), "python");
assert!(path.is_none());
}
#[test]
fn test_get_intra_file_calls_python() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = r#"
def helper():
return 42
def process(x):
result = helper()
return result * 2
def main():
value = process(10)
helper()
return value
"#;
let mut file = tempfile::Builder::new()
.suffix(".py")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls(file.path().to_str().unwrap()).unwrap();
assert!(calls.contains_key("main"), "Should have main in call map");
let main_calls = calls.get("main").unwrap();
assert!(main_calls.contains(&"process".to_string()), "main should call process");
assert!(main_calls.contains(&"helper".to_string()), "main should call helper");
assert!(calls.contains_key("process"), "Should have process in call map");
let process_calls = calls.get("process").unwrap();
assert!(process_calls.contains(&"helper".to_string()), "process should call helper");
assert!(calls.contains_key("helper"), "Should have helper in call map");
let helper_calls = calls.get("helper").unwrap();
assert!(helper_calls.is_empty(), "helper should not call any local functions");
}
#[test]
fn test_get_intra_file_calls_detailed_python() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = r#"def helper():
return 42
def process(x):
result = helper()
return result * 2
def main():
value = process(10)
helper()
return value
"#;
let mut file = tempfile::Builder::new()
.suffix(".py")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls_detailed(file.path().to_str().unwrap()).unwrap();
assert!(calls.len() >= 3, "Should have at least 3 intra-file calls, got {}", calls.len());
for window in calls.windows(2) {
assert!(
window[0].line <= window[1].line,
"Calls should be sorted by line number"
);
}
for call in &calls {
assert!(!call.caller.is_empty(), "Caller name should not be empty");
assert!(!call.callee.is_empty(), "Callee name should not be empty");
assert!(call.line > 0, "Line number should be positive");
}
}
#[test]
fn test_get_intra_file_calls_typescript() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = r#"
function helper(): number {
return 42;
}
function process(x: number): number {
const result = helper();
return result * 2;
}
function main(): number {
const value = process(10);
helper();
return value;
}
"#;
let mut file = tempfile::Builder::new()
.suffix(".ts")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls(file.path().to_str().unwrap()).unwrap();
assert!(calls.contains_key("main"), "Should have main in call map");
let main_calls = calls.get("main").unwrap();
assert!(main_calls.contains(&"process".to_string()), "main should call process");
assert!(main_calls.contains(&"helper".to_string()), "main should call helper");
}
#[test]
fn test_get_intra_file_calls_with_class_methods() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = r#"
def standalone():
return 1
class Calculator:
def add(self, a, b):
return a + b
def compute(self, x):
result = self.add(x, 1)
standalone()
return result
"#;
let mut file = tempfile::Builder::new()
.suffix(".py")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls(file.path().to_str().unwrap()).unwrap();
assert!(calls.contains_key("add"), "Should have add method in call map");
assert!(calls.contains_key("compute"), "Should have compute method in call map");
assert!(calls.contains_key("standalone"), "Should have standalone in call map");
let compute_calls = calls.get("compute").unwrap();
assert!(
compute_calls.contains(&"standalone".to_string()),
"compute should call standalone"
);
}
#[test]
fn test_get_intra_file_calls_no_external_calls() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = r#"
import os
def local_func():
return 42
def main():
# This calls os.path.join which is external
path = os.path.join("a", "b")
# This calls local_func which is internal
value = local_func()
# This calls print which is builtin/external
print(value)
return value
"#;
let mut file = tempfile::Builder::new()
.suffix(".py")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls(file.path().to_str().unwrap()).unwrap();
let main_calls = calls.get("main").unwrap();
assert_eq!(
main_calls.len(), 1,
"main should only call one local function, got {:?}", main_calls
);
assert!(
main_calls.contains(&"local_func".to_string()),
"main should call local_func"
);
}
#[test]
fn test_get_intra_file_calls_empty_file() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = "# Empty Python file with just a comment\n";
let mut file = tempfile::Builder::new()
.suffix(".py")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let calls = get_intra_file_calls(file.path().to_str().unwrap()).unwrap();
assert!(calls.is_empty(), "Empty file should have no calls");
}
#[test]
fn test_get_intra_file_calls_unsupported_language() {
use std::io::Write;
use tempfile::NamedTempFile;
let source = "Some random content";
let mut file = tempfile::Builder::new()
.suffix(".xyz")
.tempfile()
.unwrap();
file.write_all(source.as_bytes()).unwrap();
let result = get_intra_file_calls(file.path().to_str().unwrap());
assert!(result.is_err(), "Should return error for unsupported language");
assert!(matches!(result, Err(BrrrError::UnsupportedLanguage(_))));
}
#[test]
fn test_function_context_minimal() {
let ctx = FunctionContext::minimal("test_func", "src/main.rs", 10, "rust");
assert_eq!(ctx.name, "test_func");
assert_eq!(ctx.file, "src/main.rs");
assert_eq!(ctx.start_line, 10);
assert_eq!(ctx.end_line, 10);
assert_eq!(ctx.language, "rust");
assert!(ctx.signature.is_none());
assert!(ctx.docstring.is_none());
assert!(ctx.source.is_empty());
}
#[test]
fn test_function_context_from_function_info() {
let info = FunctionInfo {
name: "process".to_string(),
params: vec!["x: int".to_string(), "y: int".to_string()],
return_type: Some("int".to_string()),
docstring: Some("Process two numbers.".to_string()),
is_method: false,
is_async: false,
decorators: vec![],
line_number: 2,
end_line_number: Some(5),
language: "python".to_string(),
};
let source = "# Header\ndef process(x: int, y: int) -> int:\n \"\"\"Process two numbers.\"\"\"\n return x + y\n# Footer";
let ctx = FunctionContext::from_function_info(&info, "src/main.py", source, "python");
assert_eq!(ctx.name, "process");
assert_eq!(ctx.file, "src/main.py");
assert_eq!(ctx.start_line, 2);
assert_eq!(ctx.end_line, 5);
assert_eq!(ctx.language, "python");
assert!(ctx.signature.is_some());
assert!(ctx.signature.as_ref().unwrap().contains("process"));
assert_eq!(ctx.docstring, Some("Process two numbers.".to_string()));
assert!(ctx.source.contains("def process"));
}
#[test]
fn test_function_context_serialization() {
let ctx = FunctionContext::minimal("test", "test.py", 1, "python");
let json = serde_json::to_string(&ctx).unwrap();
assert!(json.contains("\"name\":\"test\""));
assert!(json.contains("\"file\":\"test.py\""));
assert!(json.contains("\"language\":\"python\""));
let deserialized: FunctionContext = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.name, ctx.name);
assert_eq!(deserialized.file, ctx.file);
assert_eq!(deserialized.language, ctx.language);
}
#[test]
fn test_relevant_context_new() {
let entry = FunctionContext::minimal("main", "src/main.rs", 1, "rust");
let ctx = RelevantContext::new(entry, 2);
assert_eq!(ctx.entry.name, "main");
assert_eq!(ctx.depth, 2);
assert!(ctx.callees.is_empty());
assert!(ctx.callers.is_empty());
assert_eq!(ctx.token_count, 0);
}
#[test]
fn test_relevant_context_builder_pattern() {
let entry = FunctionContext::minimal("main", "src/main.rs", 1, "rust");
let callee = FunctionContext::minimal("helper", "src/utils.rs", 10, "rust");
let caller = FunctionContext::minimal("test_main", "tests/test.rs", 5, "rust");
let ctx = RelevantContext::new(entry, 2)
.with_callee(callee)
.with_caller(caller)
.with_token_count(500);
assert_eq!(ctx.callees.len(), 1);
assert_eq!(ctx.callees[0].name, "helper");
assert_eq!(ctx.callers.len(), 1);
assert_eq!(ctx.callers[0].name, "test_main");
assert_eq!(ctx.token_count, 500);
assert_eq!(ctx.function_count(), 3);
}
#[test]
fn test_relevant_context_estimate_tokens() {
let mut entry = FunctionContext::minimal("main", "src/main.rs", 1, "rust");
entry.source = "fn main() { println!(\"hello\"); }".to_string();
let mut callee = FunctionContext::minimal("helper", "src/utils.rs", 10, "rust");
callee.source = "fn helper() {}".to_string();
let mut ctx = RelevantContext::new(entry, 1).with_callee(callee);
ctx.estimate_tokens();
assert_eq!(ctx.token_count, 11);
}
#[test]
fn test_relevant_context_to_llm_string() {
let mut entry = FunctionContext::minimal("process", "src/main.py", 10, "python");
entry.signature = Some("def process(x: int) -> int".to_string());
entry.docstring = Some("Process a number.".to_string());
let callee = FunctionContext::minimal("helper", "src/utils.py", 25, "python");
let caller = FunctionContext::minimal("main", "src/app.py", 5, "python");
let ctx = RelevantContext::new(entry, 2)
.with_callee(callee)
.with_caller(caller)
.with_token_count(100);
let output = ctx.to_llm_string();
assert!(output.contains("## Code Context: process (depth=2)"));
assert!(output.contains("### Entry Point: process (main.py:10)"));
assert!(output.contains("def process(x: int) -> int"));
assert!(output.contains("> Process a number."));
assert!(output.contains("### Calls:"));
assert!(output.contains("- helper (utils.py:25)"));
assert!(output.contains("### Called By:"));
assert!(output.contains("- main (app.py:5)"));
assert!(output.contains("Token estimate: 100"));
assert!(output.contains("Functions: 3"));
}
#[test]
fn test_relevant_context_serialization() {
let entry = FunctionContext::minimal("test", "test.py", 1, "python");
let callee = FunctionContext::minimal("helper", "helper.py", 5, "python");
let ctx = RelevantContext::new(entry, 1)
.with_callee(callee)
.with_token_count(50);
let json = serde_json::to_string(&ctx).unwrap();
assert!(json.contains("\"depth\":1"));
assert!(json.contains("\"token_count\":50"));
let deserialized: RelevantContext = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.entry.name, "test");
assert_eq!(deserialized.callees.len(), 1);
assert_eq!(deserialized.depth, 1);
assert_eq!(deserialized.token_count, 50);
}
#[test]
fn test_context_types_exported() {
fn _assert_context_types() {
let _: Option<FunctionContext> = None;
let _: Option<RelevantContext> = None;
}
}
}