1use std::collections::BTreeSet;
29
30use regex::Regex;
31use serde::{Deserialize, Serialize};
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum Lang {
36 Python,
38 Rust,
42}
43
44impl Lang {
45 #[must_use]
47 pub fn from_path(path: &str) -> Option<Self> {
48 match path.rsplit('.').next() {
49 Some("py") => Some(Self::Python),
50 Some("rs") => Some(Self::Rust),
51 _ => None,
52 }
53 }
54}
55
56#[derive(Debug, Clone, PartialEq, Eq)]
58pub struct Reference {
59 pub module: String,
62 pub name: Option<String>,
65 pub line: usize,
67}
68
69impl Reference {
70 #[must_use]
72 pub fn import_from(module: impl Into<String>, name: impl Into<String>, line: usize) -> Self {
73 Self {
74 module: module.into(),
75 name: Some(name.into()),
76 line,
77 }
78 }
79
80 #[must_use]
82 pub fn import_module(module: impl Into<String>, line: usize) -> Self {
83 Self {
84 module: module.into(),
85 name: None,
86 line,
87 }
88 }
89}
90
91#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct Definition {
94 pub name: String,
96 pub kind: DefKind,
98 pub line: usize,
100}
101
102#[derive(Debug, Clone, Copy, PartialEq, Eq)]
104pub enum DefKind {
105 Function,
107 Class,
109 Struct,
111 Enum,
113 Trait,
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
119pub enum Resolution {
120 Resolved,
122 UnknownSymbol,
124 UnknownModule,
127}
128
129#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
132#[serde(tag = "verdict", rename_all = "snake_case")]
133pub enum Verdict {
134 Resolved,
136 NotBuilt {
139 root: String,
141 },
142 Fabricated {
145 detail: String,
147 },
148}
149
150#[derive(Debug, Clone, Default)]
153pub struct SymbolIndex {
154 modules: BTreeSet<String>,
155 symbols: BTreeSet<(String, String)>,
156}
157
158impl SymbolIndex {
159 #[must_use]
161 pub fn new() -> Self {
162 Self::default()
163 }
164
165 pub fn add_module(&mut self, module: impl Into<String>) {
167 self.modules.insert(module.into());
168 }
169
170 pub fn add_symbol(&mut self, module: impl Into<String>, symbol: impl Into<String>) {
172 let module = module.into();
173 self.modules.insert(module.clone());
174 self.symbols.insert((module, symbol.into()));
175 }
176
177 pub fn add_from_source(&mut self, module: &str, source: &str, lang: Lang) {
181 for def in extract_definitions(source, lang) {
182 self.add_symbol(module, def.name);
183 }
184 }
185
186 #[must_use]
188 pub fn has_module(&self, module: &str) -> bool {
189 self.modules.contains(module)
190 }
191
192 #[must_use]
194 pub fn resolve(&self, reference: &Reference) -> Resolution {
195 if !self.modules.contains(&reference.module) {
196 return Resolution::UnknownModule;
197 }
198 match &reference.name {
199 Some(name) => {
200 if self
201 .symbols
202 .contains(&(reference.module.clone(), name.clone()))
203 {
204 Resolution::Resolved
205 } else {
206 Resolution::UnknownSymbol
207 }
208 }
209 None => Resolution::Resolved,
210 }
211 }
212
213 #[must_use]
217 pub fn classify(&self, reference: &Reference, unbuilt_roots: &BTreeSet<String>) -> Verdict {
218 match self.resolve(reference) {
219 Resolution::Resolved => Verdict::Resolved,
220 Resolution::UnknownSymbol => Verdict::Fabricated {
221 detail: format!(
222 "module `{}` has no symbol `{}`",
223 reference.module,
224 reference.name.as_deref().unwrap_or("")
225 ),
226 },
227 Resolution::UnknownModule => {
228 let root = module_root(&reference.module);
229 if unbuilt_roots.contains(root) {
230 Verdict::NotBuilt {
231 root: root.to_string(),
232 }
233 } else {
234 Verdict::Fabricated {
235 detail: format!("no module `{}`", reference.module),
236 }
237 }
238 }
239 }
240 }
241}
242
243fn module_root(module: &str) -> &str {
246 module.split(['.', ':']).next().unwrap_or(module)
247}
248
249#[must_use]
254pub fn module_is_known(module: &str, known: &BTreeSet<String>) -> bool {
255 let parts: Vec<&str> = module.split('.').collect();
256 (1..=parts.len()).any(|i| known.contains(&parts[..i].join(".")))
257}
258
259#[must_use]
264pub fn python_stdlib_modules() -> BTreeSet<String> {
265 [
266 "__future__",
269 "__main__",
270 "abc",
271 "argparse",
272 "asyncio",
273 "base64",
274 "collections",
275 "contextlib",
276 "copy",
277 "csv",
278 "dataclasses",
279 "datetime",
280 "decimal",
281 "enum",
282 "functools",
283 "glob",
284 "hashlib",
285 "io",
286 "itertools",
287 "json",
288 "logging",
289 "math",
290 "os",
291 "pathlib",
292 "random",
293 "re",
294 "shutil",
295 "subprocess",
296 "sys",
297 "tempfile",
298 "time",
299 "typing",
300 "unittest",
301 "uuid",
302 "warnings",
303 ]
304 .iter()
305 .map(|s| (*s).to_string())
306 .collect()
307}
308
309#[must_use]
311pub fn extract_references(source: &str, lang: Lang) -> Vec<Reference> {
312 match lang {
313 Lang::Python => extract_references_python(source),
314 Lang::Rust => extract_references_rust(source),
315 }
316}
317
318#[must_use]
320pub fn extract_definitions(source: &str, lang: Lang) -> Vec<Definition> {
321 match lang {
322 Lang::Python => extract_definitions_python(source),
323 Lang::Rust => extract_definitions_rust(source),
324 }
325}
326
327fn extract_references_python(source: &str) -> Vec<Reference> {
328 let from_re = Regex::new(r"^\s*from\s+([\w.-]+)\s+import\s+(.+?)\s*$").unwrap();
332 let import_re = Regex::new(r"^\s*import\s+(.+?)\s*$").unwrap();
334
335 let mut refs = Vec::new();
336 for (i, line) in source.lines().enumerate() {
337 let lineno = i + 1;
338 if let Some(caps) = from_re.captures(line) {
339 let module = caps[1].to_string();
340 if module.starts_with('.') {
344 continue;
345 }
346 let names = caps[2].trim();
347 if names == "(" {
352 refs.push(Reference::import_module(module, lineno));
353 continue;
354 }
355 let before = refs.len();
356 for item in names.split(',') {
357 let name = item
358 .split_whitespace()
359 .next()
360 .unwrap_or("")
361 .trim_matches(|c| c == '(' || c == ')');
362 if name.is_empty() || name == "*" {
363 continue;
364 }
365 refs.push(Reference::import_from(module.clone(), name, lineno));
366 }
367 if refs.len() == before {
371 refs.push(Reference::import_module(module, lineno));
372 }
373 } else if let Some(caps) = import_re.captures(line) {
374 for item in caps[1].split(',') {
375 let module = item.split_whitespace().next().unwrap_or("");
377 if !module.is_empty() {
378 refs.push(Reference::import_module(module, lineno));
379 }
380 }
381 }
382 }
383 refs
384}
385
386fn extract_definitions_python(source: &str) -> Vec<Definition> {
387 let def_re = Regex::new(r"^\s*def\s+(\w+)").unwrap();
388 let class_re = Regex::new(r"^\s*class\s+(\w+)").unwrap();
389
390 let mut defs = Vec::new();
391 for (i, line) in source.lines().enumerate() {
392 let lineno = i + 1;
393 if let Some(caps) = def_re.captures(line) {
394 defs.push(Definition {
395 name: caps[1].to_string(),
396 kind: DefKind::Function,
397 line: lineno,
398 });
399 } else if let Some(caps) = class_re.captures(line) {
400 defs.push(Definition {
401 name: caps[1].to_string(),
402 kind: DefKind::Class,
403 line: lineno,
404 });
405 }
406 }
407 defs
408}
409
410fn extract_references_rust(source: &str) -> Vec<Reference> {
411 let use_re = Regex::new(r"^\s*(?:pub\s*(?:\([^)]*\)\s*)?)?use\s+(.+?)\s*;").unwrap();
415 let mut refs = Vec::new();
416 for (i, line) in source.lines().enumerate() {
417 let lineno = i + 1;
418 if let Some(caps) = use_re.captures(line) {
419 parse_rust_use(caps[1].trim(), lineno, &mut refs);
420 }
421 }
422 refs
423}
424
425fn parse_rust_use(path: &str, lineno: usize, refs: &mut Vec<Reference>) {
428 if let Some(open) = path.find("::{") {
430 let module = &path[..open];
431 let inner = path[open + 3..].trim_end_matches('}');
432 for item in inner.split(',') {
433 let item = item.split_whitespace().next().unwrap_or("");
435 if item.is_empty() || item == "*" || item == "self" {
436 refs.push(Reference::import_module(module, lineno));
439 } else {
440 refs.push(Reference::import_from(module, item, lineno));
441 }
442 }
443 return;
444 }
445 if let Some(module) = path.strip_suffix("::*") {
447 refs.push(Reference::import_module(module, lineno));
448 return;
449 }
450 let head = path.split_whitespace().next().unwrap_or(path);
452 if let Some(idx) = head.rfind("::") {
453 refs.push(Reference::import_from(
454 &head[..idx],
455 &head[idx + 2..],
456 lineno,
457 ));
458 } else {
459 refs.push(Reference::import_module(head, lineno));
461 }
462}
463
464fn extract_definitions_rust(source: &str) -> Vec<Definition> {
465 let vis = r"(?:pub\s*(?:\([^)]*\)\s*)?)?";
467 let fn_re = Regex::new(&format!(
468 r"^\s*{vis}(?:async\s+|unsafe\s+|const\s+|extern\s+(?:\x22[^\x22]*\x22\s+)?)*fn\s+(\w+)"
469 ))
470 .unwrap();
471 let struct_re = Regex::new(&format!(r"^\s*{vis}struct\s+(\w+)")).unwrap();
472 let enum_re = Regex::new(&format!(r"^\s*{vis}enum\s+(\w+)")).unwrap();
473 let trait_re = Regex::new(&format!(r"^\s*{vis}(?:unsafe\s+)?trait\s+(\w+)")).unwrap();
474
475 let mut defs = Vec::new();
476 for (i, line) in source.lines().enumerate() {
477 let lineno = i + 1;
478 let captured = [
479 (&fn_re, DefKind::Function),
480 (&struct_re, DefKind::Struct),
481 (&enum_re, DefKind::Enum),
482 (&trait_re, DefKind::Trait),
483 ]
484 .into_iter()
485 .find_map(|(re, kind)| re.captures(line).map(|c| (c[1].to_string(), kind)));
486 if let Some((name, kind)) = captured {
487 defs.push(Definition {
488 name,
489 kind,
490 line: lineno,
491 });
492 }
493 }
494 defs
495}
496
497#[cfg(test)]
498mod tests {
499 use super::*;
500
501 fn built() -> BTreeSet<String> {
502 BTreeSet::new()
503 }
504
505 #[test]
506 fn oracle_catches_the_second_nemotron_incident() {
507 let mut idx = SymbolIndex::new();
512 idx.add_symbol("newt_agent.core", "Router");
513 idx.add_symbol("newt_agent.core", "Tier");
514 idx.add_symbol("newt_agent.data", "load_csv_to_sqlite");
515
516 let fab1 = Reference::import_from("newt_core", "classify", 1);
518 let fab2 = Reference::import_from("newt_data", "DataStore", 2);
519 assert!(matches!(
520 idx.classify(&fab1, &built()),
521 Verdict::Fabricated { .. }
522 ));
523 assert!(matches!(
524 idx.classify(&fab2, &built()),
525 Verdict::Fabricated { .. }
526 ));
527
528 let real = Reference::import_from("newt_agent.core", "Router", 3);
530 assert_eq!(idx.classify(&real, &built()), Verdict::Resolved);
531
532 let bad_symbol = Reference::import_from("newt_agent.core", "classify", 4);
534 assert!(matches!(
535 idx.classify(&bad_symbol, &built()),
536 Verdict::Fabricated { .. }
537 ));
538 }
539
540 #[test]
541 fn not_built_is_distinguished_from_fabricated() {
542 let mut idx = SymbolIndex::new();
544 idx.add_module("newt_agent");
545 let unbuilt: BTreeSet<String> = ["newt_agent".to_string()].into_iter().collect();
546
547 let r = Reference::import_from("newt_agent.core", "Router", 1);
549 assert_eq!(
550 idx.classify(&r, &unbuilt),
551 Verdict::NotBuilt {
552 root: "newt_agent".to_string()
553 }
554 );
555
556 let fab = Reference::import_from("newt_core", "classify", 2);
558 assert!(matches!(
559 idx.classify(&fab, &unbuilt),
560 Verdict::Fabricated { .. }
561 ));
562 }
563
564 #[test]
565 fn extracts_python_imports() {
566 let src = "from newt_agent.core import Router, Tier\nimport os\nimport newt_data as nd\nfrom x import *\n";
567 let refs = extract_references(src, Lang::Python);
568 assert!(refs.contains(&Reference::import_from("newt_agent.core", "Router", 1)));
569 assert!(refs.contains(&Reference::import_from("newt_agent.core", "Tier", 1)));
570 assert!(refs.contains(&Reference::import_module("os", 2)));
571 assert!(refs.contains(&Reference::import_module("newt_data", 3)));
572 assert!(!refs.iter().any(|r| r.name.as_deref() == Some("*")));
574 }
575
576 #[test]
577 fn extracts_python_definitions() {
578 let src = "def foo():\n pass\n\nclass Bar:\n pass\n";
579 let defs = extract_definitions(src, Lang::Python);
580 assert!(defs
581 .iter()
582 .any(|d| d.name == "foo" && d.kind == DefKind::Function));
583 assert!(defs
584 .iter()
585 .any(|d| d.name == "Bar" && d.kind == DefKind::Class));
586 }
587
588 #[test]
589 fn add_from_source_builds_a_resolvable_index() {
590 let mut idx = SymbolIndex::new();
591 idx.add_from_source(
592 "mypkg.mod",
593 "def helper():\n pass\nclass Widget:\n pass\n",
594 Lang::Python,
595 );
596 assert_eq!(
597 idx.resolve(&Reference::import_from("mypkg.mod", "Widget", 1)),
598 Resolution::Resolved
599 );
600 assert_eq!(
601 idx.resolve(&Reference::import_from("mypkg.mod", "Nope", 1)),
602 Resolution::UnknownSymbol
603 );
604 assert_eq!(
605 idx.resolve(&Reference::import_from("other.mod", "Widget", 1)),
606 Resolution::UnknownModule
607 );
608 }
609
610 #[test]
611 fn bare_import_only_checks_module_existence() {
612 let mut idx = SymbolIndex::new();
613 idx.add_module("os");
614 assert_eq!(
615 idx.resolve(&Reference::import_module("os", 1)),
616 Resolution::Resolved
617 );
618 assert_eq!(
619 idx.resolve(&Reference::import_module("nope", 1)),
620 Resolution::UnknownModule
621 );
622 }
623
624 #[test]
625 fn extracts_rust_use_statements() {
626 let src = "use a::b::Thing;\n\
627 pub use c::d::{One, Two as T, self, *};\n\
628 use crate::scope::*;\n\
629 use std::collections::HashMap as Map;\n";
630 let refs = extract_references(src, Lang::Rust);
631 assert!(refs.contains(&Reference::import_from("a::b", "Thing", 1)));
632 assert!(refs.contains(&Reference::import_from("c::d", "One", 2)));
633 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)));
636 assert!(refs.contains(&Reference::import_from("std::collections", "HashMap", 4)));
637 }
638
639 #[test]
640 fn extracts_rust_definitions() {
641 let src = "pub fn alpha() {}\n\
642 async fn beta() {}\n\
643 pub(crate) struct Gamma { x: u8 }\n\
644 enum Delta { A, B }\n\
645 pub trait Epsilon {}\n\
646 const NOT_A_DEF: u8 = 0;\n";
647 let defs = extract_definitions(src, Lang::Rust);
648 assert!(defs
649 .iter()
650 .any(|d| d.name == "alpha" && d.kind == DefKind::Function));
651 assert!(defs
652 .iter()
653 .any(|d| d.name == "beta" && d.kind == DefKind::Function));
654 assert!(defs
655 .iter()
656 .any(|d| d.name == "Gamma" && d.kind == DefKind::Struct));
657 assert!(defs
658 .iter()
659 .any(|d| d.name == "Delta" && d.kind == DefKind::Enum));
660 assert!(defs
661 .iter()
662 .any(|d| d.name == "Epsilon" && d.kind == DefKind::Trait));
663 assert!(!defs.iter().any(|d| d.name == "NOT_A_DEF"));
665 }
666
667 #[test]
668 fn rust_references_resolve_against_the_same_general_core() {
669 let mut idx = SymbolIndex::new();
673 idx.add_symbol("crate::router", "Router");
674 let real = &extract_references("use crate::router::Router;", Lang::Rust)[0];
675 assert_eq!(idx.resolve(real), Resolution::Resolved);
676 let fab = &extract_references("use crate::nonsense::Widget;", Lang::Rust)[0];
677 assert!(matches!(
678 idx.classify(fab, &built()),
679 Verdict::Fabricated { .. }
680 ));
681 }
682
683 #[test]
684 fn module_root_handles_python_and_rust_paths() {
685 assert_eq!(module_root("newt_agent.core"), "newt_agent");
686 assert_eq!(module_root("a::b::c"), "a");
687 assert_eq!(module_root("flat"), "flat");
688 }
689
690 #[test]
691 fn lang_from_path() {
692 assert_eq!(Lang::from_path("foo/bar.py"), Some(Lang::Python));
693 assert_eq!(Lang::from_path("src/lib.rs"), Some(Lang::Rust));
694 assert_eq!(Lang::from_path("README.md"), None);
695 }
696
697 #[test]
698 fn verdict_serializes_with_a_tag() {
699 let v = Verdict::Fabricated {
700 detail: "no module `x`".to_string(),
701 };
702 let json = serde_json::to_string(&v).unwrap();
703 assert!(json.contains("\"verdict\":\"fabricated\""));
704 }
705
706 #[test]
708 fn module_is_known_exact_single_component() {
709 let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
710 assert!(module_is_known("os", &known));
711 }
712
713 #[test]
715 fn module_is_known_prefix_covers_submodule() {
716 let known: BTreeSet<String> = ["newt_agent".to_string()].into_iter().collect();
717 assert!(module_is_known("newt_agent.core", &known));
718 }
719
720 #[test]
722 fn module_is_known_no_match() {
723 let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
724 assert!(!module_is_known("sys", &known));
725 }
726
727 #[test]
729 fn module_is_known_empty_module() {
730 let known: BTreeSet<String> = ["os".to_string()].into_iter().collect();
731 assert!(!module_is_known("", &known));
732 }
733
734 #[test]
736 fn module_is_known_substring_is_not_prefix() {
737 let known: BTreeSet<String> = ["pathlib".to_string()].into_iter().collect();
738 assert!(!module_is_known("os.path", &known));
739 }
740}