use std::collections::BTreeSet;
use regex::Regex;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Lang {
Python,
Rust,
}
impl Lang {
#[must_use]
pub fn from_path(path: &str) -> Option<Self> {
match path.rsplit('.').next() {
Some("py") => Some(Self::Python),
Some("rs") => Some(Self::Rust),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Reference {
pub module: String,
pub name: Option<String>,
pub line: usize,
}
impl Reference {
#[must_use]
pub fn import_from(module: impl Into<String>, name: impl Into<String>, line: usize) -> Self {
Self {
module: module.into(),
name: Some(name.into()),
line,
}
}
#[must_use]
pub fn import_module(module: impl Into<String>, line: usize) -> Self {
Self {
module: module.into(),
name: None,
line,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Definition {
pub name: String,
pub kind: DefKind,
pub line: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DefKind {
Function,
Class,
Struct,
Enum,
Trait,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Resolution {
Resolved,
UnknownSymbol,
UnknownModule,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "verdict", rename_all = "snake_case")]
pub enum Verdict {
Resolved,
NotBuilt {
root: String,
},
Fabricated {
detail: String,
},
}
#[derive(Debug, Clone, Default)]
pub struct SymbolIndex {
modules: BTreeSet<String>,
symbols: BTreeSet<(String, String)>,
}
impl SymbolIndex {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_module(&mut self, module: impl Into<String>) {
self.modules.insert(module.into());
}
pub fn add_symbol(&mut self, module: impl Into<String>, symbol: impl Into<String>) {
let module = module.into();
self.modules.insert(module.clone());
self.symbols.insert((module, symbol.into()));
}
pub fn add_from_source(&mut self, module: &str, source: &str, lang: Lang) {
for def in extract_definitions(source, lang) {
self.add_symbol(module, def.name);
}
}
#[must_use]
pub fn has_module(&self, module: &str) -> bool {
self.modules.contains(module)
}
#[must_use]
pub fn resolve(&self, reference: &Reference) -> Resolution {
if !self.modules.contains(&reference.module) {
return Resolution::UnknownModule;
}
match &reference.name {
Some(name) => {
if self
.symbols
.contains(&(reference.module.clone(), name.clone()))
{
Resolution::Resolved
} else {
Resolution::UnknownSymbol
}
}
None => Resolution::Resolved,
}
}
#[must_use]
pub fn classify(&self, reference: &Reference, unbuilt_roots: &BTreeSet<String>) -> Verdict {
match self.resolve(reference) {
Resolution::Resolved => Verdict::Resolved,
Resolution::UnknownSymbol => Verdict::Fabricated {
detail: format!(
"module `{}` has no symbol `{}`",
reference.module,
reference.name.as_deref().unwrap_or("")
),
},
Resolution::UnknownModule => {
let root = module_root(&reference.module);
if unbuilt_roots.contains(root) {
Verdict::NotBuilt {
root: root.to_string(),
}
} else {
Verdict::Fabricated {
detail: format!("no module `{}`", reference.module),
}
}
}
}
}
}
fn module_root(module: &str) -> &str {
module.split(['.', ':']).next().unwrap_or(module)
}
#[must_use]
pub fn module_is_known(module: &str, known: &BTreeSet<String>) -> bool {
let parts: Vec<&str> = module.split('.').collect();
(1..=parts.len()).any(|i| known.contains(&parts[..i].join(".")))
}
#[must_use]
pub fn python_stdlib_modules() -> BTreeSet<String> {
[
"__future__",
"__main__",
"abc",
"argparse",
"asyncio",
"base64",
"collections",
"contextlib",
"copy",
"csv",
"dataclasses",
"datetime",
"decimal",
"enum",
"functools",
"glob",
"hashlib",
"io",
"itertools",
"json",
"logging",
"math",
"os",
"pathlib",
"random",
"re",
"shutil",
"subprocess",
"sys",
"tempfile",
"time",
"typing",
"unittest",
"uuid",
"warnings",
]
.iter()
.map(|s| (*s).to_string())
.collect()
}
#[must_use]
pub fn extract_references(source: &str, lang: Lang) -> Vec<Reference> {
match lang {
Lang::Python => extract_references_python(source),
Lang::Rust => extract_references_rust(source),
}
}
#[must_use]
pub fn extract_definitions(source: &str, lang: Lang) -> Vec<Definition> {
match lang {
Lang::Python => extract_definitions_python(source),
Lang::Rust => extract_definitions_rust(source),
}
}
fn extract_references_python(source: &str) -> Vec<Reference> {
let from_re = Regex::new(r"^\s*from\s+([\w.-]+)\s+import\s+(.+?)\s*$").unwrap();
let import_re = Regex::new(r"^\s*import\s+(.+?)\s*$").unwrap();
let mut refs = Vec::new();
for (i, line) in source.lines().enumerate() {
let lineno = i + 1;
if let Some(caps) = from_re.captures(line) {
let module = caps[1].to_string();
if module.starts_with('.') {
continue;
}
let names = caps[2].trim();
if names == "(" {
refs.push(Reference::import_module(module, lineno));
continue;
}
let before = refs.len();
for item in names.split(',') {
let name = item
.split_whitespace()
.next()
.unwrap_or("")
.trim_matches(|c| c == '(' || c == ')');
if name.is_empty() || name == "*" {
continue;
}
refs.push(Reference::import_from(module.clone(), name, lineno));
}
if refs.len() == before {
refs.push(Reference::import_module(module, lineno));
}
} else if let Some(caps) = import_re.captures(line) {
for item in caps[1].split(',') {
let module = item.split_whitespace().next().unwrap_or("");
if !module.is_empty() {
refs.push(Reference::import_module(module, lineno));
}
}
}
}
refs
}
fn extract_definitions_python(source: &str) -> Vec<Definition> {
let def_re = Regex::new(r"^\s*def\s+(\w+)").unwrap();
let class_re = Regex::new(r"^\s*class\s+(\w+)").unwrap();
let mut defs = Vec::new();
for (i, line) in source.lines().enumerate() {
let lineno = i + 1;
if let Some(caps) = def_re.captures(line) {
defs.push(Definition {
name: caps[1].to_string(),
kind: DefKind::Function,
line: lineno,
});
} else if let Some(caps) = class_re.captures(line) {
defs.push(Definition {
name: caps[1].to_string(),
kind: DefKind::Class,
line: lineno,
});
}
}
defs
}
fn extract_references_rust(source: &str) -> Vec<Reference> {
let use_re = Regex::new(r"^\s*(?:pub\s*(?:\([^)]*\)\s*)?)?use\s+(.+?)\s*;").unwrap();
let mut refs = Vec::new();
for (i, line) in source.lines().enumerate() {
let lineno = i + 1;
if let Some(caps) = use_re.captures(line) {
parse_rust_use(caps[1].trim(), lineno, &mut refs);
}
}
refs
}
fn parse_rust_use(path: &str, lineno: usize, refs: &mut Vec<Reference>) {
if let Some(open) = path.find("::{") {
let module = &path[..open];
let inner = path[open + 3..].trim_end_matches('}');
for item in inner.split(',') {
let item = item.split_whitespace().next().unwrap_or("");
if item.is_empty() || item == "*" || item == "self" {
refs.push(Reference::import_module(module, lineno));
} else {
refs.push(Reference::import_from(module, item, lineno));
}
}
return;
}
if let Some(module) = path.strip_suffix("::*") {
refs.push(Reference::import_module(module, lineno));
return;
}
let head = path.split_whitespace().next().unwrap_or(path);
if let Some(idx) = head.rfind("::") {
refs.push(Reference::import_from(
&head[..idx],
&head[idx + 2..],
lineno,
));
} else {
refs.push(Reference::import_module(head, lineno));
}
}
fn extract_definitions_rust(source: &str) -> Vec<Definition> {
let vis = r"(?:pub\s*(?:\([^)]*\)\s*)?)?";
let fn_re = Regex::new(&format!(
r"^\s*{vis}(?:async\s+|unsafe\s+|const\s+|extern\s+(?:\x22[^\x22]*\x22\s+)?)*fn\s+(\w+)"
))
.unwrap();
let struct_re = Regex::new(&format!(r"^\s*{vis}struct\s+(\w+)")).unwrap();
let enum_re = Regex::new(&format!(r"^\s*{vis}enum\s+(\w+)")).unwrap();
let trait_re = Regex::new(&format!(r"^\s*{vis}(?:unsafe\s+)?trait\s+(\w+)")).unwrap();
let mut defs = Vec::new();
for (i, line) in source.lines().enumerate() {
let lineno = i + 1;
let captured = [
(&fn_re, DefKind::Function),
(&struct_re, DefKind::Struct),
(&enum_re, DefKind::Enum),
(&trait_re, DefKind::Trait),
]
.into_iter()
.find_map(|(re, kind)| re.captures(line).map(|c| (c[1].to_string(), kind)));
if let Some((name, kind)) = captured {
defs.push(Definition {
name,
kind,
line: lineno,
});
}
}
defs
}
#[cfg(test)]
mod tests {
use super::*;
fn built() -> BTreeSet<String> {
BTreeSet::new()
}
#[test]
fn oracle_catches_the_second_nemotron_incident() {
let mut idx = SymbolIndex::new();
idx.add_symbol("newt_agent.core", "Router");
idx.add_symbol("newt_agent.core", "Tier");
idx.add_symbol("newt_agent.data", "load_csv_to_sqlite");
let fab1 = Reference::import_from("newt_core", "classify", 1);
let fab2 = Reference::import_from("newt_data", "DataStore", 2);
assert!(matches!(
idx.classify(&fab1, &built()),
Verdict::Fabricated { .. }
));
assert!(matches!(
idx.classify(&fab2, &built()),
Verdict::Fabricated { .. }
));
let real = Reference::import_from("newt_agent.core", "Router", 3);
assert_eq!(idx.classify(&real, &built()), Verdict::Resolved);
let bad_symbol = Reference::import_from("newt_agent.core", "classify", 4);
assert!(matches!(
idx.classify(&bad_symbol, &built()),
Verdict::Fabricated { .. }
));
}
#[test]
fn not_built_is_distinguished_from_fabricated() {
let mut idx = SymbolIndex::new();
idx.add_module("newt_agent");
let unbuilt: BTreeSet<String> = ["newt_agent".to_string()].into_iter().collect();
let r = Reference::import_from("newt_agent.core", "Router", 1);
assert_eq!(
idx.classify(&r, &unbuilt),
Verdict::NotBuilt {
root: "newt_agent".to_string()
}
);
let fab = Reference::import_from("newt_core", "classify", 2);
assert!(matches!(
idx.classify(&fab, &unbuilt),
Verdict::Fabricated { .. }
));
}
#[test]
fn extracts_python_imports() {
let src = "from newt_agent.core import Router, Tier\nimport os\nimport newt_data as nd\nfrom x import *\n";
let refs = extract_references(src, Lang::Python);
assert!(refs.contains(&Reference::import_from("newt_agent.core", "Router", 1)));
assert!(refs.contains(&Reference::import_from("newt_agent.core", "Tier", 1)));
assert!(refs.contains(&Reference::import_module("os", 2)));
assert!(refs.contains(&Reference::import_module("newt_data", 3)));
assert!(!refs.iter().any(|r| r.name.as_deref() == Some("*")));
}
#[test]
fn extracts_python_definitions() {
let src = "def foo():\n pass\n\nclass Bar:\n pass\n";
let defs = extract_definitions(src, Lang::Python);
assert!(defs
.iter()
.any(|d| d.name == "foo" && d.kind == DefKind::Function));
assert!(defs
.iter()
.any(|d| d.name == "Bar" && d.kind == DefKind::Class));
}
#[test]
fn add_from_source_builds_a_resolvable_index() {
let mut idx = SymbolIndex::new();
idx.add_from_source(
"mypkg.mod",
"def helper():\n pass\nclass Widget:\n pass\n",
Lang::Python,
);
assert_eq!(
idx.resolve(&Reference::import_from("mypkg.mod", "Widget", 1)),
Resolution::Resolved
);
assert_eq!(
idx.resolve(&Reference::import_from("mypkg.mod", "Nope", 1)),
Resolution::UnknownSymbol
);
assert_eq!(
idx.resolve(&Reference::import_from("other.mod", "Widget", 1)),
Resolution::UnknownModule
);
}
#[test]
fn bare_import_only_checks_module_existence() {
let mut idx = SymbolIndex::new();
idx.add_module("os");
assert_eq!(
idx.resolve(&Reference::import_module("os", 1)),
Resolution::Resolved
);
assert_eq!(
idx.resolve(&Reference::import_module("nope", 1)),
Resolution::UnknownModule
);
}
#[test]
fn extracts_rust_use_statements() {
let src = "use a::b::Thing;\n\
pub use c::d::{One, Two as T, self, *};\n\
use crate::scope::*;\n\
use std::collections::HashMap as Map;\n";
let refs = extract_references(src, Lang::Rust);
assert!(refs.contains(&Reference::import_from("a::b", "Thing", 1)));
assert!(refs.contains(&Reference::import_from("c::d", "One", 2)));
assert!(refs.contains(&Reference::import_from("c::d", "Two", 2))); assert!(refs.contains(&Reference::import_module("c::d", 2))); assert!(refs.contains(&Reference::import_module("crate::scope", 3)));
assert!(refs.contains(&Reference::import_from("std::collections", "HashMap", 4)));
}
#[test]
fn extracts_rust_definitions() {
let src = "pub fn alpha() {}\n\
async fn beta() {}\n\
pub(crate) struct Gamma { x: u8 }\n\
enum Delta { A, B }\n\
pub trait Epsilon {}\n\
const NOT_A_DEF: u8 = 0;\n";
let defs = extract_definitions(src, Lang::Rust);
assert!(defs
.iter()
.any(|d| d.name == "alpha" && d.kind == DefKind::Function));
assert!(defs
.iter()
.any(|d| d.name == "beta" && d.kind == DefKind::Function));
assert!(defs
.iter()
.any(|d| d.name == "Gamma" && d.kind == DefKind::Struct));
assert!(defs
.iter()
.any(|d| d.name == "Delta" && d.kind == DefKind::Enum));
assert!(defs
.iter()
.any(|d| d.name == "Epsilon" && d.kind == DefKind::Trait));
assert!(!defs.iter().any(|d| d.name == "NOT_A_DEF"));
}
#[test]
fn rust_references_resolve_against_the_same_general_core() {
let mut idx = SymbolIndex::new();
idx.add_symbol("crate::router", "Router");
let real = &extract_references("use crate::router::Router;", Lang::Rust)[0];
assert_eq!(idx.resolve(real), Resolution::Resolved);
let fab = &extract_references("use crate::nonsense::Widget;", Lang::Rust)[0];
assert!(matches!(
idx.classify(fab, &built()),
Verdict::Fabricated { .. }
));
}
#[test]
fn module_root_handles_python_and_rust_paths() {
assert_eq!(module_root("newt_agent.core"), "newt_agent");
assert_eq!(module_root("a::b::c"), "a");
assert_eq!(module_root("flat"), "flat");
}
#[test]
fn lang_from_path() {
assert_eq!(Lang::from_path("foo/bar.py"), Some(Lang::Python));
assert_eq!(Lang::from_path("src/lib.rs"), Some(Lang::Rust));
assert_eq!(Lang::from_path("README.md"), None);
}
#[test]
fn verdict_serializes_with_a_tag() {
let v = Verdict::Fabricated {
detail: "no module `x`".to_string(),
};
let json = serde_json::to_string(&v).unwrap();
assert!(json.contains("\"verdict\":\"fabricated\""));
}
#[test]
fn module_is_known_exact_single_component() {
let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
assert!(module_is_known("os", &known));
}
#[test]
fn module_is_known_prefix_covers_submodule() {
let known: BTreeSet<String> = ["newt_agent".to_string()].into_iter().collect();
assert!(module_is_known("newt_agent.core", &known));
}
#[test]
fn module_is_known_no_match() {
let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
assert!(!module_is_known("sys", &known));
}
#[test]
fn module_is_known_empty_module() {
let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
assert!(!module_is_known("", &known));
}
#[test]
fn module_is_known_substring_is_not_prefix() {
let known: BTreeSet<String> = ["pathlib".to_string()].into_iter().collect();
assert!(!module_is_known("os.path", &known));
}
}